From fe55ac634b731b101687b3a4c1e5d5356e980ddd Mon Sep 17 00:00:00 2001
From: Viktor Barzin <viktorbarzin@meta.com>
Date: Sat, 14 Mar 2026 12:05:41 +0000
Subject: [PATCH] fix: add psycopg2-binary for Alembic and add connection retry
 logic

---
 pyproject.toml                    |  2 +-
 src/claude_memory/api/database.py | 20 ++++++++++++++++----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index fd2c9a8..41dd3a6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-api = ["fastapi>=0.115", "asyncpg>=0.30", "uvicorn>=0.34", "pydantic>=2.0", "alembic>=1.14", "sqlalchemy>=2.0"]
+api = ["fastapi>=0.115", "asyncpg>=0.30", "uvicorn>=0.34", "pydantic>=2.0", "alembic>=1.14", "sqlalchemy>=2.0", "psycopg2-binary>=2.9"]
 vault = ["hvac>=2.0"]
 dev = ["pytest>=8.0", "pytest-asyncio>=0.24", "ruff>=0.8", "mypy>=1.13", "httpx>=0.28", "cryptography>=43.0"]
 
diff --git a/src/claude_memory/api/database.py b/src/claude_memory/api/database.py
index 063a004..53ae545 100644
--- a/src/claude_memory/api/database.py
+++ b/src/claude_memory/api/database.py
@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import os
 
@@ -17,10 +18,8 @@ def run_migrations() -> None:
         from alembic.config import Config
 
         alembic_cfg = Config()
-        # Find migrations directory relative to this file or project root
         migrations_dir = os.environ.get("ALEMBIC_MIGRATIONS_DIR", "")
         if not migrations_dir:
-            # Check common locations
             for candidate in [
                 os.path.join(os.path.dirname(__file__), "..", "..", "..", "migrations"),
                 os.path.join(os.getcwd(), "migrations"),
@@ -43,10 +42,23 @@ def run_migrations() -> None:
 
 
 async def init_pool() -> asyncpg.Pool:
+    """Initialize connection pool with retries for database availability."""
     global pool
     run_migrations()
-    pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
-    return pool
+
+    max_retries = 5
+    for attempt in range(max_retries):
+        try:
+            pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
+            logger.info("Database pool initialized successfully")
+            return pool
+        except (asyncpg.CannotConnectNowError, OSError, ConnectionRefusedError) as e:
+            if attempt < max_retries - 1:
+                wait = 2 ** attempt
+                logger.warning("Database not ready (attempt %d/%d): %s. Retrying in %ds...", attempt + 1, max_retries, e, wait)
+                await asyncio.sleep(wait)
+            else:
+                raise
 
 
 async def close_pool() -> None: