feat: Remove long standing connections on sql state and catalog backends

Author: benjaminwilen

airbyte/caches/_catalog_backend.py

@@ -10,7 +10,6 @@
 
 import abc
 import json
-from typing import TYPE_CHECKING
 
 from sqlalchemy import Column, String
 from sqlalchemy.orm import Session, declarative_base
@@ -24,10 +23,7 @@
 )
 
 from airbyte.shared.catalog_providers import CatalogProvider
-
-
-if TYPE_CHECKING:
-    from sqlalchemy.engine import Engine
+from airbyte.shared.sql_processor import SqlConfig
 
 
 STREAMS_TABLE_NAME = "_airbyte_streams"
@@ -110,10 +106,10 @@ class SqlCatalogBackend(CatalogBackendBase):
 
     def __init__(
         self,
-        engine: Engine,
+        sql_config: SqlConfig,
         table_prefix: str,
     ) -> None:
-        self._engine: Engine = engine
+        self._sql_config = sql_config
         self._table_prefix = table_prefix
         self._ensure_internal_tables()
         self._full_catalog: ConfiguredAirbyteCatalog = ConfiguredAirbyteCatalog(
@@ -125,7 +121,7 @@ def __init__(
         self._source_catalogs: dict[str, ConfiguredAirbyteCatalog] = {}
 
     def _ensure_internal_tables(self) -> None:
-        engine = self._engine
+        engine = self._sql_config.get_sql_engine()
         SqlAlchemyModel.metadata.create_all(engine)  # type: ignore[attr-defined]
 
     def register_source(
@@ -181,7 +177,7 @@ def _save_catalog_info(
         incoming_stream_names: set[str],
     ) -> None:
         self._ensure_internal_tables()
-        engine = self._engine
+        engine = self._sql_config.get_sql_engine()
         with Session(engine) as session:
             # Delete and replace existing stream entries from the catalog cache
             table_name_entries_to_delete = [
@@ -217,7 +213,7 @@ def _fetch_streams_info(
 
         The `source_name` and `table_prefix` args are optional filters.
         """
-        engine = self._engine
+        engine = self._sql_config.get_sql_engine()
         with Session(engine) as session:
             # load all the streams
             streams: list[CachedStream] = session.query(CachedStream).all()

airbyte/caches/_state_backend.py

@@ -20,13 +20,12 @@
     StateBackendBase,
 )
 from airbyte.exceptions import PyAirbyteInputError, PyAirbyteInternalError
+from airbyte.shared.sql_processor import SqlConfig
 from airbyte.shared.state_providers import StaticInputState
 from airbyte.shared.state_writers import StateWriterBase
 
 
 if TYPE_CHECKING:
-    from sqlalchemy.engine import Engine
-
     from airbyte.shared.state_providers import StateProviderBase
 
 
@@ -136,7 +135,7 @@ def _write_state(
 
         self._state_backend._ensure_internal_tables()  # noqa: SLF001  # Non-public member access
         table_prefix = self._state_backend._table_prefix  # noqa: SLF001
-        engine = self._state_backend._engine  # noqa: SLF001
+        engine = self._state_backend._sql_config.get_sql_engine()  # noqa: SLF001
 
         # Calculate the new state model to write.
         new_state = (
@@ -187,17 +186,17 @@ class SqlStateBackend(StateBackendBase):
 
     def __init__(
         self,
-        engine: Engine,
+        sql_config: SqlConfig,
         table_prefix: str = "",
     ) -> None:
         """Initialize the state manager with a static catalog state."""
-        self._engine: Engine = engine
+        self._sql_config = sql_config
         self._table_prefix = table_prefix
         super().__init__()
 
     def _ensure_internal_tables(self) -> None:
         """Ensure the internal tables exist in the SQL database."""
-        engine = self._engine
+        engine = self._sql_config.get_sql_engine()
         SqlAlchemyModel.metadata.create_all(engine)  # type: ignore[attr-defined]
 
     def get_state_provider(
@@ -223,7 +222,7 @@ def get_state_provider(
         else:
             stream_state_model = CacheStreamStateModel
 
-        engine = self._engine
+        engine = self._sql_config.get_sql_engine()
         with Session(engine) as session:
             query = session.query(stream_state_model).filter(
                 stream_state_model.source_name == source_name

airbyte/caches/base.py

@@ -89,11 +89,11 @@ def __init__(self, **data: Any) -> None:  # noqa: ANN401
 
         # Initialize the catalog and state backends
         self._catalog_backend = SqlCatalogBackend(
-            engine=self.get_sql_engine(),
+            sql_config=self,
             table_prefix=self.table_prefix or "",
         )
         self._state_backend = SqlStateBackend(
-            engine=self.get_sql_engine(),
+            sql_config=self,
             table_prefix=self.table_prefix or "",
         )