From fe8316d533dbe70637f57b5227e0deeed625e828 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 30 Sep 2025 23:01:19 +0000 Subject: [PATCH 1/4] fix(cache): properly dispose DuckDB connections to prevent file locking - Add close() method to CacheBase to dispose all SQLAlchemy engines - Cache SQLAlchemy engine in SqlConfig and add dispose_engine() method - Implement context manager protocol (__enter__/__exit__) for clean resource management - Add __del__ for cleanup when cache is garbage collected - Add comprehensive tests demonstrating connection cleanup works Fixes #807 Co-Authored-By: AJ Steers --- airbyte/caches/base.py | 48 ++++++++++++ airbyte/shared/sql_processor.py | 31 +++++--- tests/integration_tests/test_duckdb_cache.py | 78 ++++++++++++++++++++ 3 files changed, 147 insertions(+), 10 deletions(-) diff --git a/airbyte/caches/base.py b/airbyte/caches/base.py index b0c3fd90c..030af11bd 100644 --- a/airbyte/caches/base.py +++ b/airbyte/caches/base.py @@ -3,9 +3,16 @@ from __future__ import annotations +import contextlib from pathlib import Path from typing import IO, TYPE_CHECKING, Any, ClassVar, Literal, final + +try: + from typing import Self +except ImportError: + from typing_extensions import Self + import pandas as pd import pyarrow as pa import pyarrow.dataset as ds @@ -28,6 +35,7 @@ if TYPE_CHECKING: from collections.abc import Iterator + from types import TracebackType from airbyte._message_iterators import AirbyteMessageIterator from airbyte.caches._state_backend_base import StateBackendBase @@ -108,6 +116,46 @@ def __init__(self, **data: Any) -> None: # noqa: ANN401 temp_file_cleanup=self.cleanup, ) + def close(self) -> None: + """Close all database connections and dispose of connection pools. + + This method ensures that all SQLAlchemy engines created by this cache + and its processors are properly disposed, releasing all database connections. + This is especially important for file-based databases like DuckDB, which + lock the database file until all connections are closed. + + This method is idempotent and can be called multiple times safely. + """ + if hasattr(self, "_read_processor") and self._read_processor is not None: + with contextlib.suppress(Exception): + self._read_processor.sql_config.dispose_engine() + + for backend in [self._catalog_backend, self._state_backend]: + if backend is not None and hasattr(backend, "_sql_config"): + with contextlib.suppress(Exception): + backend._sql_config.dispose_engine() # noqa: SLF001 + + with contextlib.suppress(Exception): + self.dispose_engine() + + def __enter__(self) -> Self: + """Enter context manager.""" + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit context manager and clean up resources.""" + self.close() + + def __del__(self) -> None: + """Clean up resources when cache is garbage collected.""" + with contextlib.suppress(Exception): + self.close() + @property def config_hash(self) -> str | None: """Return a hash of the cache configuration. diff --git a/airbyte/shared/sql_processor.py b/airbyte/shared/sql_processor.py index b8563e740..a602078ba 100644 --- a/airbyte/shared/sql_processor.py +++ b/airbyte/shared/sql_processor.py @@ -92,6 +92,9 @@ class SqlConfig(BaseModel, abc.ABC): table_prefix: str | None = "" """A prefix to add to created table names.""" + _engine: Engine | None = None + """Cached SQL engine instance.""" + @abc.abstractmethod def get_sql_alchemy_url(self) -> SecretString: """Returns a SQL Alchemy URL.""" @@ -133,16 +136,24 @@ def get_sql_alchemy_connect_args(self) -> dict[str, Any]: return {} def get_sql_engine(self) -> Engine: - """Return a new SQL engine to use.""" - return create_engine( - url=self.get_sql_alchemy_url(), - echo=DEBUG_MODE, - execution_options={ - "schema_translate_map": {None: self.schema_name}, - }, - future=True, - connect_args=self.get_sql_alchemy_connect_args(), - ) + """Return a cached SQL engine, creating it if necessary.""" + if self._engine is None: + self._engine = create_engine( + url=self.get_sql_alchemy_url(), + echo=DEBUG_MODE, + execution_options={ + "schema_translate_map": {None: self.schema_name}, + }, + future=True, + connect_args=self.get_sql_alchemy_connect_args(), + ) + return self._engine + + def dispose_engine(self) -> None: + """Dispose of the cached SQL engine and release all connections.""" + if self._engine is not None: + self._engine.dispose() + self._engine = None def get_vendor_client(self) -> object: """Return the vendor-specific client object. diff --git a/tests/integration_tests/test_duckdb_cache.py b/tests/integration_tests/test_duckdb_cache.py index fe4423142..70c9a1039 100644 --- a/tests/integration_tests/test_duckdb_cache.py +++ b/tests/integration_tests/test_duckdb_cache.py @@ -81,3 +81,81 @@ def duckdb_cache() -> Generator[DuckDBCache, None, None]: yield cache # TODO: Delete cache DB file after test is complete. return + + +def test_duckdb_connection_cleanup(tmp_path): + """Test that DuckDB connections are properly released after cache operations. + + This test verifies the fix for issue #807: DuckDB Cache Does Not Release Database Connections. + + Reproduction steps: + 1. Create a DuckDB cache and write data to it + 2. Delete the cache object or call close() + 3. Attempt to open a new connection to the same database file + 4. Verify the new connection succeeds without errors + """ + import duckdb + + db_path = tmp_path / "test_connection_cleanup.duckdb" + + cache = DuckDBCache(db_path=db_path) + + cache.processor._ensure_schema_exists() + + cache.processor._execute_sql( + "CREATE TABLE IF NOT EXISTS test_table (id INTEGER, name VARCHAR)" + ) + cache.processor._execute_sql("INSERT INTO test_table VALUES (1, 'test')") + + cache.close() + + try: + conn = duckdb.connect(str(db_path)) + result = conn.execute("SELECT * FROM main.test_table").fetchall() + assert len(result) == 1 + assert result[0] == (1, "test") + conn.close() + except Exception as e: + pytest.fail(f"Failed to open new connection after cache cleanup: {e}") + + +def test_duckdb_context_manager_cleanup(tmp_path): + """Test that DuckDB connections are cleaned up when using cache as context manager.""" + import duckdb + + db_path = tmp_path / "test_context_manager.duckdb" + + with DuckDBCache(db_path=db_path) as cache: + cache.processor._execute_sql( + "CREATE TABLE IF NOT EXISTS test_table (id INTEGER)" + ) + cache.processor._execute_sql("INSERT INTO test_table VALUES (1)") + + conn = duckdb.connect(str(db_path)) + result = conn.execute("SELECT * FROM main.test_table").fetchall() + assert len(result) == 1 + conn.close() + + +def test_duckdb_del_cleanup(tmp_path): + """Test that DuckDB connections are cleaned up via __del__ when cache is garbage collected.""" + import duckdb + import gc + + db_path = tmp_path / "test_del_cleanup.duckdb" + + def create_and_use_cache(): + cache = DuckDBCache(db_path=db_path) + cache.processor._execute_sql( + "CREATE TABLE IF NOT EXISTS test_table (id INTEGER)" + ) + cache.processor._execute_sql("INSERT INTO test_table VALUES (1)") + + create_and_use_cache() + + gc.collect() + + conn = duckdb.connect(str(db_path)) + result = conn.execute("SELECT * FROM main.test_table").fetchall() + assert len(result) == 1 + conn.close() From 7b8979a7cd75331d2a4c1502420eefc4e76dfac1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 30 Sep 2025 23:27:56 +0000 Subject: [PATCH 2/4] fix(lint): resolve MyPy and Ruff linting errors - Use typing_extensions.Self directly instead of try-except pattern for MyPy compatibility - Add noqa comment for PLR0904 (too many public methods) since adding close() pushed class over limit - Let Ruff auto-fix import ordering Co-Authored-By: AJ Steers --- airbyte/caches/base.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/airbyte/caches/base.py b/airbyte/caches/base.py index 030af11bd..9b6e70bd5 100644 --- a/airbyte/caches/base.py +++ b/airbyte/caches/base.py @@ -7,18 +7,13 @@ from pathlib import Path from typing import IO, TYPE_CHECKING, Any, ClassVar, Literal, final - -try: - from typing import Self -except ImportError: - from typing_extensions import Self - import pandas as pd import pyarrow as pa import pyarrow.dataset as ds from pydantic import Field, PrivateAttr from sqlalchemy import exc as sqlalchemy_exc from sqlalchemy import text +from typing_extensions import Self from airbyte_protocol.models import ConfiguredAirbyteCatalog @@ -47,7 +42,7 @@ from airbyte.strategies import WriteStrategy -class CacheBase(SqlConfig, AirbyteWriterInterface): +class CacheBase(SqlConfig, AirbyteWriterInterface): # noqa: PLR0904 """Base configuration for a cache. Caches inherit from the matching `SqlConfig` class, which provides the SQL config settings From e4c0290d38201b0e3552daa54ccb0b2c439d5a73 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 30 Sep 2025 23:54:00 +0000 Subject: [PATCH 3/4] fix(cache): allow exceptions to propagate in close() method - Remove contextlib.suppress from close() to let callers know about failures - Keep exception suppression only in __del__ for safe garbage collection - Add docstring explaining exception propagation behavior Addresses code review feedback from coderabbitai[bot] Co-Authored-By: AJ Steers --- airbyte/caches/base.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/airbyte/caches/base.py b/airbyte/caches/base.py index 9b6e70bd5..7afe8a0e3 100644 --- a/airbyte/caches/base.py +++ b/airbyte/caches/base.py @@ -120,18 +120,19 @@ def close(self) -> None: lock the database file until all connections are closed. This method is idempotent and can be called multiple times safely. + + Raises: + Exception: If any engine disposal fails, the exception will propagate + to the caller. This ensures callers are aware of cleanup failures. """ if hasattr(self, "_read_processor") and self._read_processor is not None: - with contextlib.suppress(Exception): - self._read_processor.sql_config.dispose_engine() + self._read_processor.sql_config.dispose_engine() for backend in [self._catalog_backend, self._state_backend]: if backend is not None and hasattr(backend, "_sql_config"): - with contextlib.suppress(Exception): - backend._sql_config.dispose_engine() # noqa: SLF001 + backend._sql_config.dispose_engine() # noqa: SLF001 - with contextlib.suppress(Exception): - self.dispose_engine() + self.dispose_engine() def __enter__(self) -> Self: """Enter context manager.""" From c1bd2f85b3d5dc4840bf406cbd05c559d09ffb96 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 1 Oct 2025 00:55:17 +0000 Subject: [PATCH 4/4] refactor(cache): remove hasattr() checks by ensuring attributes exist - Add _sql_config declaration to CatalogBackendBase and StateBackendBase - Refactor close() to access backends individually for better type inference - This makes the code cleaner and more Pythonic while maintaining identical behavior All tests pass and MyPy/Ruff checks pass. Co-Authored-By: AJ Steers --- airbyte/caches/_catalog_backend.py | 2 ++ airbyte/caches/_state_backend_base.py | 3 +++ airbyte/caches/base.py | 10 ++++++---- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/airbyte/caches/_catalog_backend.py b/airbyte/caches/_catalog_backend.py index 5b7494649..145f59490 100644 --- a/airbyte/caches/_catalog_backend.py +++ b/airbyte/caches/_catalog_backend.py @@ -55,6 +55,8 @@ class CatalogBackendBase(abc.ABC): - The JSON schema for each stream """ + _sql_config: SqlConfig + # Abstract implementations @abc.abstractmethod diff --git a/airbyte/caches/_state_backend_base.py b/airbyte/caches/_state_backend_base.py index 443629711..a1abfe319 100644 --- a/airbyte/caches/_state_backend_base.py +++ b/airbyte/caches/_state_backend_base.py @@ -13,6 +13,7 @@ AirbyteStreamState, ) + from airbyte.shared.sql_processor import SqlConfig from airbyte.shared.state_providers import StateProviderBase from airbyte.shared.state_writers import StateWriterBase @@ -24,6 +25,8 @@ class StateBackendBase(abc.ABC): `StateProvider` objects, which are paired to a specific source and table prefix. """ + _sql_config: SqlConfig + def __init__(self) -> None: """Initialize the state manager with a static catalog state.""" self._state_artifacts: list[AirbyteStreamState] | None = None diff --git a/airbyte/caches/base.py b/airbyte/caches/base.py index 7afe8a0e3..248451c23 100644 --- a/airbyte/caches/base.py +++ b/airbyte/caches/base.py @@ -125,12 +125,14 @@ def close(self) -> None: Exception: If any engine disposal fails, the exception will propagate to the caller. This ensures callers are aware of cleanup failures. """ - if hasattr(self, "_read_processor") and self._read_processor is not None: + if self._read_processor is not None: self._read_processor.sql_config.dispose_engine() - for backend in [self._catalog_backend, self._state_backend]: - if backend is not None and hasattr(backend, "_sql_config"): - backend._sql_config.dispose_engine() # noqa: SLF001 + if self._catalog_backend is not None: + self._catalog_backend._sql_config.dispose_engine() # noqa: SLF001 + + if self._state_backend is not None: + self._state_backend._sql_config.dispose_engine() # noqa: SLF001 self.dispose_engine()