Harden long biz_state collects: release DB during CLI and spool before flush.

Avoid idle Session disconnect after heavy timeout, and separate SSH collect from batched Postgres persist via disk spool.

Co-authored-by: Cursor <cursoragent@cursor.com>

tests/test_biz_state_collect_finalize.py

@@ -0,0 +1,232 @@
+"""Finalize / reconnect after long-collect DB disconnect."""
+
+from __future__ import annotations
+
+import unittest
+from unittest.mock import patch
+
+from sqlalchemy import create_engine
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.pool import StaticPool
+
+from netx_api.biz_state import collect_runner as runner
+from netx_api.db import Base
+from netx_api.models import BizStateBatch, BizStateTask
+
+
+class BizStateCollectFinalizeTests(unittest.TestCase):
+    def setUp(self) -> None:
+        # StaticPool: all sessions share one :memory: SQLite DB.
+        engine = create_engine(
+            "sqlite+pysqlite:///:memory:",
+            future=True,
+            connect_args={"check_same_thread": False},
+            poolclass=StaticPool,
+        )
+        TestingSession = sessionmaker(
+            bind=engine, autoflush=False, autocommit=False, expire_on_commit=False
+        )
+        Base.metadata.create_all(bind=engine)
+        self.Session = TestingSession
+        self.db = TestingSession()
+        self.task = BizStateTask(
+            id="t-finalize",
+            source="managed",
+            ne_id="ne1",
+            ne_name="PE1",
+            status="running",
+            collect_running=True,
+            interval_sec=300,
+        )
+        self.db.add(self.task)
+        self.batch = BizStateBatch(
+            id="b-finalize",
+            task_id="t-finalize",
+            status="running",
+            command_count=55,
+            row_count=6606,
+            message="",
+        )
+        self.db.add(self.batch)
+        self.db.commit()
+
+    def tearDown(self) -> None:
+        self.db.close()
+
+    def test_finalize_partial_after_heavy_timeout(self) -> None:
+        with patch.object(runner, "SessionLocal", self.Session):
+            with patch(
+                "netx_api.biz_state.compare_service.try_auto_compare_for_task",
+                return_value=None,
+            ):
+                status = runner._finalize_batch_status(
+                    batch_id="b-finalize",
+                    task_id="t-finalize",
+                    cmd_count=40,
+                    total_rows=1000,
+                    any_fail=True,
+                    any_ok=True,
+                    lane_errors=["RuntimeError: biz_state_heavy_timeout (2400s)"],
+                )
+        self.assertEqual(status, "partial")
+        self.db.expire_all()
+        batch = self.db.get(BizStateBatch, "b-finalize")
+        assert batch is not None
+        self.assertEqual(batch.status, "partial")
+        self.assertEqual(batch.command_count, 55)
+        self.assertEqual(batch.row_count, 6606)
+        self.assertIn("biz_state_heavy_timeout", batch.message or "")
+        self.assertIsNotNone(batch.ended_at)
+
+    def test_finalize_retries_once_on_operational_error(self) -> None:
+        calls = {"n": 0}
+        real_session = self.Session
+
+        class FlakySession:
+            """First commit raises OperationalError; subsequent sessions work."""
+
+            def __init__(self) -> None:
+                self._inner = real_session()
+                self._failed = False
+
+            def get(self, *args, **kwargs):
+                return self._inner.get(*args, **kwargs)
+
+            def commit(self) -> None:
+                calls["n"] += 1
+                if calls["n"] == 1:
+                    self._failed = True
+                    raise OperationalError(
+                        "UPDATE",
+                        {},
+                        Exception(
+                            "consuming input failed: server closed the connection unexpectedly"
+                        ),
+                    )
+                self._inner.commit()
+
+            def rollback(self) -> None:
+                try:
+                    self._inner.rollback()
+                except Exception:
+                    pass
+
+            def close(self) -> None:
+                self._inner.close()
+
+            def connection(self):
+                # Avoid invalidate() wiping StaticPool's only connection in tests.
+                raise RuntimeError("skip invalidate in test")
+
+            def add(self, *args, **kwargs):
+                return self._inner.add(*args, **kwargs)
+
+        def session_factory():
+            return FlakySession()
+
+        with patch.object(runner, "SessionLocal", session_factory):
+            with patch(
+                "netx_api.biz_state.compare_service.try_auto_compare_for_task",
+                return_value=None,
+            ):
+                status = runner._finalize_batch_status(
+                    batch_id="b-finalize",
+                    task_id="t-finalize",
+                    cmd_count=55,
+                    total_rows=6606,
+                    any_fail=True,
+                    any_ok=True,
+                    lane_errors=["RuntimeError: biz_state_heavy_timeout (2400s)"],
+                )
+        self.assertEqual(status, "partial")
+        self.assertEqual(calls["n"], 2)
+        self.db.expire_all()
+        batch = self.db.get(BizStateBatch, "b-finalize")
+        assert batch is not None
+        self.assertEqual(batch.status, "partial")
+        self.assertIn("biz_state_heavy_timeout", batch.message or "")
+
+    def test_fail_batch_retries_on_stale_connection(self) -> None:
+        calls = {"n": 0}
+        real_session = self.Session
+
+        class FlakySession:
+            def __init__(self) -> None:
+                self._inner = real_session()
+
+            def get(self, *args, **kwargs):
+                return self._inner.get(*args, **kwargs)
+
+            def commit(self) -> None:
+                calls["n"] += 1
+                if calls["n"] == 1:
+                    raise OperationalError(
+                        "UPDATE",
+                        {},
+                        Exception("server closed the connection unexpectedly"),
+                    )
+                self._inner.commit()
+
+            def rollback(self) -> None:
+                try:
+                    self._inner.rollback()
+                except Exception:
+                    pass
+
+            def close(self) -> None:
+                self._inner.close()
+
+            def connection(self):
+                raise RuntimeError("skip invalidate in test")
+
+        with patch.object(runner, "SessionLocal", FlakySession):
+            runner._fail_batch_status("b-finalize", "RuntimeError: boom")
+        self.assertEqual(calls["n"], 2)
+        self.db.expire_all()
+        batch = self.db.get(BizStateBatch, "b-finalize")
+        assert batch is not None
+        self.assertEqual(batch.status, "failed")
+        self.assertIn("boom", batch.message or "")
+
+    def test_run_db_with_reconnect_reraises_non_stale(self) -> None:
+        calls = {"n": 0}
+
+        class BoomSession:
+            def commit(self) -> None:
+                pass
+
+            def rollback(self) -> None:
+                pass
+
+            def close(self) -> None:
+                pass
+
+            def connection(self):
+                raise RuntimeError("no conn")
+
+        def fn(db) -> None:
+            calls["n"] += 1
+            raise ValueError("not a disconnect")
+
+        with patch.object(runner, "SessionLocal", BoomSession):
+            with self.assertRaises(ValueError):
+                runner._run_db_with_reconnect(fn, label="test")
+        self.assertEqual(calls["n"], 1)
+
+    def test_is_stale_db_connection(self) -> None:
+        self.assertTrue(
+            runner._is_stale_db_connection(
+                OperationalError("x", {}, Exception("server closed the connection"))
+            )
+        )
+        self.assertTrue(
+            runner._is_stale_db_connection(
+                RuntimeError("consuming input failed: server closed the connection unexpectedly")
+            )
+        )
+        self.assertFalse(runner._is_stale_db_connection(ValueError("nope")))
+
+
+if __name__ == "__main__":
+    unittest.main()