Defer DB connection until after jixia analysis completes

database/__init__.py

@@ -6,11 +6,22 @@
 from jixia.structs import parse_name
 
 from .informalize import generate_informal
-from .jixia_db import load_data
+from .jixia_db import run_analysis, load_data
 from .vector_db import create_vector_db
 from .create_schema import create_schema
 
 
+def _connect():
+    return psycopg.connect(
+        os.environ["CONNECTION_STRING"],
+        autocommit=True,
+        keepalives=1,
+        keepalives_idle=30,
+        keepalives_interval=10,
+        keepalives_count=5,
+    )
+
+
 def main():
     parser = ArgumentParser()
     subparser = parser.add_subparsers()
@@ -42,32 +53,31 @@ def main():
 
     args = parser.parse_args()
 
-    # The jixia analysis phase can run 30-60+ min with no DB traffic at all.
-    # GitHub Actions runners sit behind Azure's outbound NAT, whose default
-    # idle timeout (~4 min) silently drops the connection before that phase
-    # ends. TCP keepalives (well under 4 min) keep the NAT mapping alive.
-    with psycopg.connect(
-        os.environ["CONNECTION_STRING"],
-        autocommit=True,
-        keepalives=1,
-        keepalives_idle=30,
-        keepalives_interval=10,
-        keepalives_count=5,
-    ) as conn:
-        if args.command == "schema":
+    if args.command == "schema":
+        with _connect() as conn:
             create_schema(conn)
-        elif args.command == "jixia":
-            # jixia runs each module with cwd=project_root, so the module file
-            # path must be absolute — a relative root would be resolved twice.
-            project = LeanProject(os.path.abspath(args.project_root))
-            prefixes = [parse_name(p) for p in args.prefixes.split(",")]
-            load_data(project, prefixes, conn)
-        elif args.command == "informal":
+    elif args.command == "jixia":
+        # jixia runs each module with cwd=project_root, so the module file
+        # path must be absolute — a relative root would be resolved twice.
+        project = LeanProject(os.path.abspath(args.project_root))
+        prefixes = [parse_name(p) for p in args.prefixes.split(",")]
+        # Run analysis (30-90+ min of pure CPU work) before opening the DB
+        # connection — GitHub Actions runners sit behind Azure's outbound NAT
+        # (~4 min idle timeout), and Heroku's own connection handling has also
+        # been observed to drop connections left idle that long. Connecting
+        # only once we're ready to write avoids relying on either surviving
+        # a long idle period.
+        analysis = run_analysis(project, prefixes)
+        with _connect() as conn:
+            load_data(project, analysis, conn)
+    elif args.command == "informal":
+        with _connect() as conn:
             generate_informal(
                 conn,
                 batch_size=args.batch_size,
                 limit_level=args.limit_level,
                 limit_num_per_level=args.limit_num_per_level,
             )
-        elif args.command == "vector-db":
+    elif args.command == "vector-db":
+        with _connect() as conn:
             create_vector_db(conn, os.environ["CHROMA_PATH"], batch_size=args.batch_size)

database/jixia_db.py

@@ -37,7 +37,29 @@ def _get_range(declaration: Declaration):
         return None
 
 
-def load_data(project: LeanProject, prefixes: list[LeanName], conn: Connection):
+def run_analysis(project: LeanProject, prefixes: list[LeanName]) -> list[tuple[Path, list[LeanName]]]:
+    # Pure CPU analysis, no DB connection involved; this can take 30-90+ min.
+    lean_sysroot = Path(os.environ["LEAN_SYSROOT"])
+    lean_src = lean_sysroot / "src" / "lean"
+    # Each jixia worker loads the full Mathlib environment (~2-3 GB), so the
+    # default thread count (CPUs + 4) can exhaust memory and get the process
+    # OOM-killed. Cap it via JIXIA_MAX_WORKERS in memory-constrained CI.
+    max_workers_env = os.environ.get("JIXIA_MAX_WORKERS")
+    max_workers = int(max_workers_env) if max_workers_env else None
+
+    module_lists = []
+    for d in project.root, lean_src:
+        results = project.batch_run_jixia(
+            base_dir=d,
+            prefixes=prefixes,
+            plugins=["module", "declaration", "symbol"],
+            max_workers=max_workers,
+        )
+        module_lists.append((d, [r[0] for r in results]))
+    return module_lists
+
+
+def load_data(project: LeanProject, analysis: list[tuple[Path, list[LeanName]]], conn: Connection):
     def load_module(data: Iterable[LeanName], base_dir: Path):
         values = []
         for m in data:
@@ -174,22 +196,8 @@ def topological_sort():
             """)
 
     with conn.cursor() as cursor:
-        lean_sysroot = Path(os.environ["LEAN_SYSROOT"])
-        lean_src = lean_sysroot / "src" / "lean"
-        # Each jixia worker loads the full Mathlib environment (~2-3 GB), so the
-        # default thread count (CPUs + 4) can exhaust memory and get the process
-        # OOM-killed. Cap it via JIXIA_MAX_WORKERS in memory-constrained CI.
-        max_workers_env = os.environ.get("JIXIA_MAX_WORKERS")
-        max_workers = int(max_workers_env) if max_workers_env else None
         all_modules = []
-        for d in project.root, lean_src:
-            results = project.batch_run_jixia(
-                base_dir=d,
-                prefixes=prefixes,
-                plugins=["module", "declaration", "symbol"],
-                max_workers=max_workers,
-            )
-            modules = [r[0] for r in results]
+        for d, modules in analysis:
             load_module(modules, d)
             all_modules += modules