Replace raw SQL schema with SQLAlchemy Table definitions

optopsy/ui/app.py

@@ -136,72 +136,7 @@ def _patched_plotly_post_init(self):
 
 cl.Plotly.__post_init__ = _patched_plotly_post_init
 
-_DB_SCHEMA_STATEMENTS = [
-    """CREATE TABLE IF NOT EXISTS users (
-        id TEXT PRIMARY KEY,
-        identifier TEXT NOT NULL UNIQUE,
-        "createdAt" TEXT NOT NULL,
-        metadata TEXT DEFAULT '{}'
-    )""",
-    """CREATE TABLE IF NOT EXISTS threads (
-        id TEXT PRIMARY KEY,
-        "userId" TEXT,
-        "userIdentifier" TEXT,
-        "createdAt" TEXT,
-        name TEXT,
-        metadata TEXT,
-        tags TEXT,
-        FOREIGN KEY("userId") REFERENCES users(id)
-    )""",
-    """CREATE TABLE IF NOT EXISTS steps (
-        id TEXT PRIMARY KEY,
-        name TEXT,
-        type TEXT,
-        "threadId" TEXT NOT NULL,
-        "parentId" TEXT,
-        streaming INTEGER DEFAULT 0,
-        "waitForAnswer" INTEGER,
-        "isError" INTEGER,
-        metadata TEXT DEFAULT '{}',
-        tags TEXT,
-        input TEXT,
-        output TEXT,
-        "createdAt" TEXT,
-        start TEXT,
-        "end" TEXT,
-        generation TEXT DEFAULT '{}',
-        "defaultOpen" INTEGER DEFAULT 0,
-        "showInput" TEXT,
-        language TEXT,
-        FOREIGN KEY("threadId") REFERENCES threads(id)
-    )""",
-    """CREATE TABLE IF NOT EXISTS feedbacks (
-        id TEXT PRIMARY KEY,
-        "forId" TEXT NOT NULL,
-        value REAL,
-        comment TEXT,
-        FOREIGN KEY("forId") REFERENCES steps(id)
-    )""",
-    """CREATE TABLE IF NOT EXISTS elements (
-        id TEXT PRIMARY KEY,
-        "threadId" TEXT NOT NULL,
-        type TEXT,
-        "chainlitKey" TEXT,
-        url TEXT,
-        "objectKey" TEXT,
-        name TEXT,
-        display TEXT,
-        size TEXT,
-        language TEXT,
-        page TEXT,
-        "forId" TEXT,
-        mime TEXT,
-        props TEXT DEFAULT '{}',
-        "autoPlay" TEXT,
-        "playerConfig" TEXT,
-        FOREIGN KEY("threadId") REFERENCES threads(id)
-    )""",
-]
+from optopsy.ui.models import metadata as _db_metadata
 
 
 def _get_async_conninfo() -> str:
@@ -242,17 +177,18 @@ def _get_sync_conninfo() -> str:
 def _init_db_sync() -> None:
     """Create tables synchronously at module import time.
 
-    Uses SQLAlchemy so the same DDL works for both SQLite and PostgreSQL.
-    The sync Postgres path requires ``psycopg2`` (included in the ``ui`` extra).
+    Uses SQLAlchemy's ``metadata.create_all()`` so the DDL is generated from
+    the Table definitions in ``models.py`` — emitting native ``UUID``,
+    ``JSONB``, ``TEXT[]``, and ``BOOLEAN`` on PostgreSQL while falling back
+    to ``TEXT`` / ``INTEGER`` on SQLite.
 
     Retries up to 5 times with exponential backoff so the app survives
     transient database unavailability (e.g. Railway starting the DB service
     concurrently with the app).
     """
     import time
 
-    from sqlalchemy import create_engine, text
-    from sqlalchemy.exc import OperationalError, ProgrammingError
+    from sqlalchemy import create_engine
 
     sync_url = _get_sync_conninfo()
 
@@ -265,39 +201,7 @@ def _init_db_sync() -> None:
     for attempt in range(max_retries):
         engine = create_engine(sync_url)
         try:
