fix(persistence): postgres timestamp insertion (#2844)

Arize-ai · Apr 10, 2024 · 3477bb9 · 3477bb9
1 parent 3a6666c
commit 3477bb9
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py b/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py
@@ -27,13 +27,13 @@ def upgrade() -> None:
         sa.Column("description", sa.String, nullable=True),
         sa.Column(
             "created_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
         ),
         sa.Column(
             "updated_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
             onupdate=sa.func.now(),
@@ -46,8 +46,8 @@ def upgrade() -> None:
         # TODO(mikeldking): might not be the right place for this
         sa.Column("session_id", sa.String, nullable=True),
         sa.Column("trace_id", sa.String, nullable=False, unique=True),
-        sa.Column("start_time", sa.DateTime(timezone=True), nullable=False, index=True),
-        sa.Column("end_time", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("start_time", sa.TIMESTAMP(timezone=True), nullable=False, index=True),
+        sa.Column("end_time", sa.TIMESTAMP(timezone=True), nullable=False),
         sa.Column("latency_ms", sa.Float, nullable=False),
     )
 
@@ -59,8 +59,8 @@ def upgrade() -> None:
         sa.Column("parent_span_id", sa.String, nullable=True, index=True),
         sa.Column("name", sa.String, nullable=False),
         sa.Column("kind", sa.String, nullable=False),
-        sa.Column("start_time", sa.DateTime(timezone=True), nullable=False),
-        sa.Column("end_time", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("start_time", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.Column("end_time", sa.TIMESTAMP(timezone=True), nullable=False),
         sa.Column("attributes", sa.JSON, nullable=False),
         sa.Column("events", sa.JSON, nullable=False),
         sa.Column(
@@ -99,13 +99,13 @@ def upgrade() -> None:
         ),
         sa.Column(
             "created_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
         ),
         sa.Column(
             "updated_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
             onupdate=sa.func.now(),
@@ -138,13 +138,13 @@ def upgrade() -> None:
         ),
         sa.Column(
             "created_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
         ),
         sa.Column(
             "updated_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
             onupdate=sa.func.now(),
@@ -178,13 +178,13 @@ def upgrade() -> None:
         ),
         sa.Column(
             "created_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
         ),
         sa.Column(
             "updated_at",
-            sa.DateTime(timezone=True),
+            sa.TIMESTAMP(timezone=True),
             nullable=False,
             server_default=sa.func.now(),
             onupdate=sa.func.now(),

diff --git a/src/phoenix/db/models.py b/src/phoenix/db/models.py
@@ -3,8 +3,8 @@
 
 from sqlalchemy import (
     JSON,
+    TIMESTAMP,
     CheckConstraint,
-    DateTime,
     Dialect,
     ForeignKey,
     MetaData,
@@ -33,7 +33,7 @@ class UtcTimeStamp(TypeDecorator[datetime]):
     """
 
     cache_ok = True
-    impl = DateTime
+    impl = TIMESTAMP(timezone=True)
     _LOCAL_TIMEZONE = datetime.now(timezone.utc).astimezone().tzinfo
 
     def process_bind_param(