From b977bf6725dfb286318b4392e9a3dd391c03ff4f Mon Sep 17 00:00:00 2001
From: Jiang <huarch97@outlook.com>
Date: Tue, 21 Jul 2026 11:26:21 +0800
Subject: [PATCH] fix(db): validate cached project connections

---
 app/native/wndb/connection.py      | 16 +++++++--
 tests/unit/test_wndb_connection.py | 57 ++++++++++++++++++++++++++----
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/app/native/wndb/connection.py b/app/native/wndb/connection.py
index 7dd5de2..c8d1a67 100644
--- a/app/native/wndb/connection.py
+++ b/app/native/wndb/connection.py
@@ -20,6 +20,17 @@ def _close_connection(connection: pg.Connection) -> None:
         connection.close()
 
 
+def _is_healthy(connection: pg.Connection) -> bool:
+    if _is_closed(connection):
+        return False
+    try:
+        with connection.cursor() as cur:
+            cur.execute("SELECT 1")
+    except pg.Error:
+        return False
+    return True
+
+
 def _get_project_lock(name: str) -> RLock:
     with _registry_lock:
         lock = _project_locks.get(name)
@@ -32,7 +43,7 @@ def _get_project_lock(name: str) -> RLock:
 def open_connection(name: str) -> pg.Connection:
     with _get_project_lock(name):
         connection = g_conn_dict.get(name)
-        if connection is None or _is_closed(connection):
+        if connection is None or not _is_healthy(connection):
             if connection is not None:
                 _close_connection(connection)
             connection = pg.connect(
@@ -47,8 +58,9 @@ def is_connection_open(name: str) -> bool:
         connection = g_conn_dict.get(name)
         if connection is None:
             return False
-        if _is_closed(connection):
+        if not _is_healthy(connection):
             del g_conn_dict[name]
+            _close_connection(connection)
             return False
         return True
 
diff --git a/tests/unit/test_wndb_connection.py b/tests/unit/test_wndb_connection.py
index ea68d90..b8fc6f7 100644
--- a/tests/unit/test_wndb_connection.py
+++ b/tests/unit/test_wndb_connection.py
@@ -6,9 +6,8 @@ from app.native.wndb import project
 
 
 class _FakeCursor:
-    def __init__(self, rows):
-        self.rows = rows
-        self.executed = []
+    def __init__(self, connection):
+        self.connection = connection
 
     def __enter__(self):
         return self
@@ -17,22 +16,26 @@ class _FakeCursor:
         return False
 
     def execute(self, sql):
-        self.executed.append(sql)
+        self.connection.executed.append(sql)
+        if self.connection.fail_ping and sql == "SELECT 1":
+            raise connection.pg.OperationalError("server closed the connection")
 
     def fetchall(self):
-        return self.rows
+        return self.connection.rows
 
 
 class _FakeConnection:
-    def __init__(self, rows=None, *, closed=False):
+    def __init__(self, rows=None, *, closed=False, fail_ping=False):
         self.rows = list(rows or [])
         self.closed = closed
+        self.fail_ping = fail_ping
+        self.executed = []
         self.close_calls = 0
 
     def cursor(self, row_factory=None):
         if self.closed:
             raise RuntimeError("the connection is closed")
-        return _FakeCursor(self.rows)
+        return _FakeCursor(self)
 
     def close(self):
         self.close_calls += 1
@@ -55,6 +58,19 @@ def test_is_project_open_drops_closed_cached_connection():
     assert "fengyang" not in connection.g_conn_dict
 
 
+def test_open_connection_reuses_healthy_cached_connection(monkeypatch):
+    cached = _FakeConnection()
+    connection.g_conn_dict["fengyang"] = cached
+
+    def fail_connect(*, conninfo, autocommit):
+        raise AssertionError("cached connection should be reused")
+
+    monkeypatch.setattr(connection.pg, "connect", fail_connect)
+
+    assert connection.open_connection("fengyang") is cached
+    assert cached.executed == ["SELECT 1"]
+
+
 def test_read_all_reopens_closed_cached_connection(monkeypatch):
     stale = _FakeConnection(closed=True)
     fresh = _FakeConnection(rows=[{"key": "DURATION", "value": "01:00:00"}])
@@ -76,3 +92,30 @@ def test_read_all_reopens_closed_cached_connection(monkeypatch):
     assert rows == [{"key": "DURATION", "value": "01:00:00"}]
     assert opened == [("dbname=fengyang", True)]
     assert connection.g_conn_dict["fengyang"] is fresh
+    assert fresh.executed == ["select * from times"]
+
+
+def test_read_all_reopens_cached_connection_when_health_check_fails(monkeypatch):
+    stale = _FakeConnection(fail_ping=True)
+    fresh = _FakeConnection(rows=[{"scheme_name": "base"}])
+    connection.g_conn_dict["fengyang"] = stale
+
+    opened = []
+
+    def fake_connect(*, conninfo, autocommit):
+        opened.append((conninfo, autocommit))
+        return fresh
+
+    monkeypatch.setattr(connection.pg, "connect", fake_connect)
+    monkeypatch.setattr(
+        connection, "get_pgconn_string", lambda db_name: f"dbname={db_name}"
+    )
+
+    rows = database.read_all("fengyang", "select * from scheme_list")
+
+    assert rows == [{"scheme_name": "base"}]
+    assert stale.executed == ["SELECT 1"]
+    assert stale.close_calls == 1
+    assert opened == [("dbname=fengyang", True)]
+    assert connection.g_conn_dict["fengyang"] is fresh
+    assert fresh.executed == ["select * from scheme_list"]