feat(persistence): json binary for postgres (#2849)

Arize-ai · Apr 10, 2024 · 29351bf · 29351bf
1 parent 45f084d
commit 29351bf
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py b/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py
@@ -10,13 +10,20 @@
 
 import sqlalchemy as sa
 from alembic import op
+from sqlalchemy import JSON
+from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
 revision: str = "cf03bd6bae1d"
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
+JSON_ = JSON().with_variant(
+    postgresql.JSONB(),  # type: ignore
+    "postgresql",
+)
+
 
 def upgrade() -> None:
     projects_table = op.create_table(
@@ -61,8 +68,8 @@ def upgrade() -> None:
         sa.Column("kind", sa.String, 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("attributes", JSON_, nullable=False),
+        sa.Column("events", JSON_, nullable=False),
         sa.Column(
             "status",
             sa.String,
@@ -87,7 +94,7 @@ def upgrade() -> None:
         sa.Column("label", sa.String, nullable=True),
         sa.Column("score", sa.Float, nullable=True),
         sa.Column("explanation", sa.String, nullable=True),
-        sa.Column("metadata", sa.JSON, nullable=False),
+        sa.Column("metadata", JSON_, nullable=False),
         sa.Column(
             "annotator_kind",
             sa.String,
@@ -126,7 +133,7 @@ def upgrade() -> None:
         sa.Column("label", sa.String, nullable=True),
         sa.Column("score", sa.Float, nullable=True),
         sa.Column("explanation", sa.String, nullable=True),
-        sa.Column("metadata", sa.JSON, nullable=False),
+        sa.Column("metadata", JSON_, nullable=False),
         sa.Column(
             "annotator_kind",
             sa.String,
@@ -166,7 +173,7 @@ def upgrade() -> None:
         sa.Column("label", sa.String, nullable=True),
         sa.Column("score", sa.Float, nullable=True),
         sa.Column("explanation", sa.String, nullable=True),
-        sa.Column("metadata", sa.JSON, nullable=False),
+        sa.Column("metadata", JSON_, nullable=False),
         sa.Column(
             "annotator_kind",
             sa.String,

diff --git a/src/phoenix/db/models.py b/src/phoenix/db/models.py
@@ -13,6 +13,7 @@
     func,
     insert,
 )
+from sqlalchemy.dialects import postgresql
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.orm import (
     DeclarativeBase,
@@ -22,6 +23,11 @@
     relationship,
 )
 
+JSON_ = JSON().with_variant(
+    postgresql.JSONB(),  # type: ignore
+    "postgresql",
+)
+
 
 class UtcTimeStamp(TypeDecorator[datetime]):
     """TODO(persistence): Figure out how to reliably store and retrieve
@@ -72,8 +78,8 @@ class Base(DeclarativeBase):
         }
     )
     type_annotation_map = {
-        Dict[str, Any]: JSON,
-        List[Dict[str, Any]]: JSON,
+        Dict[str, Any]: JSON_,
+        List[Dict[str, Any]]: JSON_,
     }