fix: harden SQL connection lifecycle

grand/backends/_sqlbackend.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from threading import RLock
 from typing import Hashable, Generator
 import time
 
@@ -13,10 +14,27 @@
 _DEFAULT_SQL_STR_LEN = 64
 
 
+class _LockedConnection:
+    def __init__(self, connection, lock):
+        self._connection = connection
+        self._lock = lock
+
+    def execute(self, *args, **kwargs):
+        with self._lock:
+            return self._connection.execute(*args, **kwargs)
+
+    def __getattr__(self, name):
+        return getattr(self._connection, name)
+
+
 class SQLBackend(Backend):
     """
     A graph datastore that uses a SQL-like store for persistance and queries.
 
+    Operations on one backend instance are serialized because SQLAlchemy
+    connections cannot be used concurrently. Use separate backend instances
+    when parallel database operations are required.
+
     """
 
     def __init__(
@@ -57,7 +75,9 @@ def __init__(
 
         sqlalchemy_kwargs = sqlalchemy_kwargs or {}
         self._engine = sqlalchemy.create_engine(db_url, **sqlalchemy_kwargs)
-        self._connection = self._engine.connect()
+        self._lock = RLock()
+        self._connection = _LockedConnection(self._engine.connect(), self._lock)
+        self._closed = False
         self._transaction_depth = 0
         self._metadata = sqlalchemy.MetaData()
 
@@ -106,31 +126,46 @@ def __init__(
 
     @contextmanager
     def _mutation(self):
-        if self._transaction_depth:
-            yield
-            return
-        try:
-            yield
-            self._connection.commit()
-        except Exception:
-            self._connection.rollback()
-            raise
+        with self._lock:
+            self._ensure_open()
+            if self._transaction_depth:
+                yield
+                return
+            try:
+                yield
+                self._connection.commit()
+            except Exception:
+                self._connection.rollback()
+                raise
 
     @contextmanager
     def transaction(self):
         """Group multiple mutations into one atomic commit."""
-        outermost = self._transaction_depth == 0
-        self._transaction_depth += 1
-        try:
-            yield self
-            if outermost:
-                self._connection.commit()
-        except Exception:
-            if outermost:
-                self._connection.rollback()
-            raise
-        finally:
-            self._transaction_depth -= 1
+        with self._lock:
+            self._ensure_open()
+            outermost = self._transaction_depth == 0
+            self._transaction_depth += 1
+            try:
+                yield self
+                if outermost:
+                    self._connection.commit()
+            except Exception:
+                if outermost:
+                    self._connection.rollback()
+                raise
+            finally:
+                self._transaction_depth -= 1
+
+    def _ensure_open(self):
+        if self._closed:
+            raise RuntimeError("SQLBackend is closed")
+
+    def __enter__(self):
+        self._ensure_open()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
 
     def is_directed(self) -> bool:
         """
@@ -791,8 +826,17 @@ def ingest_from_edgelist_dataframe(
         }
 
     def commit(self):
-        if self._connection.in_transaction():
-            self._connection.commit()
+        with self._lock:
+            self._ensure_open()
+            if self._connection.in_transaction():
+                self._connection.commit()
 
     def close(self):
-        self._connection.close()
+        with self._lock:
+            if self._closed:
+                return
+            if self._connection.in_transaction():
+                self._connection.rollback()
+            self._connection.close()
+            self._engine.dispose()
+            self._closed = True

grand/backends/test_sql_transactions.py

@@ -1,4 +1,6 @@
 import pytest
+from concurrent.futures import ThreadPoolExecutor
+from unittest.mock import Mock
 
 sqlalchemy = pytest.importorskip("sqlalchemy")
 
@@ -134,3 +136,41 @@ def fail_edge_insert(statement, *args, **kwargs):
     assert not backend.has_node("A")
     assert not backend.has_node("B")
     backend.close()
+
+
+def test_context_manager_closes_connection_and_disposes_engine(tmp_path):
+    backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}")
+    dispose = backend._engine.dispose
+    backend._engine.dispose = Mock(wraps=dispose)
+
+    with backend as entered:
+        assert entered is backend
+        backend.add_node("A", {})
+
+    assert backend._closed
+    assert backend._connection.closed
+    backend._engine.dispose.assert_called_once_with()
+    with pytest.raises(RuntimeError, match="closed"):
+        backend.add_node("B", {})
+
+
+def test_close_is_idempotent(tmp_path):
+    backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}")
+
+    backend.close()
+    backend.close()
+
+    assert backend._closed
+
+
+def test_concurrent_mutations_do_not_share_connection_simultaneously(tmp_path):
+    backend = SQLBackend(
+        db_url=f"sqlite:///{tmp_path / 'graph.db'}",
+        sqlalchemy_kwargs={"connect_args": {"check_same_thread": False}},
+    )
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        list(executor.map(lambda node: backend.add_node(node, {}), range(20)))
+
+    assert backend.get_node_count() == 20
+    backend.close()