Merge pull request #2796 from Arize-ai/2793-postgres

feat(persistece): configurable sql engine

Dockerfile

@@ -36,7 +36,7 @@ COPY ./ /phoenix/
 COPY --from=frontend-builder /phoenix/src/phoenix/server/static/ /phoenix/src/phoenix/server/static/
 # Delete symbolic links used during development.
 RUN find src/ -xtype l -delete  
-RUN pip install --target ./env .[container]
+RUN pip install --target ./env .[container, pg]
 
 # The production image is distroless, meaning that it is a minimal image that
 # contains only the necessary dependencies to run the application. This is

pyproject.toml

@@ -89,6 +89,11 @@ llama-index = [
   "llama-index-callbacks-arize-phoenix>=0.1.2",
   "openinference-instrumentation-llama-index>=1.2.0",
 ]
+
+pg = [
+  "asyncpg",
+]
+
 container = [
   "prometheus-client",
 ]

src/phoenix/config.py

@@ -22,6 +22,10 @@
 """
 The project name to use when logging traces and evals. defaults to 'default'.
 """
+ENV_PHOENIX_SQL_DATABASE = "__DANGEROUS__PHOENIX_SQL_DATABASE"
+"""
+The database URL to use when logging traces and evals.
+"""
 ENV_SPAN_STORAGE_TYPE = "__DANGEROUS__PHOENIX_SPAN_STORAGE_TYPE"
 """
 **EXPERIMENTAL**
@@ -152,6 +156,14 @@ def get_env_span_storage_type() -> Optional["SpanStorageType"]:
         )
 
 
+def get_env_database_connection_str() -> str:
+    env_url = os.getenv(ENV_PHOENIX_SQL_DATABASE)
+    if env_url is None:
+        working_dir = get_working_dir()
+        return f"sqlite:///{working_dir}/phoenix.db"
+    return env_url
+
+
 class SpanStorageType(Enum):
     TEXT_FILES = "text-files"
 

src/phoenix/db/alembic.ini

@@ -94,17 +94,17 @@ keys = console
 keys = generic
 
 [logger_root]
-level = WARN
+level = DEBUG
 handlers = console
 qualname =
 
 [logger_sqlalchemy]
-level = WARN
+level = DEBUG
 handlers =
 qualname = sqlalchemy.engine
 
 [logger_alembic]
-level = WARN
+level = DEBUG
 handlers =
 qualname = alembic
 

src/phoenix/db/engines.py

@@ -7,7 +7,7 @@
 from typing import Any, Union
 
 import numpy as np
-from sqlalchemy import URL, event
+from sqlalchemy import URL, event, make_url
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
 from phoenix.db.migrate import migrate
@@ -24,18 +24,49 @@ def set_sqlite_pragma(connection: Connection, _: Any) -> None:
     cursor.close()
 
 
-def aiosqlite_engine(
+def get_db_url(driver: str = "sqlite+aiosqlite", database: Union[str, Path] = ":memory:") -> URL:
+    return URL.create(driver, database=str(database))
+
+
+def create_engine(connection_str: str, echo: bool = False) -> AsyncEngine:
+    """
+    Factory to create a SQLAlchemy engine from a URL string.
+    """
+    url = make_url(connection_str)
+    if not url.database:
+        raise ValueError("Failed to parse database from connection string")
+    if "sqlite" in url.drivername:
+        # Split the URL to get the database name
+        return aio_sqlite_engine(database=url.database, echo=echo)
+    if "postgresql" in url.drivername:
+        return aio_postgresql_engine(url=url, echo=echo)
+    raise ValueError(f"Unsupported driver: {url.drivername}")
+
+
+def aio_sqlite_engine(
     database: Union[str, Path] = ":memory:",
     echo: bool = False,
 ) -> AsyncEngine:
-    driver_name = "sqlite+aiosqlite"
-    url = URL.create(driver_name, database=str(database))
+    url = get_db_url(driver="sqlite+aiosqlite", database=database)
     engine = create_async_engine(url=url, echo=echo, json_serializer=_dumps)
     event.listen(engine.sync_engine, "connect", set_sqlite_pragma)
     if str(database) == ":memory:":
         asyncio.run(init_models(engine))
     else:
-        migrate(url)
+        migrate(engine.url)
+    return engine
+
+
+def aio_postgresql_engine(
+    url: URL,
+    echo: bool = False,
+) -> AsyncEngine:
+    # Swap out the engine
+    async_url = url.set(drivername="postgresql+asyncpg")
+    engine = create_async_engine(url=async_url, echo=echo, json_serializer=_dumps)
+    # TODO(persistence): figure out the postgres pragma
+    # event.listen(engine.sync_engine, "connect", set_pragma)
+    migrate(engine.url)
     return engine
 
 

src/phoenix/server/main.py

@@ -17,6 +17,7 @@
 
 from phoenix.config import (
     EXPORT_DIR,
+    get_env_database_connection_str,
     get_env_host,
     get_env_port,
     get_pids_path,
@@ -26,7 +27,7 @@
 from phoenix.core.traces import Traces
 from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
 from phoenix.datasets.fixtures import FIXTURES, get_datasets
-from phoenix.db.engines import aiosqlite_engine
+from phoenix.db.engines import create_engine
 from phoenix.pointcloud.umap_parameters import (
     DEFAULT_MIN_DIST,
     DEFAULT_N_NEIGHBORS,
@@ -70,7 +71,7 @@
 |  🚀 Phoenix Server 🚀
 |  Phoenix UI: http://{host}:{port}
 |  Log traces: /v1/traces over HTTP
-|  Storage location: {working_dir}
+|  Storage: {storage}
 """
 
 
@@ -268,7 +269,8 @@ def _send_spans(spans: Iterable[Span], url: str) -> None:
         start_prometheus()
 
     working_dir = get_working_dir().resolve()
-    engine = aiosqlite_engine(working_dir / "phoenix.db")
+    db_connection_str = get_env_database_connection_str()
+    engine = create_engine(db_connection_str)
     app = create_app(
         engine=engine,
         export_path=export_path,
@@ -290,7 +292,7 @@ def _send_spans(spans: Iterable[Span], url: str) -> None:
         "version": phoenix_version,
         "host": host,
         "port": port,
-        "working_dir": working_dir,
+        "storage": db_connection_str,
     }
     print(_WELCOME_MESSAGE.format(**config))
 

src/phoenix/session/session.py

@@ -37,7 +37,7 @@
 from phoenix.core.model_schema_adapter import create_model_from_datasets
 from phoenix.core.traces import Traces
 from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
-from phoenix.db.engines import aiosqlite_engine
+from phoenix.db.engines import aio_sqlite_engine
 from phoenix.pointcloud.umap_parameters import get_umap_parameters
 from phoenix.server.app import create_app
 from phoenix.server.thread_server import ThreadServer
@@ -310,7 +310,7 @@ def __init__(
             ).start()
         # Initialize an app service that keeps the server running
         self.app = create_app(
-            engine=aiosqlite_engine(),
+            engine=aio_sqlite_engine(),
             export_path=self.export_path,
             model=self.model,
             corpus=self.corpus,