feat: use timezone aware timestamps (UTC)

harp_apps/janitor/tests/test_worker.py

@@ -10,15 +10,9 @@ class TestJanitorWorker(SqlalchemyStorageTestFixtureMixin):
     async def test_delete_old_transactions(self, storage: SqlAlchemyStorage):
         worker = JanitorWorker(storage)
 
-        await self.create_transaction(
-            storage, started_at=datetime.now(UTC).replace(tzinfo=None) - OLD_AFTER - timedelta(hours=1)
-        )
-        await self.create_transaction(
-            storage, started_at=datetime.now(UTC).replace(tzinfo=None) - OLD_AFTER - timedelta(minutes=1)
-        )
-        await self.create_transaction(
-            storage, started_at=datetime.now(UTC).replace(tzinfo=None) - OLD_AFTER + timedelta(minutes=1)
-        )
+        await self.create_transaction(storage, started_at=datetime.now(UTC) - OLD_AFTER - timedelta(hours=1))
+        await self.create_transaction(storage, started_at=datetime.now(UTC) - OLD_AFTER - timedelta(minutes=1))
+        await self.create_transaction(storage, started_at=datetime.now(UTC) - OLD_AFTER + timedelta(minutes=1))
 
         async with storage.session_factory() as session:
             assert (await worker.compute_metrics(session))["storage.transactions"] == 3
@@ -58,15 +52,9 @@ async def test_delete_orphan_blobs(self, storage: SqlAlchemyStorage):
     async def test_delete_old_transactions_but_keep_flagged_ones(self, storage: SqlAlchemyStorage):
         worker = JanitorWorker(storage)
 
-        t1 = await self.create_transaction(
-            storage, started_at=datetime.now(UTC).replace(tzinfo=None) - OLD_AFTER - timedelta(hours=1)
-        )
-        await self.create_transaction(
-            storage, started_at=datetime.now(UTC).replace(tzinfo=None) - OLD_AFTER - timedelta(minutes=1)
-        )
-        await self.create_transaction(
-            storage, started_at=datetime.now(UTC).replace(tzinfo=None) - OLD_AFTER + timedelta(minutes=1)
-        )
+        t1 = await self.create_transaction(storage, started_at=datetime.now(UTC) - OLD_AFTER - timedelta(hours=1))
+        await self.create_transaction(storage, started_at=datetime.now(UTC) - OLD_AFTER - timedelta(minutes=1))
+        await self.create_transaction(storage, started_at=datetime.now(UTC) - OLD_AFTER + timedelta(minutes=1))
 
         user = await storage.users.find_one_by_username("anonymous")
         await storage.flags.create({"type": 1, "user_id": user.id, "transaction_id": t1.id})

harp_apps/sqlalchemy_storage/models/blobs.py

@@ -1,4 +1,4 @@
-from sqlalchemy import DateTime, LargeBinary, String, delete, func, select
+from sqlalchemy import TIMESTAMP, LargeBinary, String, delete, func, select
 from sqlalchemy.orm import aliased, mapped_column
 
 from harp.models import Blob as BlobModel
@@ -13,7 +13,7 @@ class Blob(Base):
     id = mapped_column(String(40), primary_key=True, unique=True)
     data = mapped_column(LargeBinary())
     content_type = mapped_column(String(64))
-    created_at = mapped_column(DateTime(), server_default=func.now())
+    created_at = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
 
 
 class BlobsRepository(Repository[Blob]):

harp_apps/sqlalchemy_storage/models/messages.py

@@ -1,7 +1,7 @@
 from datetime import UTC
 from typing import TYPE_CHECKING
 
-from sqlalchemy import DateTime, ForeignKey, Index, Integer, String
+from sqlalchemy import TIMESTAMP, ForeignKey, Index, Integer, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from harp.http import get_serializer_for
@@ -21,7 +21,7 @@ class Message(Base):
     summary = mapped_column(String(255))
     headers = mapped_column(String(40))
     body = mapped_column(String(40))
-    created_at = mapped_column(DateTime())
+    created_at = mapped_column(TIMESTAMP(timezone=True))
 
     transaction_id = mapped_column(ForeignKey("sa_transactions.id", ondelete="CASCADE"))
     transaction: Mapped["Transaction"] = relationship(back_populates="messages")
@@ -39,7 +39,7 @@ def to_model(self):
             summary=self.summary,
             headers=self.headers,
             body=self.body,
-            created_at=self.created_at,
+            created_at=self.created_at.replace(tzinfo=UTC) if self.created_at else self.created_at,
         )
 
     @classmethod
@@ -52,7 +52,7 @@ def from_models(cls, transaction, message, headers, content):
         obj.summary = serializer.summary
         obj.headers = headers.id
         obj.body = content.id
-        obj.created_at = message.created_at.astimezone(UTC).replace(tzinfo=None)
+        obj.created_at = message.created_at.astimezone(UTC)
         return obj
 
 

harp_apps/sqlalchemy_storage/models/metrics.py

