airbyte-lib: Stream state (#34778)

Co-authored-by: Aaron Steers <[email protected]>

Author: 2 people

airbyte-lib/airbyte_lib/_file_writers/base.py

@@ -18,6 +18,10 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
+    from airbyte_protocol.models import (
+        AirbyteStateMessage,
+    )
+
 
 DEFAULT_BATCH_SIZE = 10000
 
@@ -109,3 +113,14 @@ def cleanup_batch(
         Subclasses should override `_cleanup_batch` instead.
         """
         self._cleanup_batch(stream_name, batch_id, batch_handle)
+
+    @overrides
+    def _finalize_state_messages(
+        self,
+        stream_name: str,
+        state_messages: list[AirbyteStateMessage],
+    ) -> None:
+        """
+        State messages are not used in file writers, so this method is a no-op.
+        """
+        pass

airbyte-lib/airbyte_lib/_processors.py

@@ -78,6 +78,9 @@ def __init__(
             )
             raise TypeError(err_msg)
 
+        self.source_catalog: ConfiguredAirbyteCatalog | None = None
+        self._source_name: str | None = None
+
         self._pending_batches: dict[str, dict[str, Any]] = defaultdict(lambda: {}, {})
         self._finalized_batches: dict[str, dict[str, Any]] = defaultdict(lambda: {}, {})
 
@@ -301,6 +304,16 @@ def _finalize_batches(
 
             return batches_to_finalize
 
+    @abc.abstractmethod
+    def _finalize_state_messages(
+        self,
+        stream_name: str,
+        state_messages: list[AirbyteStateMessage],
+    ) -> None:
+        """Handle state messages.
+        Might be a no-op if the processor doesn't handle incremental state."""
+        pass
+
     @final
     @contextlib.contextmanager
     def _finalizing_batches(
@@ -318,6 +331,7 @@ def _finalizing_batches(
 
         progress.log_batches_finalizing(stream_name, len(batches_to_finalize))
         yield batches_to_finalize
+        self._finalize_state_messages(stream_name, state_messages_to_finalize)
         progress.log_batches_finalized(stream_name, len(batches_to_finalize))
 
         self._finalized_batches[stream_name].update(batches_to_finalize)

airbyte-lib/airbyte_lib/caches/_catalog_manager.py

@@ -6,11 +6,13 @@
 import json
 from typing import TYPE_CHECKING, Callable
 
-from sqlalchemy import Column, String
+from sqlalchemy import Column, DateTime, String
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session
+from sqlalchemy.sql import func
 
 from airbyte_protocol.models import (
+    AirbyteStateMessage,
     AirbyteStream,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
@@ -25,6 +27,9 @@
     from sqlalchemy.engine import Engine
 
 STREAMS_TABLE_NAME = "_airbytelib_streams"
+STATE_TABLE_NAME = "_airbytelib_state"
+
+GLOBAL_STATE_STREAM_NAMES = ["_GLOBAL", "_LEGACY"]
 
 Base = declarative_base()
 
@@ -38,7 +43,24 @@ class CachedStream(Base):  # type: ignore[valid-type,misc]
     catalog_metadata = Column(String)
 
 
+class StreamState(Base):  # type: ignore[valid-type,misc]
+    __tablename__ = STATE_TABLE_NAME
+
+    source_name = Column(String)
+    stream_name = Column(String)
+    table_name = Column(String, primary_key=True)
+    state_json = Column(String)
+    last_updated = Column(DateTime(timezone=True), onupdate=func.now(), default=func.now())
+
+
 class CatalogManager:
+    """
+    A class to manage the stream catalog of data synced to a cache:
+    * What streams exist and to what tables they map
+    * The JSON schema for each stream
+    * The state of each stream if available
+    """
+
     def __init__(
         self,
         engine: Engine,
@@ -68,6 +90,56 @@ def _ensure_internal_tables(self) -> None:
         engine = self._engine
         Base.metadata.create_all(engine)
 
+    def save_state(
+        self,
+        source_name: str,
+        state: AirbyteStateMessage,
+        stream_name: str,
+    ) -> None:
+        self._ensure_internal_tables()
+        engine = self._engine
+        with Session(engine) as session:
+            session.query(StreamState).filter(
+                StreamState.table_name == self._table_name_resolver(stream_name)
+            ).delete()
+            session.commit()
+            session.add(
+                StreamState(
+                    source_name=source_name,
+                    stream_name=stream_name,
+                    table_name=self._table_name_resolver(stream_name),
+                    state_json=state.json(),
+                )
+            )
+            session.commit()
+
+    def get_state(
+        self,
+        source_name: str,
+        streams: list[str],
+    ) -> list[dict] | None:
+        self._ensure_internal_tables()
+        engine = self._engine
+        with Session(engine) as session:
+            states = (
+                session.query(StreamState)
+                .filter(
+                    StreamState.source_name == source_name,
+                    StreamState.stream_name.in_([*streams, *GLOBAL_STATE_STREAM_NAMES]),
+                )
+                .all()
+            )
+            if not states:
+                return None
+            # Only return the states if the table name matches what the current cache
+            # would generate. Otherwise consider it part of a different cache.
+            states = [
+                state
+                for state in states
+                if state.table_name == self._table_name_resolver(state.stream_name)
+            ]
+            return [json.loads(state.state_json) for state in states]
+
     def register_source(
         self,
         source_name: str,

airbyte-lib/airbyte_lib/caches/base.py

@@ -46,6 +46,7 @@
     from sqlalchemy.sql.base import Executable
 
     from airbyte_protocol.models import (
+        AirbyteStateMessage,
         ConfiguredAirbyteCatalog,
     )
 
@@ -561,6 +562,36 @@ def _finalize_batches(
             # Return the batch handles as measure of work completed.
             return batches_to_finalize
 
+    @overrides
+    def _finalize_state_messages(
+        self,
+        stream_name: str,
+        state_messages: list[AirbyteStateMessage],
+    ) -> None:
+        """Handle state messages by passing them to the catalog manager."""
+        if not self._catalog_manager:
+            raise exc.AirbyteLibInternalError(
+                message="Catalog manager should exist but does not.",
+            )
+        if state_messages and self._source_name:
+            self._catalog_manager.save_state(
+                source_name=self._source_name,
+                stream_name=stream_name,
+                state=state_messages[-1],
+            )
+
+    def get_state(self) -> list[dict]:
+        """Return the current state of the source."""
+        if not self._source_name:
+            return []
+        if not self._catalog_manager:
+            raise exc.AirbyteLibInternalError(
+                message="Catalog manager should exist but does not.",
+            )
+        return (
+            self._catalog_manager.get_state(self._source_name, list(self._streams_with_data)) or []
+        )
+
     def _execute_sql(self, sql: str | TextClause | Executable) -> CursorResult:
         """Execute the given SQL statement."""
         if isinstance(sql, str):
@@ -881,6 +912,7 @@ def register_source(
 
         This method is called by the source when it is initialized.
         """
+        self._source_name = source_name
         self._ensure_schema_exists()
         super().register_source(
             source_name,

airbyte-lib/airbyte_lib/source.py

@@ -13,6 +13,7 @@
 from airbyte_protocol.models import (
     AirbyteCatalog,
     AirbyteMessage,
+    AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
     ConnectorSpecification,
@@ -46,7 +47,7 @@
 
 
 @contextmanager
-def as_temp_files(files_contents: list[dict]) -> Generator[list[str], Any, None]:
+def as_temp_files(files_contents: list[Any]) -> Generator[list[str], Any, None]:
     """Write the given contents to temporary files and yield the file paths as strings."""
     temp_files: list[Any] = []
     try:
@@ -140,6 +141,10 @@ def set_config(
 
         self._config_dict = config
 
+    def get_config(self) -> dict[str, Any]:
+        """Get the config for the connector."""
+        return self._config
+
     @property
     def _config(self) -> dict[str, Any]:
         if self._config_dict is None:
@@ -252,7 +257,7 @@ def configured_catalog(self) -> ConfiguredAirbyteCatalog:
                 # TODO: Set sync modes and primary key to a sensible adaptive default
                 ConfiguredAirbyteStream(
                     stream=stream,
-                    sync_mode=SyncMode.full_refresh,
+                    sync_mode=SyncMode.incremental,
                     destination_sync_mode=DestinationSyncMode.overwrite,
                     primary_key=stream.source_defined_primary_key,
                 )
@@ -309,7 +314,6 @@ def _with_missing_columns(records: Iterable[dict[str, Any]]) -> Iterator[dict[st
                 self._read_with_catalog(
                     streaming_cache_info,
                     configured_catalog,
-                    force_full_refresh=True,  # Always full refresh when skipping the cache
                 ),
             )
         )
@@ -360,8 +364,7 @@ def uninstall(self) -> None:
     def _read(
         self,
         cache_info: CacheTelemetryInfo,
-        *,
-        force_full_refresh: bool,
+        state: list[AirbyteStateMessage] | None = None,
     ) -> Iterable[AirbyteMessage]:
         """
         Call read on the connector.
@@ -379,15 +382,14 @@ def _read(
         yield from self._read_with_catalog(
             cache_info,
             catalog=self.configured_catalog,
-            force_full_refresh=force_full_refresh,
+            state=state,
         )
 
     def _read_with_catalog(
         self,
         cache_info: CacheTelemetryInfo,
         catalog: ConfiguredAirbyteCatalog,
-        *,
-        force_full_refresh: bool,
+        state: list[AirbyteStateMessage] | None = None,
     ) -> Iterator[AirbyteMessage]:
         """Call read on the connector.
 
@@ -397,21 +399,27 @@ def _read_with_catalog(
         * Listen to the messages and return the AirbyteRecordMessages that come along.
         * Send out telemetry on the performed sync (with information about which source was used and
           the type of the cache)
-
-        TODO: When we add support for incremental syncs, we should only send `--state <state_file>`
-              if force_full_refresh is False.
         """
-        _ = force_full_refresh  # TODO: Use this decide whether to send `--state <state_file>`
         source_tracking_information = self.executor.get_telemetry_info()
         send_telemetry(source_tracking_information, cache_info, SyncState.STARTED)
         try:
-            with as_temp_files([self._config, catalog.json()]) as [
+            with as_temp_files(
+                [self._config, catalog.json(), json.dumps(state) if state else "[]"]
+            ) as [
                 config_file,
                 catalog_file,
+                state_file,
             ]:
                 yield from self._execute(
-                    # TODO: Add support for incremental syncs by sending `--state <state_file>`
-                    ["read", "--config", config_file, "--catalog", catalog_file],
+                    [
+                        "read",
+                        "--config",
+                        config_file,
+                        "--catalog",
+                        catalog_file,
+                        "--state",
+                        state_file,
+                    ],
                 )
         except Exception:
             send_telemetry(
@@ -520,11 +528,12 @@ def read(
             incoming_source_catalog=self.configured_catalog,
             stream_names=set(self.get_selected_streams()),
         )
+        state = cache.get_state() if not force_full_refresh else None
         cache.process_airbyte_messages(
             self._tally_records(
                 self._read(
                     cache.get_telemetry_info(),
-                    force_full_refresh=force_full_refresh,
+                    state=state,
                 ),
             ),
             write_strategy=write_strategy,