Merge pull request #30 from h4gen/sdk-v2

Sdk v2

Author: h4gen

debug_test_output.log

@@ -0,0 +1,49 @@
+
+import asyncio
+import logging
+from pg_replica import connect
+
+# Configure logging to see the output
+logging.basicConfig(level=logging.INFO)
+
+async def run_demo():
+    print("🚀 Starting God Tier DX Demo...")
+    
+    # 1. Connect (Mocking the DB connection for this script if needed, 
+    # but let's assume the client initializes correctly)
+    # We use a dummy URL since we might not have a real DB running in this specific interaction
+    client = connect(sink_url="postgresql://user:pass@localhost:5432/db", sync=False)
+    
+    print("✅ Connected.")
+    
+    # 2. Fluent Access
+    products = client.products
+    print(f"✅ Accessed pipeline: {products.name}")
+    
+    # 3. Configure (Plan)
+    print("\n--- Planning Configuration ---")
+    changeset = await products.configure(
+        model="openai/small",
+        columns=["name", "description"],
+        template="Product: $name\nContext: $chunk"
+    )
+    print(changeset)
+    
+    # 4. Apply (Simulated)
+    print("\n--- Applying Configuration ---")
+    await changeset.apply()
+    print("✅ Configuration applied.")
+    
+    # 5. Branching (SearchOps)
+    print("\n--- Creating Branch ---")
+    # interactive=False to avoid waiting for sync in this quick test
+    branch = await products.branch("experiment-v2", model="voyage/large", interactive=False)
+    print(f"✅ Branch created: {branch.name}")
+    
+    # 6. Promotion
+    print("\n--- Promoting Branch ---")
+    await products.promote("experiment-v2")
+    print("✅ Promotion triggered.")
+
+if __name__ == "__main__":
+    asyncio.run(run_demo())

demo_notebook.py

@@ -24,6 +24,8 @@ dependencies = [
     "prometheus-client>=0.23.1",
     "fastapi>=0.128.0",
     "uvicorn>=0.40.0",
+    "rich>=14.2.0",
+    "tqdm>=4.67.1",
 ]
 
 [project.optional-dependencies]

pyproject.toml

@@ -112,7 +112,8 @@ async def search(
             embedding=embedding,
             limit=limit,
             config=config,
-            conn_provider=self._get_conn
+            conn_provider=self._get_conn,
+            target_name=target_name
         )
 
     async def get_status(self) -> dict[str, Any]:

src/pg_replica/client.py

@@ -56,11 +56,17 @@ async def close_pools():
     global _source_pool, _sink_pool
     if _source_pool:
         logger.info("Closing source connection pool...")
-        await _source_pool.close()
+        try:
+            await _source_pool.close()
+        except (asyncio.CancelledError, Exception):
+            pass
         _source_pool = None
     if _sink_pool:
         logger.info("Closing sink connection pool...")
-        await _sink_pool.close()
+        try:
+            await _sink_pool.close()
+        except (asyncio.CancelledError, Exception):
+            pass
         _sink_pool = None
 
 