-            with engine.begin() as conn:
-                for stmt in _DB_SCHEMA_STATEMENTS:
-                    conn.execute(text(stmt))
-
-            # Add columns introduced in newer Chainlit versions.
-            # Each ALTER TABLE runs in its own transaction so that a
-            # "column already exists" error on PostgreSQL doesn't abort
-            # subsequent statements (PostgreSQL marks the whole
-            # transaction as failed after any error).  Using separate
-            # transactions avoids SAVEPOINTs, which are unreliable with
-            # pysqlite's default transaction handling.
-            is_pg = sync_url.startswith("postgresql")
-            for col, definition in [
-                ("defaultOpen", "INTEGER DEFAULT 0"),
-                ("waitForAnswer", "INTEGER"),
-            ]:
-                try:
-                    with engine.begin() as conn:
-                        if is_pg:
-                            conn.execute(
-                                text(
-                                    f'ALTER TABLE steps ADD COLUMN IF NOT EXISTS "{col}" {definition}'
-                                )
-                            )
-                        else:
-                            conn.execute(
-                                text(
-                                    f'ALTER TABLE steps ADD COLUMN "{col}" {definition}'
-                                )
-                            )
-                except (OperationalError, ProgrammingError):
-                    pass  # column already exists
-
+            _db_metadata.create_all(engine)
             return  # success
         except Exception:
             engine.dispose()

optopsy/ui/models.py

@@ -0,0 +1,116 @@
+"""SQLAlchemy table definitions for the Chainlit persistence layer.
+
+Defines a single schema that emits dialect-appropriate DDL for both SQLite
+and PostgreSQL via ``metadata.create_all(engine)``.
+
+On PostgreSQL, columns use native ``UUID``, ``JSONB``, ``TEXT[]``, and
+``BOOLEAN`` types.  On SQLite the same definitions fall back to ``TEXT``
+and ``INTEGER`` — no hand-written SQL required.
+"""
+
+from sqlalchemy import (
+    Boolean,
+    Column,
+    ForeignKey,
+    Integer,
+    MetaData,
+    Table,
+    Text,
+)
+from sqlalchemy.dialects.postgresql import ARRAY as PG_ARRAY
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.dialects.postgresql import UUID as PG_UUID
+
+metadata = MetaData()
+
+# ---------------------------------------------------------------------------
+# Dialect-aware column types
+#
+# On PostgreSQL these resolve to native UUID / JSONB / TEXT[].
+# On SQLite they fall back to plain TEXT so Chainlit's raw-SQL data layer
+# (which serialises everything to strings) keeps working unchanged.
+# ---------------------------------------------------------------------------
+_UUID = Text().with_variant(PG_UUID(as_uuid=False), "postgresql")
+_JSON = Text().with_variant(JSONB(), "postgresql")
+_TAGS = Text().with_variant(PG_ARRAY(Text), "postgresql")
+
+users = Table(
+    "users",
+    metadata,
+    Column("id", _UUID, primary_key=True),
+    Column("identifier", Text, nullable=False, unique=True),
+    Column("metadata", _JSON, nullable=False),
+    Column("createdAt", Text),
+)
+
+threads = Table(
+    "threads",
+    metadata,
+    Column("id", _UUID, primary_key=True),
+    Column("createdAt", Text),
+    Column("name", Text),
+    Column("userId", _UUID, ForeignKey("users.id", ondelete="CASCADE")),
+    Column("userIdentifier", Text),
+    Column("tags", _TAGS),
+    Column("metadata", _JSON),
+)
+
+steps = Table(
+    "steps",
+    metadata,
+    Column("id", _UUID, primary_key=True),
+    Column("name", Text, nullable=False),
+    Column("type", Text, nullable=False),
+    Column(
+        "threadId",
+        _UUID,
+        ForeignKey("threads.id", ondelete="CASCADE"),
+        nullable=False,
+    ),
+    Column("parentId", _UUID),
+    Column("streaming", Boolean, nullable=False),
+    Column("waitForAnswer", Boolean),
+    Column("isError", Boolean),
+    Column("metadata", _JSON),
+    Column("tags", _TAGS),
+    Column("input", Text),
+    Column("output", Text),
+    Column("createdAt", Text),
+    Column("command", Text),
+    Column("start", Text),
+    Column("end", Text),
+    Column("generation", _JSON),
+    Column("showInput", Text),
+    Column("language", Text),
+    Column("indent", Integer),
+    Column("defaultOpen", Boolean),
+)
+
+elements = Table(
+    "elements",
+    metadata,
+    Column("id", _UUID, primary_key=True),
+    Column("threadId", _UUID, ForeignKey("threads.id", ondelete="CASCADE")),
+    Column("type", Text),
+    Column("url", Text),
+    Column("chainlitKey", Text),
+    Column("name", Text, nullable=False),
+    Column("display", Text),
+    Column("objectKey", Text),
+    Column("size", Text),
+    Column("page", Integer),
+    Column("language", Text),
+    Column("forId", _UUID),
+    Column("mime", Text),
+    Column("props", _JSON),
+)
+
+feedbacks = Table(
+    "feedbacks",
+    metadata,
+    Column("id", _UUID, primary_key=True),
+    Column("forId", _UUID, nullable=False),
+    Column("threadId", _UUID, ForeignKey("threads.id", ondelete="CASCADE")),
+    Column("value", Integer, nullable=False),
+    Column("comment", Text),
+)