diff --git a/airbyte/caches/_catalog_backend.py b/airbyte/caches/_catalog_backend.py index ffd3c8c78..5b7494649 100644 --- a/airbyte/caches/_catalog_backend.py +++ b/airbyte/caches/_catalog_backend.py @@ -27,7 +27,7 @@ if TYPE_CHECKING: - from sqlalchemy.engine import Engine + from airbyte.shared.sql_processor import SqlConfig STREAMS_TABLE_NAME = "_airbyte_streams" @@ -110,10 +110,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 +125,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 +181,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 +217,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() diff --git a/airbyte/caches/_state_backend.py b/airbyte/caches/_state_backend.py index 07cffe389..7fef5bbfc 100644 --- a/airbyte/caches/_state_backend.py +++ b/airbyte/caches/_state_backend.py @@ -25,8 +25,7 @@ if TYPE_CHECKING: - from sqlalchemy.engine import Engine - + from airbyte.shared.sql_processor import SqlConfig 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 diff --git a/airbyte/caches/base.py b/airbyte/caches/base.py index 1bfab644a..12d4a3ad9 100644 --- a/airbyte/caches/base.py +++ b/airbyte/caches/base.py @@ -90,11 +90,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 "", )