OS-6741. Replaced datetime utcnow with custom lib

## Description

Replaced datetime utcnow with custom lib

## Related issue number

Replaced datetime utcnow with custom lib

## Special notes

<!-- Please provide additional information if required. -->

## Checklist

* [ ] The pull request title is a good summary of the changes
* [ ] Unit tests for the changes exist
* [ ] New and existing unit tests pass locally

arcee/Dockerfile

@@ -6,6 +6,7 @@ ENV PYTHONPATH /usr/src/app/
 
 COPY arcee/arcee_receiver/requirements.txt arcee/arcee_receiver/requirements.txt
 COPY optscale_client/aconfig_cl optscale_client/aconfig_cl
+COPY tools/optscale_time tools/optscale_time
 
 RUN pip install --no-cache-dir -r /usr/src/app/arcee/arcee_receiver/requirements.txt
 

arcee/arcee_receiver/requirements.txt

@@ -8,3 +8,4 @@ mongodb-migrations==1.2.1
 pydantic==2.4.2
 # OptScale packages
 -e optscale_client/aconfig_cl
+-e tools/optscale_time

arcee/arcee_receiver/server.py

@@ -33,6 +33,7 @@
     get_metrics as _get_task_metrics)
 
 from optscale_client.aconfig_cl.aconfig_cl import AConfigCl
+import tools.optscale_time as opttime
 
 app = Sanic("arcee")
 
@@ -648,7 +649,7 @@ async def update_run(request, body: RunPatchIn, run_id: str):
     d = body.model_dump(exclude_unset=True, exclude={'finish'})
     # TODO: remove "finish" from PATCH payload. Set ts based on "state"
     if body.finish:
-        d.update({"finish": int(datetime.utcnow().timestamp())})
+        d.update({"finish": opttime.utcnow_timestamp()})
     hyperparameters = d.get("hyperparameters", {})
     if hyperparameters:
         existing_hyperparams = r.get("hyperparameters", {})