@@ -183,6 +189,7 @@ async def get_source_column_types(
 async def setup_source(settings: Settings, config: SearchPipeline, target_name: str):
     """Remotely initialize the source publication."""
     pub_name = f"pub_{target_name}"
+    print(f"DEBUG: setup_source called for {pub_name}")
     logger.info(f"Setting up remote source publication {pub_name}...")
 
     # Pre-flight readiness check to avoid race conditions in tests
@@ -199,6 +206,7 @@ async def setup_source(settings: Settings, config: SearchPipeline, target_name:
                 else ""
             )
 
+            # Ensure publication exists (idempotent check)
             await cur.execute(
                 f"SELECT 1 FROM pg_publication WHERE pubname = '{pub_name}'"
             )
@@ -210,6 +218,7 @@ async def setup_source(settings: Settings, config: SearchPipeline, target_name:
                     f"CREATE PUBLICATION {pub_name} FOR TABLE {config.ingest.table} ({cols}){where_clause}"
                 )
             else:
+                logger.debug(f"Publication {pub_name} exists. Updating definition...")
                 await cur.execute(
                     f"ALTER PUBLICATION {pub_name} SET TABLE {config.ingest.table} ({cols}){where_clause}"
                 )
@@ -662,7 +671,7 @@ async def ensure_embedding_cache_table(settings: Settings, config: SearchPipelin
             )
 
 async def cleanup_vectorizer_infrastructure(
-    settings: Settings, config: SearchPipeline, vectorizer_name: str
+    settings: Settings, config: SearchPipeline, target_name: str, vectorizer_name: str
 ):
     """Robustly clean up all infrastructure for a specific vectorizer."""
     logger.info(f"Robust cleanup for vectorizer {vectorizer_name}...")
@@ -679,13 +688,13 @@ async def cleanup_vectorizer_infrastructure(
                 -- 1. Check if ANY view is using this as its target (safety)
                 SELECT table_name INTO live_target 
                 FROM information_schema.view_table_usage 
-                WHERE view_name = '{config.ingest.table}_search' 
+                WHERE view_name = '{target_name}_search' 
                 AND table_name IN ('{vectorizer_name}', '{embedding_view}') 
                 LIMIT 1;
                 
                 -- 2. If it's live, we MUST drop the replica view first
                 IF live_target IS NOT NULL THEN
-                    EXECUTE 'DROP VIEW IF EXISTS ' || quote_ident('{config.ingest.table}_search') || ' CASCADE';
+                    EXECUTE 'DROP VIEW IF EXISTS ' || quote_ident('{target_name}_search') || ' CASCADE';
                 END IF;
 
                 -- 3. Drop the pgai vectorizer if it exists
@@ -771,7 +780,7 @@ async def atomic_view_swap(
         await conn.set_autocommit(False)
         try:
             async with conn.cursor() as cur:
-                await cur.execute(f"DROP VIEW IF EXISTS {config.ingest.table}_search")
+                await cur.execute(f"DROP VIEW IF EXISTS {target_name}_search")
 
                 extra_cols = ",\n                            ".join([f"r.{c}" for c in config.ingest.columns if c != config.ingest.p_key])
                 if extra_cols:
@@ -784,7 +793,7 @@ async def atomic_view_swap(
                     logger.info("Implementing Hybrid View with RRF scoring support...")
                     await cur.execute(
                         f"""
-                        CREATE VIEW {config.ingest.table}_search AS
+                        CREATE VIEW {target_name}_search AS
                         SELECT 
                             r.{config.ingest.p_key},
                             e.chunk as chunk,
@@ -798,7 +807,7 @@ async def atomic_view_swap(
                     # VECTOR SEARCH ONLY
                     await cur.execute(
                         f"""
-                        CREATE VIEW {config.ingest.table}_search AS
+                        CREATE VIEW {target_name}_search AS
                         SELECT 
                             r.{config.ingest.p_key},
                             e.chunk as chunk,
@@ -827,7 +836,7 @@ async def ensure_outbox_infrastructure(settings: Settings):
             # 1. The Outbox: Transactional log of vectorized changes
             await cur.execute(
                 """
-                CREATE TABLE IF NOT EXISTS _sink_outbox (
+                CREATE TABLE IF NOT EXISTS public._sink_outbox (
                     id BIGSERIAL PRIMARY KEY,
                     target_name TEXT NOT NULL,
                     version_id TEXT NOT NULL,
@@ -840,7 +849,7 @@ async def ensure_outbox_infrastructure(settings: Settings):
             )
             # Index for the MirrorWorker to poll efficiently
             await cur.execute(
-                "CREATE INDEX IF NOT EXISTS idx_sink_outbox_id ON _sink_outbox(id)"
+                "CREATE INDEX IF NOT EXISTS idx_sink_outbox_id ON public._sink_outbox(id)"
             )
 
             # 2. Mirror Registry: Track progress of external sinks
@@ -926,10 +935,10 @@ async def setup_outbox_trigger(
                 CREATE OR REPLACE FUNCTION {trigger_fn_name}() RETURNS TRIGGER AS $$
                 BEGIN
                     IF (TG_OP = 'DELETE') THEN
-                        INSERT INTO _sink_outbox (target_name, version_id, source_id, action)
+                        INSERT INTO public._sink_outbox (target_name, version_id, source_id, action)
                         VALUES ('{target_name}', '{version_id}', OLD.{config.ingest.p_key}::text, 'DELETE');
                     ELSE
-                        INSERT INTO _sink_outbox (target_name, version_id, source_id, action, payload)
+                        INSERT INTO public._sink_outbox (target_name, version_id, source_id, action, payload)
                         VALUES (
                             '{target_name}', 
                             '{version_id}', 
@@ -949,14 +958,21 @@ async def setup_outbox_trigger(
 
             # 2. Attach Trigger to pgai STORE table
             # We use AFTER INSERT OR UPDATE OR DELETE
-            await cur.execute(
-                f"""
-                DROP TRIGGER IF EXISTS {trigger_name} ON {vectorizer_name};
-                CREATE TRIGGER {trigger_name}
-                AFTER INSERT OR UPDATE OR DELETE ON {vectorizer_name}
-                FOR EACH ROW EXECUTE FUNCTION {trigger_fn_name}();
-                """
-            )
+            # Use try/except to handle race where worker hasn't created table yet
+            try:
+                await cur.execute(
+                    f"""
+                    DROP TRIGGER IF EXISTS {trigger_name} ON {vectorizer_name};
+                    CREATE TRIGGER {trigger_name}
+                    AFTER INSERT OR UPDATE OR DELETE ON {vectorizer_name}
+                    FOR EACH ROW EXECUTE FUNCTION {trigger_fn_name}();
+                    """
+                )
+            except Exception as e:
+                if "does not exist" in str(e):
+                    logger.warning(f"Deferred trigger setup for {vectorizer_name}: Table not yet created by worker.")
+                    return # Retry next loop
+                raise e
 
             # 3. Backfill existing rows (Handling the Race Condition)
             # Since the vectorizer might have already processed rows before we attached the trigger,
@@ -1339,10 +1355,17 @@ async def find_and_fix_ghost_records(settings: Settings, config: SearchPipeline,
         logger.debug(f"No records found for anti-entropy in {target_name}")
         return
 
-    min_id_raw, max_id_raw = min(all_ids), max(all_ids)
     source_types = await get_source_column_types(settings, config)
     id_type = source_types.get(config.ingest.p_key, "TEXT")
 
+    # Ensure all IDs are of the same type for comparison
+    if id_type in ("INT", "BIGINT"):
+        all_ids = [int(x) for x in all_ids]
+    else:
+        all_ids = [str(x) for x in all_ids]
+
+    min_id_raw, max_id_raw = min(all_ids), max(all_ids)
+
     # 2. Strategy: Set Comparison for UUIDs/Strings or Small Tables
     if id_type not in ("INT", "BIGINT"):
         async with await get_source_conn() as s_conn:
@@ -1405,7 +1428,7 @@ async def drop_subscription_completely(settings: Settings, config: SearchPipelin
     try:
         async with await connect_db(settings.resolved_sink_url) as conn:
             await conn.set_autocommit(True)
-            await conn.execute(f"DROP VIEW IF EXISTS {config.ingest.table}_search CASCADE")
+            await conn.execute(f"DROP VIEW IF EXISTS {target_name}_search CASCADE")
 
             # Retry loop for dropping subscription to handle "sync in progress"
             async def try_drop_subscription():
@@ -1439,7 +1462,7 @@ async def try_drop_subscription():
 
             # Cleanup vectorizers
             async with conn.cursor() as cur:
-                await cur.execute("SELECT id FROM ai.vectorizer WHERE name LIKE %s", (f"{config.ingest.table}_store%",))
+                await cur.execute("SELECT id FROM ai.vectorizer WHERE name LIKE %s", (f"{target_name}_store%",))
                 for (vid,) in await cur.fetchall():
                     await cur.execute(f"SELECT ai.drop_vectorizer({vid}, drop_all => true)")
     except Exception as e:

src/pg_replica/database.py

@@ -149,6 +149,26 @@ async def _replication_loop(self):
             except asyncio.TimeoutError:
                 continue
 
+    async def _supervised_run(self, name: str, factory):
+        """Supervises a worker task, restarting it on failure."""
+        logger.info(f"Starting supervised worker: {name}")
+        while not self._stop_event.is_set():
+            try:
+                worker = factory()
+                # Run the worker. If it returns, it finished (unexpected for long-running).
+                # If it raises, we catch and restart.
+                await worker.run()
+            except asyncio.CancelledError:
+                logger.info(f"Worker {name} cancelled.")
+                break
+            except Exception as e:
+                logger.error(f"Worker {name} crashed: {e}. Restarting in 2s...", exc_info=True)
+                try:
+                    await asyncio.sleep(2.0)
+                except asyncio.CancelledError:
+                    break
+        logger.info(f"Worker {name} stopped.")
+
     async def start(self):
         """Start all managed services."""
         if self.settings.sink_url == "local":
@@ -178,11 +198,23 @@ async def try_reconcile():
         except asyncio.TimeoutError as e:
             raise RuntimeError(str(e))
 
-        worker = Worker(db_url=self.settings.resolved_sink_url, poll_interval=timedelta(seconds=2.0))
-        self._tasks.append(asyncio.create_task(worker.run(), name="pgai_worker"))
+        # Supervise pgai worker
+        self._tasks.append(asyncio.create_task(
+            self._supervised_run(
+                "pgai_worker", 
+                lambda: Worker(db_url=self.settings.resolved_sink_url, poll_interval=timedelta(seconds=2.0))
+            ), 
+            name="pgai_worker_supervisor"
+        ))
 
-        mirror_worker = MirrorWorker(self.settings)
-        self._tasks.append(asyncio.create_task(mirror_worker.run(), name="mirror_worker"))
+        # Supervise mirror worker
+        self._tasks.append(asyncio.create_task(
+            self._supervised_run(
+                "mirror_worker",
+                lambda: MirrorWorker(self.settings)
+            ),
+            name="mirror_worker_supervisor"
+        ))
 
         self._tasks.append(asyncio.create_task(self._replication_loop(), name="watchdog"))