diff --git a/src/pyseekdb/client/admin_client.py b/src/pyseekdb/client/admin_client.py index aef416be..b3d2956b 100644 --- a/src/pyseekdb/client/admin_client.py +++ b/src/pyseekdb/client/admin_client.py @@ -156,6 +156,10 @@ def fork_database(self, source_name: str, destination_name: str, tenant: str = D """Proxy to server implementation""" return self._server.fork_database(source_name=source_name, destination_name=destination_name, tenant=tenant) + def close(self) -> None: + """Close the underlying client and release its resources.""" + self._server.close() + def __repr__(self): """Return the developer-readable representation.""" return f"" @@ -247,6 +251,10 @@ def count_collection(self) -> int: """Proxy to server implementation - collection operations only""" return self._server.count_collection() + def close(self) -> None: + """Close the underlying client and release its resources.""" + self._server.close() + def __repr__(self): """Return the developer-readable representation.""" return f"" diff --git a/src/pyseekdb/client/base_connection.py b/src/pyseekdb/client/base_connection.py index 0af978ee..3d3f47f2 100644 --- a/src/pyseekdb/client/base_connection.py +++ b/src/pyseekdb/client/base_connection.py @@ -32,6 +32,10 @@ def _cleanup(self): """Internal cleanup method to close connection and release resources""" pass + def close(self) -> None: + """Close the client connection and release owned resources.""" + self._cleanup() + @abstractmethod def _execute(self, sql: str) -> Any: """Execute SQL statement (basic functionality)""" @@ -56,13 +60,13 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): """Context manager support: automatic resource cleanup""" - self._cleanup() + self.close() def __del__(self): """Destructor: ensure connection is closed to prevent resource leaks""" try: - if hasattr(self, "_connection") and self.is_connected(): - self._cleanup() + if hasattr(self, "_connection"): + self.close() except Exception as exc: # Ignore all exceptions in destructor # Avoid issues during interpreter shutdown diff --git a/src/pyseekdb/client/client_seekdb_embedded.py b/src/pyseekdb/client/client_seekdb_embedded.py index fb714e40..42efc8ea 100644 --- a/src/pyseekdb/client/client_seekdb_embedded.py +++ b/src/pyseekdb/client/client_seekdb_embedded.py @@ -5,9 +5,13 @@ import logging import os +import threading from collections.abc import Sequence from typing import Any +import pymysql +from pymysql.cursors import DictCursor + # Try to import pylibseekdb - it may not be available on all platforms try: import pylibseekdb as seekdb # type: ignore[import-not-found] @@ -25,6 +29,64 @@ logger = logging.getLogger(__name__) +class _NativeEmbeddedBackend: + """Connect through pylibseekdb's native Python connection wrapper.""" + + supports_dbapi_cursor = False + + @staticmethod + def connect(instance: Any, database: str, connection_kwargs: dict[str, Any]) -> Any: + del connection_kwargs + if instance is not None: + return instance.connect(database=database, autocommit=True) + return seekdb.connect(database=database, autocommit=True) # type: ignore[union-attr] + + @staticmethod + def is_connection_open(connection: Any) -> bool: + return connection is not None + + +class _PyMySQLEmbeddedBackend: + """Connect to an owned SeekDB instance through its MySQL endpoint.""" + + supports_dbapi_cursor = True + + @staticmethod + def connect(instance: Any, database: str, connection_kwargs: dict[str, Any]) -> pymysql.Connection: + if instance is None: + raise RuntimeError("pylibseekdb returned no SeekdbInstance for a PyMySQL embedded connection") + + options = dict(instance.connection_options()) + kwargs = dict(connection_kwargs) + + # The endpoint and user belong to the lifecycle handle. Never let + # caller-provided values redirect this connection to another instance. + for key in ("host", "port", "unix_socket", "user"): + kwargs.pop(key, None) + kwargs.update(options) + + # Database selection remains caller-owned. BaseClient relies on a + # dictionary cursor and autocommit behavior matching RemoteServerClient. + kwargs.pop("db", None) + kwargs["database"] = database + kwargs.setdefault("charset", "utf8mb4") + kwargs["cursorclass"] = DictCursor + kwargs["autocommit"] = True + return pymysql.connect(**kwargs) + + @staticmethod + def is_connection_open(connection: Any) -> bool: + return connection is not None and bool(getattr(connection, "open", False)) + + +def _create_embedded_backend() -> _NativeEmbeddedBackend | _PyMySQLEmbeddedBackend: + """Select the safest connection backend from pylibseekdb's capabilities.""" + instance_type = getattr(seekdb, "SeekdbInstance", None) + if instance_type is not None and callable(getattr(instance_type, "connection_options", None)): + return _PyMySQLEmbeddedBackend() + return _NativeEmbeddedBackend() + + class SeekdbEmbeddedClient(BaseClient): """Embedded seekdb client (lazy connection) @@ -58,6 +120,10 @@ def __init__(self, path: str = "./seekdb.db", database: str = "test", **kwargs): raise ValueError(f"Path exists but is not a directory: {self.path}") self.database = database + self._connection_kwargs = dict(kwargs) + self._backend = _create_embedded_backend() + self._connection_lock = threading.RLock() + self._instance = None self._connection = None self._initialized = False @@ -65,41 +131,73 @@ def __init__(self, path: str = "./seekdb.db", database: str = "test", **kwargs): # ==================== Connection Management ==================== - def _ensure_connection(self) -> Any: # seekdb.Connection + def _ensure_connection(self) -> Any: """Ensure connection is established (internal method)""" - if not self._initialized: - # 1. open seekdb + with self._connection_lock: + if self._backend.is_connection_open(self._connection): + return self._connection + self._connection = None + + if not self._initialized: + try: + self._instance = seekdb.open(db_dir=self.path) # type: ignore[attr-defined] + logger.info(f"✅ seekdb opened: {self.path}") + except Exception as exc: + # pylibseekdb < 1.4 exposes only a process-wide module API. + # Keep that legacy API working, but never hide this error for + # the instance API where each db_dir has independent ownership. + if hasattr(seekdb, "SeekdbInstance") or "initialized twice" not in str(exc): + raise + logger.debug(f"seekdb already opened through the legacy module API: {exc}") + self._initialized = True + try: - seekdb.open(db_dir=self.path) # type: ignore[attr-defined] - logger.info(f"✅ seekdb opened: {self.path}") - except Exception as e: - if "initialized twice" not in str(e): - raise - logger.debug(f"seekdb already opened: {e}") - - self._initialized = True - - # 3. Create connection - if self._connection is None: - self._connection = seekdb.connect( # type: ignore[attr-defined] - database=self.database, autocommit=True - ) + connection = self._backend.connect(self._instance, self.database, self._connection_kwargs) + except Exception: + if self._instance is not None: + try: + self._instance.close() + except Exception as close_exc: + logger.warning("Failed to close seekdb instance after connection error: %s", close_exc) + finally: + self._instance = None + self._initialized = False + raise + + self._connection = connection logger.info(f"✅ Connected to database: {self.database}") + return self._connection - return self._connection - - def _cleanup(self): - """Internal cleanup method: close connection)""" - if self._connection is not None: - self._connection.close() + def _cleanup(self) -> None: + """Close the connection and its owning pylibseekdb instance.""" + with self._connection_lock: + connection = self._connection + instance = self._instance self._connection = None - logger.info(f"Connection closed: path={self.path}, database={self.database}") + self._instance = None + + # Legacy pylibseekdb has only a process-wide module instance. Leave + # its initialization state intact because another pyseekdb client may + # still be using it. The object API provides safe per-instance close. + if instance is not None: + self._initialized = False + + try: + if connection is not None: + connection.close() + finally: + if instance is not None: + instance.close() + + if connection is not None or instance is not None: + logger.info(f"Connection closed: path={self.path}, database={self.database}") def is_connected(self) -> bool: """Check connection status""" - return self._connection is not None and self._initialized + with self._connection_lock: + return self._initialized and self._backend.is_connection_open(self._connection) - def get_raw_connection(self) -> Any: # seekdb.Connection + def get_raw_connection(self) -> Any: """Get raw connection object""" return self._ensure_connection() @@ -109,28 +207,27 @@ def mode(self) -> str: return "SeekdbEmbeddedClient" def _use_context_manager_for_cursor(self) -> bool: - """ - Override to use try/finally instead of context manager for cursor - (seekdb embedded client doesn't support context manager) - """ - return False + """Use DB-API cursor contexts for PyMySQL, not for the legacy native cursor.""" + return self._backend.supports_dbapi_cursor def _execute_query_with_cursor( # noqa: C901 self, conn: Any, sql: str, params: list[Any], use_context_manager: bool = True ) -> list[dict[str, Any]]: """ - Execute SQL query and return normalized rows - Override base class to handle pyseekdb cursor which doesn't support parameterized queries + Execute SQL through DB-API for PyMySQL or adapt the legacy native cursor. Args: conn: Database connection sql: SQL query string with %s placeholders - params: Query parameters to embed in SQL - use_context_manager: Whether to use context manager (ignored for embedded client) + params: Query parameters + use_context_manager: Whether the selected cursor supports a context manager Returns: List of normalized row dictionaries """ + if self._backend.supports_dbapi_cursor: + return super()._execute_query_with_cursor(conn, sql, params, use_context_manager) + # pyseekdb.Cursor.execute() only accepts SQL string, not parameters # Embed parameters directly into SQL embedded_sql = render_sql_with_params(sql, params) diff --git a/tests/integration_tests/test_get_or_create_collection_multiprocess.py b/tests/integration_tests/test_get_or_create_collection_multiprocess.py index 5c511995..3b84defe 100644 --- a/tests/integration_tests/test_get_or_create_collection_multiprocess.py +++ b/tests/integration_tests/test_get_or_create_collection_multiprocess.py @@ -12,7 +12,6 @@ from __future__ import annotations import contextlib -import gc import importlib import importlib.metadata import multiprocessing as mp @@ -478,7 +477,7 @@ def mixed_in_thread(thread_id: int) -> dict[str, int]: output.put({"ok": False, "process_id": process_id, "error_type": type(exc).__name__, "error": str(exc)}) -def _build_client_config(mode: str) -> tuple[dict[str, Any], Path | None]: +def _build_client_config(mode: str) -> tuple[dict[str, Any], Path | None, Any]: """Build client config.""" database = f"test_mp_{uuid.uuid4().hex[:8]}" temp_db_path: Path | None = None @@ -512,16 +511,16 @@ def _build_client_config(mode: str) -> tuple[dict[str, Any], Path | None]: admin = _make_admin_client(client_config) admin.create_database(database) - del admin - gc.collect() - return client_config, temp_db_path + return client_config, temp_db_path, admin @pytest.fixture def multiprocess_db(_mode): """Multiprocess db.""" - client_config, temp_db_path = _build_client_config(_mode) + client_config, temp_db_path, admin = _build_client_config(_mode) yield client_config + with contextlib.suppress(Exception): + admin.close() if temp_db_path is not None: with contextlib.suppress(Exception): shutil.rmtree(temp_db_path, ignore_errors=True) @@ -545,6 +544,8 @@ def crud_collection(multiprocess_db): with contextlib.suppress(Exception): client.delete_collection(collection_name) + with contextlib.suppress(Exception): + client.close() def _seed_collection_rows( diff --git a/tests/unit_tests/test_embedded_client_lifecycle.py b/tests/unit_tests/test_embedded_client_lifecycle.py new file mode 100644 index 00000000..748d87fd --- /dev/null +++ b/tests/unit_tests/test_embedded_client_lifecycle.py @@ -0,0 +1,314 @@ +"""Unit tests for embedded pylibseekdb instance ownership.""" + +from __future__ import annotations + +import importlib +import sys +from types import ModuleType +from typing import Any + +import pytest + +from pyseekdb.client.admin_client import _AdminClientProxy, _ClientProxy + + +class _FakeConnection: + def __init__(self, path: str, events: list[tuple[Any, ...]]) -> None: + self.path = path + self._events = events + + def close(self) -> None: + self._events.append(("connection.close", self.path)) + + +class _FakePyMySQLCursor: + description = (("value",),) + + def __init__(self, events: list[tuple[Any, ...]]) -> None: + self._events = events + + def __enter__(self): + self._events.append(("cursor.enter",)) + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self._events.append(("cursor.exit",)) + + def execute(self, sql: str, params: list[Any]) -> None: + self._events.append(("cursor.execute", sql, params)) + + def fetchall(self) -> list[dict[str, int]]: + return [{"value": 7}] + + +class _FakePyMySQLConnection: + def __init__(self, events: list[tuple[Any, ...]]) -> None: + self._events = events + self.open = True + + def cursor(self) -> _FakePyMySQLCursor: + return _FakePyMySQLCursor(self._events) + + def close(self) -> None: + self._events.append(("pymysql.close",)) + self.open = False + + +class _FakeInstance: + def __init__(self, path: str, events: list[tuple[Any, ...]], fail_connect: bool = False) -> None: + self.path = path + self._events = events + self._fail_connect = fail_connect + + def connect(self, *, database: str, autocommit: bool) -> _FakeConnection: + self._events.append(("instance.connect", self.path, database, autocommit)) + if self._fail_connect: + raise RuntimeError("connect failed") + return _FakeConnection(self.path, self._events) + + def close(self) -> None: + self._events.append(("instance.close", self.path)) + + +class _FakeInstanceApi: + SeekdbInstance = _FakeInstance + + def __init__(self, *, fail_first_connect: bool = False) -> None: + self.events: list[tuple[Any, ...]] = [] + self._fail_first_connect = fail_first_connect + + def open(self, *, db_dir: str) -> _FakeInstance: + self.events.append(("open", db_dir)) + fail_connect = self._fail_first_connect + self._fail_first_connect = False + return _FakeInstance(db_dir, self.events, fail_connect=fail_connect) + + def connect(self, *, database: str, autocommit: bool) -> _FakeConnection: + raise AssertionError("the module-level connection must not be used by the instance API") + + +class _FakeOptionsInstance: + def __init__(self, path: str, events: list[tuple[Any, ...]]) -> None: + self.path = path + self._events = events + + def connection_options(self) -> dict[str, Any]: + self._events.append(("instance.connection_options", self.path)) + return {"user": "root", "unix_socket": f"{self.path}/run/sql.sock"} + + def connect(self, *, database: str, autocommit: bool) -> _FakeConnection: + raise AssertionError("the native connection must not be used when connection_options is available") + + def close(self) -> None: + self._events.append(("instance.close", self.path)) + + +class _FakeOptionsApi: + SeekdbInstance = _FakeOptionsInstance + + def __init__(self) -> None: + self.events: list[tuple[Any, ...]] = [] + + def open(self, *, db_dir: str) -> _FakeOptionsInstance: + self.events.append(("open", db_dir)) + return _FakeOptionsInstance(db_dir, self.events) + + def connection_options(self) -> dict[str, Any]: + raise AssertionError("the module-level connection options must not be used") + + +class _FakeLegacyApi: + def __init__(self) -> None: + self.events: list[tuple[Any, ...]] = [] + + def open(self, *, db_dir: str) -> None: + self.events.append(("open", db_dir)) + + def connect(self, *, database: str, autocommit: bool) -> _FakeConnection: + self.events.append(("module.connect", database, autocommit)) + return _FakeConnection("legacy", self.events) + + +@pytest.fixture +def embedded_module(monkeypatch: pytest.MonkeyPatch): + """Import the embedded module without loading the native pylibseekdb extension.""" + module_name = "pyseekdb.client.client_seekdb_embedded" + monkeypatch.setitem(sys.modules, "pylibseekdb", ModuleType("pylibseekdb")) + sys.modules.pop(module_name, None) + module = importlib.import_module(module_name) + yield module + sys.modules.pop(module_name, None) + + +def _install_fake_seekdb(embedded_module: Any, monkeypatch: pytest.MonkeyPatch, fake_seekdb: Any) -> None: + monkeypatch.setattr(embedded_module, "seekdb", fake_seekdb) + monkeypatch.setattr(embedded_module, "_PYLIBSEEKDB_AVAILABLE", True) + + +def test_clients_use_and_close_their_own_seekdb_instances(tmp_path, monkeypatch, embedded_module) -> None: + fake_seekdb = _FakeInstanceApi() + _install_fake_seekdb(embedded_module, monkeypatch, fake_seekdb) + path_a = str((tmp_path / "a").resolve()) + path_b = str((tmp_path / "b").resolve()) + + client_a = embedded_module.SeekdbEmbeddedClient(path=path_a, database="db_a") + client_b = embedded_module.SeekdbEmbeddedClient(path=path_b, database="db_b") + + assert client_a.get_raw_connection().path == path_a + assert client_b.get_raw_connection().path == path_b + + client_a.close() + assert client_b.is_connected() + client_b.close() + client_b.close() + + assert fake_seekdb.events == [ + ("open", path_a), + ("instance.connect", path_a, "db_a", True), + ("open", path_b), + ("instance.connect", path_b, "db_b", True), + ("connection.close", path_a), + ("instance.close", path_a), + ("connection.close", path_b), + ("instance.close", path_b), + ] + + +def test_connect_failure_releases_new_instance_and_allows_retry(tmp_path, monkeypatch, embedded_module) -> None: + fake_seekdb = _FakeInstanceApi(fail_first_connect=True) + _install_fake_seekdb(embedded_module, monkeypatch, fake_seekdb) + path = str((tmp_path / "db").resolve()) + client = embedded_module.SeekdbEmbeddedClient(path=path) + + with pytest.raises(RuntimeError, match="connect failed"): + client.get_raw_connection() + + assert client._instance is None + assert not client._initialized + assert not client.is_connected() + assert fake_seekdb.events[-1] == ("instance.close", path) + + assert client.get_raw_connection().path == path + client.close() + assert [event for event in fake_seekdb.events if event[0] == "open"] == [("open", path), ("open", path)] + + +def test_connection_options_backend_uses_pymysql_and_closes_in_lifecycle_order( + tmp_path, monkeypatch, embedded_module +) -> None: + fake_seekdb = _FakeOptionsApi() + _install_fake_seekdb(embedded_module, monkeypatch, fake_seekdb) + connect_calls: list[dict[str, Any]] = [] + + def fake_connect(**kwargs: Any) -> _FakePyMySQLConnection: + connect_calls.append(kwargs) + fake_seekdb.events.append(("pymysql.connect",)) + return _FakePyMySQLConnection(fake_seekdb.events) + + monkeypatch.setattr(embedded_module.pymysql, "connect", fake_connect) + path = str((tmp_path / "db").resolve()) + client = embedded_module.SeekdbEmbeddedClient( + path=path, + database="app", + host="wrong-host", + user="wrong-user", + read_timeout=5, + ) + + connection = client.get_raw_connection() + + assert isinstance(connection, _FakePyMySQLConnection) + assert client.mode == "SeekdbEmbeddedClient" + assert client.is_connected() + assert connect_calls == [ + { + "read_timeout": 5, + "user": "root", + "unix_socket": f"{path}/run/sql.sock", + "database": "app", + "charset": "utf8mb4", + "cursorclass": embedded_module.DictCursor, + "autocommit": True, + } + ] + + rows = client._execute_query_with_cursor(connection, "SELECT %s AS value", [7], True) + assert rows == [{"value": 7}] + + client.close() + client.close() + + assert fake_seekdb.events == [ + ("open", path), + ("instance.connection_options", path), + ("pymysql.connect",), + ("cursor.enter",), + ("cursor.execute", "SELECT %s AS value", [7]), + ("cursor.exit",), + ("pymysql.close",), + ("instance.close", path), + ] + + +def test_pymysql_connect_failure_releases_instance_and_allows_retry(tmp_path, monkeypatch, embedded_module) -> None: + fake_seekdb = _FakeOptionsApi() + _install_fake_seekdb(embedded_module, monkeypatch, fake_seekdb) + attempts = 0 + + def fake_connect(**kwargs: Any) -> _FakePyMySQLConnection: + nonlocal attempts + del kwargs + attempts += 1 + if attempts == 1: + raise RuntimeError("pymysql connect failed") + return _FakePyMySQLConnection(fake_seekdb.events) + + monkeypatch.setattr(embedded_module.pymysql, "connect", fake_connect) + path = str((tmp_path / "db").resolve()) + client = embedded_module.SeekdbEmbeddedClient(path=path) + + with pytest.raises(RuntimeError, match="pymysql connect failed"): + client.get_raw_connection() + + assert client._instance is None + assert not client._initialized + assert fake_seekdb.events[-1] == ("instance.close", path) + + assert isinstance(client.get_raw_connection(), _FakePyMySQLConnection) + client.close() + assert [event for event in fake_seekdb.events if event[0] == "open"] == [("open", path), ("open", path)] + + +def test_legacy_module_api_remains_compatible(tmp_path, monkeypatch, embedded_module) -> None: + fake_seekdb = _FakeLegacyApi() + _install_fake_seekdb(embedded_module, monkeypatch, fake_seekdb) + path = str((tmp_path / "legacy").resolve()) + client = embedded_module.SeekdbEmbeddedClient(path=path) + + client.get_raw_connection() + client.close() + client.get_raw_connection() + client.close() + + assert [event for event in fake_seekdb.events if event[0] == "open"] == [("open", path)] + assert [event for event in fake_seekdb.events if event[0] == "module.connect"] == [ + ("module.connect", "test", True), + ("module.connect", "test", True), + ] + + +@pytest.mark.parametrize("proxy_type", [_ClientProxy, _AdminClientProxy]) +def test_public_proxy_close_delegates_to_server(proxy_type) -> None: + class _Server: + def __init__(self) -> None: + self.close_count = 0 + + def close(self) -> None: + self.close_count += 1 + + server = _Server() + proxy = proxy_type(server) + + proxy.close() + + assert server.close_count == 1