@@ -683,7 +684,7 @@ async def create_run_milestone(request, run_id: str):
     d = {
         "_id": str(uuid.uuid4()),
         "run_id": run_id,
-        "timestamp": int(datetime.utcnow().timestamp()),
+        "timestamp": opttime.utcnow_timestamp(),
         "milestone": milestone,
     }
     await db.milestone.insert_one(
@@ -1170,7 +1171,7 @@ async def create_token(request):
     d = {
         "_id": str(uuid.uuid4()),
         "token": token,
-        "created": int(datetime.utcnow().timestamp()),
+        "created": opttime.utcnow_timestamp(),
         "deleted_at": 0,
     }
     await db.token.insert_one(
@@ -1201,8 +1202,7 @@ async def delete_token(request, token: str):
     await db.token.update_one(
         {"_id": token_id}, {
             '$set': {
-                "deleted_at": int(
-                    datetime.utcnow().timestamp()),
+                "deleted_at": int(opttime.utcnow_timestamp()),
             }
         })
     return json({"deleted": True, "id": token_id})
@@ -1256,7 +1256,7 @@ async def create_stage(request, run_id: str):
     d = {
         "_id": str(uuid.uuid4()),
         "run_id": run_id,
-        "timestamp": int(datetime.utcnow().timestamp()),
+        "timestamp": opttime.utcnow_timestamp(),
         "name": stage_name,
     }
     await db.stage.insert_one(
@@ -1313,7 +1313,7 @@ async def create_proc_data(request, run_id: str):
     d = {
         "_id": str(uuid.uuid4()),
         "run_id": run_id,
-        "timestamp": int(datetime.utcnow().timestamp()),
+        "timestamp": opttime.utcnow_timestamp(),
         'instance_id': instance,
         "proc_stats": proc_stats,
     }
@@ -1619,7 +1619,7 @@ async def delete_leaderboard(request, id_: str):
     if not o:
         raise SanicException("Leaderboard not found", status_code=404)
     await db.leaderboard.update_one({"_id": id_}, {'$set': {
-        "deleted_at": int(datetime.utcnow().timestamp())
+        "deleted_at": opttime.utcnow_timestamp()
     }})
     return json('', status=204)
 
@@ -1950,7 +1950,7 @@ async def delete_dataset(request, id_: str):
     if await _dataset_used_in_leaderboard(db, id_):
         raise SanicException("Dataset used in leaderboard", status_code=409)
     await db.dataset.update_one({"_id": id_}, {'$set': {
-        "deleted_at": int(datetime.utcnow().timestamp())
+        "deleted_at": opttime.utcnow_timestamp()
     }})
     return json('', status=204)
 

auth/Dockerfile

@@ -10,6 +10,7 @@ RUN apt-get update && apt-get install -y libsodium-dev \
     && rm -rf /var/lib/apt/lists/*
 
 COPY tools/optscale_exceptions tools/optscale_exceptions
+COPY tools/optscale_time tools/optscale_time
 COPY optscale_client/config_client optscale_client/config_client
 COPY optscale_client/rest_api_client optscale_client/rest_api_client
 COPY auth/zoho_integrator auth/zoho_integrator

auth/auth_server/alembic/versions/d7d1cf2f5182_initial.py

@@ -11,7 +11,7 @@
 import uuid
 import string
 import random
-from datetime import datetime
+from datetime import datetime, timezone
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy import (Column, String, Integer, ForeignKey, Table,
@@ -40,7 +40,7 @@ def gen_salt():
 
 
 def get_current_timestamp():
-    return int(datetime.utcnow().timestamp())
+    return int(datetime.now(tz=timezone.utc).timestamp())
 
 
 def create_tables():

auth/auth_server/auth_token/token_store.py

@@ -1,5 +1,4 @@
 import logging
-import datetime
 from sqlalchemy import and_, or_
 from auth.auth_server.auth_token.macaroon import MacaroonToken
 from auth.auth_server.exceptions import Err
@@ -8,6 +7,7 @@
 from auth.auth_server.utils import get_context_values, get_digest
 from tools.optscale_exceptions.common_exc import (UnauthorizedException,
                                                   ForbiddenException)
+import tools.optscale_time as opttime
 
 LOG = logging.getLogger(__name__)
 
@@ -25,7 +25,7 @@ def session(self):
 
     def check_token_valid(self, token_str):
         token = self.session.query(Token).filter(
-            Token.valid_until >= datetime.datetime.utcnow(),
+            Token.valid_until >= opttime.utcnow(),
             Token.digest == get_digest(token_str),
         ).all()
         if not token:

auth/auth_server/controllers/assignment.py

@@ -1,5 +1,4 @@
 import logging
-import datetime
 from sqlalchemy import and_
 
 from auth.auth_server.controllers.base import BaseController
@@ -11,7 +10,7 @@
 from tools.optscale_exceptions.common_exc import (ForbiddenException,
                                                   NotFoundException,
                                                   WrongArgumentsException)
-
+import tools.optscale_time as opttime
 
 LOG = logging.getLogger(__name__)
 
@@ -106,7 +105,7 @@ def delete(self, item_id, **kwargs):
             token = kwargs.get('token')
             user = self.get_user(token)
             self._check_assign_ability(token, user, item)
-        item.deleted_at = datetime.datetime.utcnow().timestamp()
+        item.deleted_at = opttime.utcnow_timestamp()
         self.session.add(item)
         self.session.commit()
 

auth/auth_server/controllers/base.py

@@ -1,6 +1,5 @@
 import time
 import logging
-from datetime import datetime
 import requests
 from ordered_set import OrderedSet
 from sqlalchemy import and_
@@ -18,6 +17,7 @@
                                                   ForbiddenException)
 from tools.optscale_exceptions.http_exc import handle503
 from optscale_client.rest_api_client.client_v2 import Client as RestApiClient
+import tools.optscale_time as opttime
 
 LOG = logging.getLogger(__name__)
 
@@ -106,7 +106,7 @@ def _get_model_type(self):
 
     def get_user(self, token):
         token = self.session.query(Token).get(get_digest(token))
-        if not token or not token.valid_until > datetime.utcnow():
+        if not token or not token.valid_until > opttime.utcnow():
             raise UnauthorizedException(Err.OA0023, [])
         return token.user
 
@@ -311,7 +311,7 @@ def nested_dict_iter(nested):
     def use_verification_code(self, email, code):
         if not email or not code:
             return
-        now = datetime.utcnow()
+        now = opttime.utcnow()
         return self.session.query(VerificationCode).filter(
             and_(
                 VerificationCode.email == email,

auth/auth_server/controllers/token.py

@@ -15,6 +15,7 @@
 from tools.optscale_exceptions.common_exc import (WrongArgumentsException,
                                                   ForbiddenException,
                                                   NotFoundException)
+from tools.optscale_time import utcnow
 from optscale_client.config_client.client import etcd
 
 LOG = logging.getLogger(__name__)
@@ -98,7 +99,7 @@ def create_token_by_user_id(self, **kwargs):
     def create_user_token(self, user, **kwargs):
         model_type = self._get_model_type()
         LOG.info("Creating %s with parameters %s", model_type.__name__, kwargs)
-        now = datetime.datetime.utcnow()
+        now = utcnow()
         macaroon_token = MacaroonToken(user.salt, user.id).create(
             xstr(kwargs.get('register', False)),
             xstr(kwargs.get('provider', 'optscale'))

auth/auth_server/controllers/user.py

@@ -1,4 +1,3 @@
-import datetime
 import logging
 import re
 from sqlalchemy.sql import func
@@ -15,6 +14,7 @@
     check_action, hash_password, is_email_format, get_input,
     check_string_attribute, check_bool_attribute, is_hystax_email,
     is_demo_email)
+from tools.optscale_time import utcnow
 from tools.optscale_exceptions.common_exc import (
     WrongArgumentsException, ForbiddenException, NotFoundException,
     ConflictException)
@@ -185,7 +185,7 @@ def delete(self, item_id, **kwargs):
                                  'DELETE_USER', item.type.name, item.scope_id
                                  ) or self._is_self_edit(user, item_id)):
                 raise ForbiddenException(Err.OA0012, [])
-        item.deleted_at = datetime.datetime.utcnow().timestamp()
+        item.deleted_at = utcnow().timestamp()
         self.session.add(item)
         self.session.commit()
 

auth/auth_server/controllers/verification_code.py

@@ -1,5 +1,5 @@
 import logging
-from datetime import datetime, timedelta
+from datetime import timedelta
 from sqlalchemy import and_, exists
 from sqlalchemy.exc import IntegrityError
 from auth.auth_server.controllers.base import BaseController
@@ -10,6 +10,7 @@
 from auth.auth_server.utils import get_digest
 from tools.optscale_exceptions.common_exc import (WrongArgumentsException,
                                                   ForbiddenException)
+from tools.optscale_time import utcnow
 
 LOG = logging.getLogger(__name__)
 VERIFICATION_CODE_LIFETIME_HRS = 1
@@ -43,7 +44,7 @@ def _check_input(self, **input_):
 
     def _check_generation_timeout(self, email):
         model = self._get_model_type()
-        timeout = datetime.utcnow() - timedelta(
+        timeout = utcnow() - timedelta(
             minutes=GENERATION_THRESHOLD_MIN)
         code_exists = self.session.query(exists().where(and_(
             model.email == email,
@@ -69,7 +70,7 @@ def _invalidate_verification_codes(self, email, deleted_at):
     def create_verification_code(self, email, code):
         model_type = self._get_model_type()
         LOG.info("Creating %s for %s", model_type.__name__, email)
-        now = datetime.utcnow()
+        now = utcnow()
         now_ts = int(now.timestamp())
         params = {
             'email': email,

auth/auth_server/models/models.py

@@ -3,7 +3,6 @@
 import uuid
 import string
 import random
-from datetime import datetime
 from sqlalchemy import Enum, UniqueConstraint
 from sqlalchemy.ext.declarative.base import _declarative_constructor
 from sqlalchemy.ext.declarative import declarative_base, declared_attr
@@ -14,6 +13,7 @@
 
 from auth.auth_server.utils import as_dict, ModelEncoder
 from auth.auth_server.models.exceptions import InvalidTreeException
+from tools.optscale_time import utcnow, utcnow_timestamp
 
 
 def gen_id():
@@ -26,7 +26,7 @@ def gen_salt():
 
 
 def get_current_timestamp():
-    return int(datetime.utcnow().timestamp())
+    return utcnow_timestamp()
 
 
 class PermissionKeys(Enum):
@@ -230,7 +230,7 @@ class Token(Base):
 
     digest = Column(String(32), primary_key=True, nullable=False)
     user_id = Column(String(36), ForeignKey('user.id'))
-    created_at = Column(TIMESTAMP, nullable=False, default=datetime.utcnow)
+    created_at = Column(TIMESTAMP, nullable=False, default=utcnow)
     valid_until = Column(TIMESTAMP, nullable=False, index=True)
     ip = Column(String(39), nullable=False)
     user = relationship("User", backref="tokens")

auth/auth_server/tests/unittests/test_api_role.py

@@ -1,5 +1,4 @@
 # pylint: disable=C0302
-import datetime
 import random
 import string
 from unittest.mock import patch
@@ -8,6 +7,7 @@
                                             Assignment, ActionGroup)
 from auth.auth_server.models.models import gen_salt
 from auth.auth_server.utils import hash_password
+from tools.optscale_time import utcnow_timestamp
 
 
 RES_INFO_URL = "auth.auth_server.controllers.base." \
@@ -233,8 +233,7 @@ def test_get_roles_deleted_agrp(self, p_get_hierarchy, p_get_context,
         self.assertIsNotNone(response['actions'].get(
             self.roles_action_group.name))
         session = self.db_session
-        self.roles_action_group.deleted_at = int(
-            datetime.datetime.utcnow().timestamp())
+        self.roles_action_group.deleted_at = utcnow_timestamp()
         session.add(self.roles_action_group)
         session.commit()
         _, response = self.client.role_get(test_role1.id)
@@ -261,8 +260,7 @@ def test_get_roles_deleted_action(self, p_get_hierarchy, p_get_context,
             self.cloud_sites_action_group.name].get(
             self.list_css_action.name))
         session = self.db_session
-        self.list_css_action.deleted_at = int(
-            datetime.datetime.utcnow().timestamp())
+        self.list_css_action.deleted_at = utcnow_timestamp()
         session.add(self.list_css_action)
         session.commit()
         _, response = self.client.role_get(test_role1.id)

auth/auth_server/tests/unittests/test_api_token.py

@@ -1,10 +1,11 @@
-from datetime import datetime, timedelta
+from datetime import timedelta
 import uuid
 from auth.auth_server.tests.unittests.test_api_base import TestAuthBase
 from auth.auth_server.models.models import (Type, User, gen_salt,
                                             VerificationCode)
 from auth.auth_server.utils import hash_password, get_digest
 from auth.auth_server.tests.unittests.utils import extract_caveats
+from tools.optscale_time import utcnow
 
 
 class TestTokenApi(TestAuthBase):
@@ -125,7 +126,7 @@ def test_token_by_verification_code(self):
         session.add(partner_user)
 
         code_1, code_3, code_4 = 123456, 234567, 345678
-        now = datetime.utcnow()
+        now = utcnow()
         vc_1 = VerificationCode(
             email='[email protected]', valid_until=now + timedelta(hours=1),
             code=get_digest(str(code_1)))

auth/auth_server/tests/unittests/test_api_verification_code.py

@@ -3,6 +3,7 @@
 from auth.auth_server.tests.unittests.test_api_base import TestAuthBase
 from auth.auth_server.models.models import Type, VerificationCode
 from auth.auth_server.utils import get_digest
+from tools.optscale_time import utcnow
 
 
 class TestVerificationCodeApi(TestAuthBase):
@@ -51,7 +52,7 @@ def test_invalid_parameters(self):
 
     def test_invalidation(self):
         code_1 = 1
-        dt = datetime.utcnow()
+        dt = utcnow()
         with freeze_time(dt):
             _, verification_code_1 = self.client.verification_code_create(
                 self.admin_user.email, code_1)
@@ -94,7 +95,7 @@ def test_create_immutable(self):
             self.assertEqual(resp['error']['params'], [field])
 
     def test_generation_time_threshold(self):
-        dt = datetime.utcnow()
+        dt = utcnow()
         with freeze_time(dt):
             code, _ = self.client.verification_code_create(
                 self.admin_user.email, 1)

auth/requirements.txt

@@ -13,5 +13,6 @@ pyyaml==6.0.1
 zcrmsdk==3.1.0
 # OptScale packages
 -e tools/optscale_exceptions
+-e tools/optscale_time
 -e optscale_client/config_client
 -e optscale_client/rest_api_client