Fix: Add "Destinations 1.0" handling to PyAirbyte: sync IDs, generation IDs, and stream success statuses

airbyte/_future_cdk/catalog_providers.py

@@ -10,6 +10,8 @@
 
 from typing import TYPE_CHECKING, Any, final
 
+import ulid
+
 from airbyte_protocol.models import (
  ConfiguredAirbyteCatalog,
 )
@@ -44,7 +46,24 @@ def __init__(
  Since the catalog is passed by reference, the catalog manager may be updated with new
  streams as they are discovered.
  """
- self._catalog: ConfiguredAirbyteCatalog = configured_catalog
+ self._catalog: ConfiguredAirbyteCatalog = self.validate_catalog(configured_catalog)
+
+ @staticmethod
+ def validate_catalog(catalog: ConfiguredAirbyteCatalog) -> None:
+ """Validate the catalog to ensure it is valid.
+
+ This requires ensuring that `generationId` and `minGenerationId` are both set. If
+ not, both values will be set to `1`.
+ """
+ for stream in catalog.streams:
+ if stream.generation_id is None:
+ stream.generation_id = 1
+ if stream.minimum_generation_id is None:
+ stream.minimum_generation_id = 1
+ if stream.sync_id is None:
+ stream.sync_id = 1 # This should ideally increment monotonically with each sync.
+
+ return catalog
 
  @property
  def configured_catalog(self) -> ConfiguredAirbyteCatalog:

airbyte/_message_iterators.py

@@ -18,6 +18,7 @@
 )
 
 from airbyte.constants import AB_EXTRACTED_AT_COLUMN
+from airbyte.progress import _new_stream_success_message
 
 
 if TYPE_CHECKING:
@@ -90,6 +91,8 @@ def generator() -> Generator[AirbyteMessage, None, None]:
  state=state_provider.get_stream_state(stream_name),
  )
 
+ yield _new_stream_success_message(stream_name)
+
  return cls(generator())
 
  @classmethod

airbyte/progress.py

@@ -30,7 +30,15 @@
 from rich.markdown import Markdown as RichMarkdown
 from typing_extensions import Literal
 
-from airbyte_protocol.models import AirbyteStreamStatus, Type
+from airbyte_protocol.models import (
+ AirbyteMessage,
+ AirbyteStreamStatus,
+ AirbyteStreamStatusTraceMessage,
+ AirbyteTraceMessage,
+ StreamDescriptor,
+ TraceType,
+ Type,
+)
 
 from airbyte._util import meta
 from airbyte._util.telemetry import EventState, EventType, send_telemetry
@@ -40,8 +48,6 @@
  from collections.abc import Generator, Iterable
  from types import ModuleType
 
- from airbyte_protocol.models import AirbyteMessage
-
  from airbyte._message_iterators import AirbyteMessageIterator
  from airbyte.caches.base import CacheBase
  from airbyte.destinations.base import Destination
@@ -68,6 +74,24 @@
  IS_NOTEBOOK = False
 
 
+def _new_stream_success_message(stream_name: str) -> AirbyteMessage:
+ """Return a new stream success message."""
+ return AirbyteMessage(
+ type=Type.TRACE,
+ trace=AirbyteTraceMessage(
+ type=TraceType.STREAM_STATUS,
+ stream=stream_name,
+ emitted_at=pendulum.now().float_timestamp,
+ stream_status=AirbyteStreamStatusTraceMessage(
+ stream_descriptor=StreamDescriptor(
+ name=stream_name,
+ ),
+ status=AirbyteStreamStatus.COMPLETE,
+ ),
+ ),
+ )
+
+
 class ProgressStyle(Enum):
  """An enum of progress bar styles."""
 
@@ -201,6 +225,8 @@ def __init__(
  def tally_records_read(
  self,
  messages: Iterable[AirbyteMessage],
+ *,
+ auto_close_streams: bool = False,
  ) -> Generator[AirbyteMessage, Any, None]:
  """This method simply tallies the number of records processed and yields the messages."""
  # Update the display before we start.
@@ -247,6 +273,11 @@ def tally_records_read(
  # Update the display.
  self._update_display()
 
+ if auto_close_streams:
+ for stream_name in self._unclosed_stream_names:
+ yield _new_stream_success_message(stream_name)
+ self._log_stream_read_end(stream_name)
+
  def tally_pending_writes(
  self,
  messages: IO[str] | AirbyteMessageIterator,
@@ -342,6 +373,15 @@ def _log_stream_read_end(self, stream_name: str) -> None:
  )
  self.stream_read_end_times[stream_name] = time.time()
 
+ @property
+ def _unclosed_stream_names(self) -> list[str]:
+ """Return a list of streams that have not yet been fully read."""
+ return [
+ stream_name
+ for stream_name in self.stream_read_counts
+ if stream_name not in self.stream_read_end_times
+ ]
+
  def log_success(
  self,
  ) -> None:

examples/run_bigquery_destination.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""
+Usage:
+ poetry install
+ poetry run python examples/run_bigquery_destination.py
+"""
+
+from __future__ import annotations
+
+import tempfile
+import warnings
+from typing import cast
+
+import airbyte as ab
+from airbyte.secrets.base import SecretString
+from airbyte.secrets.google_gsm import GoogleGSMSecretManager
+
+warnings.filterwarnings("ignore", message="Cannot create BigQuery Storage client")
+
+
+AIRBYTE_INTERNAL_GCP_PROJECT = "dataline-integration-testing"
+SECRET_NAME = "SECRET_DESTINATION-BIGQUERY_CREDENTIALS__CREDS"
+
+bigquery_destination_secret: dict = cast(
+ SecretString,
+ GoogleGSMSecretManager(
+ project=AIRBYTE_INTERNAL_GCP_PROJECT,
+ credentials_json=ab.get_secret("GCP_GSM_CREDENTIALS"),
+ ).get_secret(SECRET_NAME),
+).parse_json()
+
+
+def main() -> None:
+ source = ab.get_source(
+ "source-faker",
+ config={"count": 1000, "seed": 0, "parallelism": 1, "always_updated": False},
+ install_if_missing=True,
+ )
+ source.check()
+ source.select_all_streams()
+
+ with tempfile.NamedTemporaryFile(mode="w+", delete=False, encoding="utf-8") as temp:
+ # Write credentials to the temp file
+ temp.write(bigquery_destination_secret["credentials_json"])
+ temp.flush()
+ temp.close()
+
+ destination = ab.get_destination(
+ "destination-bigquery",
+ config={**bigquery_destination_secret, "dataset_id": "pyairbyte_tests"},
+ )
+ write_result = destination.write(
+ source,
+ # cache=False, # Toggle comment to test with/without caching
+ )
+
+
+if __name__ == "__main__":
+ main()