From 4e25e5dff3c036e5daa63e34b353b7167fee0982 Mon Sep 17 00:00:00 2001
From: Roger Yang <roger.yang@arize.com>
Date: Wed, 10 Apr 2024 12:17:47 -0700
Subject: [PATCH] fix: postgres timestamp insertion

---
 .../migrations/versions/cf03bd6bae1d_init.py  | 24 +++++++++----------
 src/phoenix/db/models.py                      |  4 ++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py b/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py
index 803114dd8b..644a46891c 100644
--- 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
index 1994db4be6..8e25f9e4b5 100644
--- 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(