@@ -36,7 +36,7 @@ def __init__(self, session_factory, /):
 
     @with_session
     async def insert_values(self, values: dict, /, session):
-        now = datetime.now(UTC).replace(tzinfo=None)
+        now = datetime.now(UTC)
         for name, value in values.items():
             metric = await self.find_or_create_one({"name": name}, session=session)
             await self.values.create({"metric_id": metric.id, "value": value, "created_at": now}, session=session)

harp_apps/sqlalchemy_storage/models/transactions.py

@@ -1,7 +1,7 @@
 from datetime import UTC, datetime, timedelta
 from typing import TYPE_CHECKING, List
 
-from sqlalchemy import Column, DateTime, Float, ForeignKey, Index, Integer, String, Table, exists, insert
+from sqlalchemy import TIMESTAMP, Column, Float, ForeignKey, Index, Integer, String, Table, exists, insert
 from sqlalchemy.orm import Mapped, joinedload, mapped_column, relationship, selectinload
 
 from harp.models.transactions import Transaction as TransactionModel
@@ -27,8 +27,8 @@ class Transaction(Base):
     id = mapped_column(String(27), primary_key=True, unique=True)
     type = mapped_column(String(10), index=True)
     endpoint = mapped_column(String(32), nullable=True, index=True)
-    started_at = mapped_column(DateTime(), index=True)
-    finished_at = mapped_column(DateTime(), nullable=True)
+    started_at = mapped_column(TIMESTAMP(timezone=True), index=True)
+    finished_at = mapped_column(TIMESTAMP(timezone=True), nullable=True)
     elapsed = mapped_column(Float(), nullable=True)
     apdex = mapped_column(Integer(), nullable=True)
     x_method = mapped_column(String(16), nullable=True, index=True)
@@ -60,8 +60,8 @@ def to_model(self, with_user_flags=False):
             id=self.id,
             type=self.type,
             endpoint=self.endpoint,
-            started_at=self.started_at,
-            finished_at=self.finished_at,
+            started_at=self.started_at.replace(tzinfo=UTC),
+            finished_at=self.finished_at.replace(tzinfo=UTC) if self.finished_at else self.finished_at,
             elapsed=self.elapsed,
             apdex=self.apdex,
             extras=dict(
@@ -123,7 +123,7 @@ def select(self, /, *, with_messages=False, with_user_flags=False, with_tags=Fal
         return query
 
     def delete_old(self, old_after: timedelta):
-        threshold = (datetime.now(UTC) - old_after).replace(tzinfo=None)
+        threshold = datetime.now(UTC) - old_after
         no_flags = ~exists().where(UserFlag.transaction_id == self.Type.id)
         return self.delete().where((self.Type.started_at < threshold) & no_flags)
 
@@ -136,7 +136,7 @@ async def create(self, values: dict | TransactionModel, /, *, session=None):
                 id=values.id,
                 type=values.type,
                 endpoint=values.endpoint,
-                started_at=values.started_at.replace(tzinfo=None),
+                started_at=values.started_at,
                 x_method=values.extras.get("method"),
                 tags=values.tags,
             )

harp_apps/sqlalchemy_storage/storage.py

@@ -293,7 +293,7 @@ async def transactions_grouped_by_time_bucket(
             query = query.where(Transaction.endpoint == endpoint)
 
         if start_datetime:
-            query = query.where(Transaction.started_at >= start_datetime.astimezone(UTC).replace(tzinfo=None))
+            query = query.where(Transaction.started_at >= start_datetime.astimezone(UTC))
 
         query = query.group_by(s_date).order_by(s_date.asc())
         async with self.begin() as session:
@@ -313,7 +313,7 @@ async def transactions_grouped_by_time_bucket(
     async def get_usage(self):
         async with self.begin() as session:
             query = select(count(Transaction.id)).where(
-                Transaction.started_at > (datetime.now(UTC) - timedelta(hours=24)).replace(tzinfo=None)
+                Transaction.started_at > (datetime.now(UTC) - timedelta(hours=24))
             )
             return (await session.execute(query)).scalar_one_or_none()
 
@@ -401,7 +401,7 @@ async def finalize_transaction():
                     update(Transaction)
                     .where(Transaction.id == event.transaction.id)
                     .values(
-                        finished_at=event.transaction.finished_at.astimezone(UTC).replace(tzinfo=None),
+                        finished_at=event.transaction.finished_at.astimezone(UTC),
                         elapsed=event.transaction.elapsed,
                         apdex=tpdex(event.transaction.elapsed),
                         x_status_class=event.transaction.extras.get("status_class"),

harp_apps/sqlalchemy_storage/utils/testing/mixins.py

@@ -39,7 +39,7 @@ async def create_transaction(self, storage: SqlAlchemyStorage, **kwargs):
                         "id": generate_transaction_id_ksuid(),
                         "type": "http",
                         "endpoint": "/",
-                        "started_at": datetime.now(UTC).replace(tzinfo=None),
+                        "started_at": datetime.now(UTC),
                     },
                     **kwargs,
                 }