Adding eugenes suggestions

backend/app/checkpoint.py

@@ -59,15 +59,20 @@ async def setup(self) -> None:
                 f"{os.environ['POSTGRES_DB']}"
             )
 
-            conn = AsyncConnectionPool(
+            pool = AsyncConnectionPool(
                 conninfo=conninfo,
                 kwargs={"autocommit": True, "prepare_threshold": 0},
+                open=False  # Don't open in constructor
             )
+            await pool.open()
 
             self.async_postgres_saver = AsyncPostgresSaver(
-                conn=conn, pipe=self.pipe, serde=self.serde
+                conn=pool, pipe=self.pipe, serde=self.serde
             )
 
+            # Setup will create/migrate the tables if they don't exist
+            await self.async_postgres_saver.setup()
+
             logger.warning("Checkpoint setup complete.")
         except Exception as e:
             logger.error(f"Failed to set up AsyncPostgresCheckpoint: {e}")

backend/migrations/000002_checkpoints_update_schema.up.sql

@@ -1,41 +1,11 @@
--- Drop existing checkpoints-related tables if they exist
-ALTER TABLE IF EXISTS checkpoints RENAME TO old_checkpoints;
-DROP TABLE IF EXISTS checkpoint_writes;
-DROP TABLE IF EXISTS checkpoint_blobs;
-DROP TABLE IF EXISTS checkpoint_migrations;
+ALTER TABLE checkpoints
+    ADD COLUMN IF NOT EXISTS thread_ts TIMESTAMPTZ,
+    ADD COLUMN IF NOT EXISTS parent_ts TIMESTAMPTZ;
 
-CREATE TABLE IF NOT EXISTS checkpoints (
-    thread_id TEXT NOT NULL,
-    checkpoint_ns TEXT NOT NULL DEFAULT '',
-    checkpoint_id TEXT NOT NULL DEFAULT '',
-    parent_checkpoint_id TEXT,
-    type TEXT,
-    checkpoint JSONB,
-    metadata JSONB NOT NULL DEFAULT '{}',
-    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
-    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
-);
+UPDATE checkpoints
+    SET thread_ts = CURRENT_TIMESTAMP AT TIME ZONE 'UTC'
+WHERE thread_ts IS NULL;
 
--- Create checkpoint_blobs table
-CREATE TABLE IF NOT EXISTS checkpoint_blobs (
-    thread_id TEXT NOT NULL,
-    checkpoint_ns TEXT NOT NULL DEFAULT '',
-    channel TEXT NOT NULL,
-    version TEXT NOT NULL,
-    type TEXT NOT NULL,
-    blob BYTEA,
-    PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
-);
-
--- Create checkpoint_writes table
-CREATE TABLE IF NOT EXISTS checkpoint_writes (
-    thread_id TEXT NOT NULL,
-    checkpoint_ns TEXT NOT NULL DEFAULT '',
-    checkpoint_id TEXT NOT NULL,
-    task_id TEXT NOT NULL,
-    idx INTEGER NOT NULL,
-    channel TEXT NOT NULL,
-    type TEXT,
-    blob BYTEA NOT NULL,
-    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
-);
+ALTER TABLE checkpoints
+    DROP CONSTRAINT IF EXISTS checkpoints_pkey,
+    ADD PRIMARY KEY (thread_id, thread_ts)

backend/migrations/000005_advanced_checkpoints_schema.down.sql

@@ -0,0 +1,12 @@
+-- Drop the blob storage table
+DROP TABLE IF EXISTS checkpoint_blobs;
+
+-- Drop the writes tracking table
+DROP TABLE IF EXISTS checkpoint_writes;
+
+-- Drop the new checkpoints table that was created by the application
+DROP TABLE IF EXISTS checkpoints;
+
+-- Restore the original checkpoints table by renaming old_checkpoints back
+-- This preserves the original data that was saved before the migration
+ALTER TABLE old_checkpoints RENAME TO checkpoints;

backend/migrations/000005_advanced_checkpoints_schema.up.sql

@@ -0,0 +1,9 @@
+-- BREAKING CHANGE WARNING:
+-- This migration represents a transition from pickle-based checkpointing to a new checkpoint system.
+-- As a result, any threads created before this migration will not be usable/clickable in the UI.
+-- old thread data remains in old_checkpoints table but cannot be accessed by the new version.
+
+-- Rename existing checkpoints table to preserve current data
+-- This is necessary because the application will create a new checkpoints table
+-- with an updated schema during runtime initialization.
+ALTER TABLE checkpoints RENAME TO old_checkpoints;

backend/pyproject.toml

@@ -5,6 +5,7 @@ description = ""
 authors = ["Your Name <[email protected]>"]
 readme = "README.md"
 packages = [{include = "app"}]
+package-mode = false
 
 [tool.poetry.dependencies]
 python = "^3.9.0,<3.12"