airbyte.sources.base

   1# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
   2"""Base class implementation for sources."""
   3
   4from __future__ import annotations
   5
   6import sys
   7import threading
   8import time
   9import warnings
  10from itertools import islice
  11from typing import TYPE_CHECKING, Any, Literal
  12
  13import yaml
  14from rich import print  # noqa: A004  # Allow shadowing the built-in
  15from rich.console import Console
  16from rich.markdown import Markdown
  17from rich.markup import escape
  18from rich.table import Table
  19from typing_extensions import override
  20
  21from airbyte_protocol.models import (
  22    AirbyteCatalog,
  23    AirbyteMessage,
  24    ConfiguredAirbyteCatalog,
  25    ConfiguredAirbyteStream,
  26    DestinationSyncMode,
  27    SyncMode,
  28    Type,
  29)
  30
  31from airbyte import exceptions as exc
  32from airbyte._connector_base import ConnectorBase
  33from airbyte._message_iterators import AirbyteMessageIterator
  34from airbyte._util.temp_files import as_temp_files
  35from airbyte.caches.util import get_default_cache
  36from airbyte.datasets._lazy import LazyDataset
  37from airbyte.progress import ProgressStyle, ProgressTracker
  38from airbyte.records import StreamRecord, StreamRecordHandler
  39from airbyte.results import ReadResult
  40from airbyte.shared.catalog_providers import CatalogProvider
  41from airbyte.strategies import WriteStrategy
  42
  43
  44if TYPE_CHECKING:
  45    from collections.abc import Generator, Iterable, Iterator
  46
  47    from airbyte_protocol.models import (
  48        AirbyteStream,
  49        ConnectorSpecification,
  50    )
  51
  52    from airbyte._executors.base import Executor
  53    from airbyte.caches import CacheBase
  54    from airbyte.callbacks import ConfigChangeCallback
  55    from airbyte.datasets._inmemory import InMemoryDataset
  56    from airbyte.documents import Document
  57    from airbyte.shared.state_providers import StateProviderBase
  58    from airbyte.shared.state_writers import StateWriterBase
  59
  60from airbyte.constants import (
  61    AB_EXTRACTED_AT_COLUMN,
  62    AB_META_COLUMN,
  63    AB_RAW_ID_COLUMN,
  64)
  65
  66
  67class Source(ConnectorBase):  # noqa: PLR0904
  68    """A class representing a source that can be called."""
  69
  70    connector_type = "source"
  71
  72    def __init__(
  73        self,
  74        executor: Executor,
  75        name: str,
  76        config: dict[str, Any] | None = None,
  77        *,
  78        config_change_callback: ConfigChangeCallback | None = None,
  79        streams: str | list[str] | None = None,
  80        validate: bool = False,
  81        cursor_key_overrides: dict[str, str] | None = None,
  82        primary_key_overrides: dict[str, str | list[str]] | None = None,
  83    ) -> None:
  84        """Initialize the source.
  85
  86        If config is provided, it will be validated against the spec if validate is True.
  87        """
  88        self._to_be_selected_streams: list[str] | str = []
  89        """Used to hold selection criteria before catalog is known."""
  90
  91        super().__init__(
  92            executor=executor,
  93            name=name,
  94            config=config,
  95            config_change_callback=config_change_callback,
  96            validate=validate,
  97        )
  98        self._config_dict: dict[str, Any] | None = None
  99        self._last_log_messages: list[str] = []
 100        self._discovered_catalog: AirbyteCatalog | None = None
 101        self._selected_stream_names: list[str] = []
 102
 103        self._cursor_key_overrides: dict[str, str] = {}
 104        """A mapping of lower-cased stream names to cursor key overrides."""
 105
 106        self._primary_key_overrides: dict[str, list[str]] = {}
 107        """A mapping of lower-cased stream names to primary key overrides."""
 108
 109        if config is not None:
 110            self.set_config(config, validate=validate)
 111        if streams is not None:
 112            self.select_streams(streams)
 113        if cursor_key_overrides is not None:
 114            self.set_cursor_keys(**cursor_key_overrides)
 115        if primary_key_overrides is not None:
 116            self.set_primary_keys(**primary_key_overrides)
 117
 118    def set_streams(self, streams: list[str]) -> None:
 119        """Deprecated. See select_streams()."""
 120        warnings.warn(
 121            "The 'set_streams' method is deprecated and will be removed in a future version. "
 122            "Please use the 'select_streams' method instead.",
 123            DeprecationWarning,
 124            stacklevel=2,
 125        )
 126        self.select_streams(streams)
 127
 128    def set_cursor_key(
 129        self,
 130        stream_name: str,
 131        cursor_key: str,
 132    ) -> None:
 133        """Set the cursor for a single stream.
 134
 135        Note:
 136        - This does not unset previously set cursors.
 137        - The cursor key must be a single field name.
 138        - Not all streams support custom cursors. If a stream does not support custom cursors,
 139          the override may be ignored.
 140        - Stream names are case insensitive, while field names are case sensitive.
 141        - Stream names are not validated by PyAirbyte. If the stream name
 142          does not exist in the catalog, the override may be ignored.
 143        """
 144        self._cursor_key_overrides[stream_name.lower()] = cursor_key
 145
 146    def set_cursor_keys(
 147        self,
 148        **kwargs: str,
 149    ) -> None:
 150        """Override the cursor key for one or more streams.
 151
 152        Usage:
 153            ```python
 154            source.set_cursor_keys(
 155                stream1="cursor1",
 156                stream2="cursor2",
 157            )
 158            ```
 159
 160        Note:
 161        - This does not unset previously set cursors.
 162        - The cursor key must be a single field name.
 163        - Not all streams support custom cursors. If a stream does not support custom cursors,
 164          the override may be ignored.
 165        - Stream names are case insensitive, while field names are case sensitive.
 166        - Stream names are not validated by PyAirbyte. If the stream name
 167          does not exist in the catalog, the override may be ignored.
 168        """
 169        self._cursor_key_overrides.update({k.lower(): v for k, v in kwargs.items()})
 170
 171    def set_primary_key(
 172        self,
 173        stream_name: str,
 174        primary_key: str | list[str],
 175    ) -> None:
 176        """Set the primary key for a single stream.
 177
 178        Note:
 179        - This does not unset previously set primary keys.
 180        - The primary key must be a single field name or a list of field names.
 181        - Not all streams support overriding primary keys. If a stream does not support overriding
 182          primary keys, the override may be ignored.
 183        - Stream names are case insensitive, while field names are case sensitive.
 184        - Stream names are not validated by PyAirbyte. If the stream name
 185          does not exist in the catalog, the override may be ignored.
 186        """
 187        self._primary_key_overrides[stream_name.lower()] = (
 188            primary_key if isinstance(primary_key, list) else [primary_key]
 189        )
 190
 191    def set_primary_keys(
 192        self,
 193        **kwargs: str | list[str],
 194    ) -> None:
 195        """Override the primary keys for one or more streams.
 196
 197        This does not unset previously set primary keys.
 198
 199        Usage:
 200            ```python
 201            source.set_primary_keys(
 202                stream1="pk1",
 203                stream2=["pk1", "pk2"],
 204            )
 205            ```
 206
 207        Note:
 208        - This does not unset previously set primary keys.
 209        - The primary key must be a single field name or a list of field names.
 210        - Not all streams support overriding primary keys. If a stream does not support overriding
 211          primary keys, the override may be ignored.
 212        - Stream names are case insensitive, while field names are case sensitive.
 213        - Stream names are not validated by PyAirbyte. If the stream name
 214          does not exist in the catalog, the override may be ignored.
 215        """
 216        self._primary_key_overrides.update(
 217            {k.lower(): v if isinstance(v, list) else [v] for k, v in kwargs.items()}
 218        )
 219
 220    def _log_warning_preselected_stream(self, streams: str | list[str]) -> None:
 221        """Logs a warning message indicating stream selection which are not selected yet."""
 222        if streams == "*":
 223            print(
 224                "Warning: Config is not set yet. All streams will be selected after config is set.",
 225                file=sys.stderr,
 226            )
 227        else:
 228            print(
 229                "Warning: Config is not set yet. "
 230                f"Streams to be selected after config is set: {streams}",
 231                file=sys.stderr,
 232            )
 233
 234    def select_all_streams(self) -> None:
 235        """Select all streams.
 236
 237        This is a more streamlined equivalent to:
 238        > source.select_streams(source.get_available_streams()).
 239        """
 240        if self._config_dict is None:
 241            self._to_be_selected_streams = "*"
 242            self._log_warning_preselected_stream(self._to_be_selected_streams)
 243            return
 244
 245        self._selected_stream_names = self.get_available_streams()
 246
 247    def select_streams(self, streams: str | list[str]) -> None:
 248        """Select the stream names that should be read from the connector.
 249
 250        Args:
 251            streams: A list of stream names to select. If set to "*", all streams will be selected.
 252
 253        Currently, if this is not set, all streams will be read.
 254        """
 255        if self._config_dict is None:
 256            self._to_be_selected_streams = streams
 257            self._log_warning_preselected_stream(streams)
 258            return
 259
 260        if streams == "*":
 261            self.select_all_streams()
 262            return
 263
 264        if isinstance(streams, str):
 265            # If a single stream is provided, convert it to a one-item list
 266            streams = [streams]
 267
 268        available_streams = self.get_available_streams()
 269        for stream in streams:
 270            if stream not in available_streams:
 271                raise exc.AirbyteStreamNotFoundError(
 272                    stream_name=stream,
 273                    connector_name=self.name,
 274                    available_streams=available_streams,
 275                )
 276        self._selected_stream_names = streams
 277
 278    def get_selected_streams(self) -> list[str]:
 279        """Get the selected streams.
 280
 281        If no streams are selected, return an empty list.
 282        """
 283        return self._selected_stream_names
 284
 285    def set_config(
 286        self,
 287        config: dict[str, Any],
 288        *,
 289        validate: bool = True,
 290    ) -> None:
 291        """Set the config for the connector.
 292
 293        If validate is True, raise an exception if the config fails validation.
 294
 295        If validate is False, validation will be deferred until check() or validate_config()
 296        is called.
 297        """
 298        if validate:
 299            self.validate_config(config)
 300
 301        self._config_dict = config
 302
 303        if self._to_be_selected_streams:
 304            self.select_streams(self._to_be_selected_streams)
 305            self._to_be_selected_streams = []
 306
 307    def _discover(self) -> AirbyteCatalog:
 308        """Call discover on the connector.
 309
 310        This involves the following steps:
 311        - Write the config to a temporary file
 312        - execute the connector with discover --config <config_file>
 313        - Listen to the messages and return the first AirbyteCatalog that comes along.
 314        - Make sure the subprocess is killed when the function returns.
 315        """
 316        with as_temp_files([self._hydrated_config]) as [config_file]:
 317            for msg in self._execute(["discover", "--config", config_file]):
 318                if msg.type == Type.CATALOG and msg.catalog:
 319                    return msg.catalog
 320            raise exc.AirbyteConnectorMissingCatalogError(
 321                connector_name=self.name,
 322                log_text=self._last_log_messages,
 323            )
 324
 325    def get_available_streams(self) -> list[str]:
 326        """Get the available streams from the spec."""
 327        return [s.name for s in self.discovered_catalog.streams]
 328
 329    def _get_incremental_stream_names(self) -> list[str]:
 330        """Get the name of streams that support incremental sync."""
 331        return [
 332            stream.name
 333            for stream in self.discovered_catalog.streams
 334            if SyncMode.incremental in stream.supported_sync_modes
 335        ]
 336
 337    @override
 338    def _get_spec(self, *, force_refresh: bool = False) -> ConnectorSpecification:
 339        """Call spec on the connector.
 340
 341        This involves the following steps:
 342        * execute the connector with spec
 343        * Listen to the messages and return the first AirbyteCatalog that comes along.
 344        * Make sure the subprocess is killed when the function returns.
 345        """
 346        if force_refresh or self._spec is None:
 347            for msg in self._execute(["spec"]):
 348                if msg.type == Type.SPEC and msg.spec:
 349                    self._spec = msg.spec
 350                    break
 351
 352        if self._spec:
 353            return self._spec
 354
 355        raise exc.AirbyteConnectorMissingSpecError(
 356            connector_name=self.name,
 357            log_text=self._last_log_messages,
 358        )
 359
 360    @property
 361    def config_spec(self) -> dict[str, Any]:
 362        """Generate a configuration spec for this connector, as a JSON Schema definition.
 363
 364        This function generates a JSON Schema dictionary with configuration specs for the
 365        current connector, as a dictionary.
 366
 367        Returns:
 368            dict: The JSON Schema configuration spec as a dictionary.
 369        """
 370        return self._get_spec(force_refresh=True).connectionSpecification
 371
 372    @property
 373    def _yaml_spec(self) -> str:
 374        """Get the spec as a yaml string.
 375
 376        For now, the primary use case is for writing and debugging a valid config for a source.
 377
 378        This is private for now because we probably want better polish before exposing this
 379        as a stable interface. This will also get easier when we have docs links with this info
 380        for each connector.
 381        """
 382        spec_obj: ConnectorSpecification = self._get_spec()
 383        spec_dict: dict[str, Any] = spec_obj.model_dump(exclude_unset=True)
 384        # convert to a yaml string
 385        return yaml.dump(spec_dict)
 386
 387    @property
 388    def docs_url(self) -> str:
 389        """Get the URL to the connector's documentation."""
 390        return "https://docs.airbyte.com/integrations/sources/" + self.name.lower().replace(
 391            "source-", ""
 392        )
 393
 394    @property
 395    def discovered_catalog(self) -> AirbyteCatalog:
 396        """Get the raw catalog for the given streams.
 397
 398        If the catalog is not yet known, we call discover to get it.
 399        """
 400        if self._discovered_catalog is None:
 401            self._discovered_catalog = self._discover()
 402
 403        return self._discovered_catalog
 404
 405    @property
 406    def configured_catalog(self) -> ConfiguredAirbyteCatalog:
 407        """Get the configured catalog for the given streams.
 408
 409        If the raw catalog is not yet known, we call discover to get it.
 410
 411        If no specific streams are selected, we return a catalog that syncs all available streams.
 412
 413        TODO: We should consider disabling by default the streams that the connector would
 414        disable by default. (For instance, streams that require a premium license are sometimes
 415        disabled by default within the connector.)
 416        """
 417        # Ensure discovered catalog is cached before we start
 418        _ = self.discovered_catalog
 419
 420        # Filter for selected streams if set, otherwise use all available streams:
 421        streams_filter: list[str] = self._selected_stream_names or self.get_available_streams()
 422        return self.get_configured_catalog(streams=streams_filter)
 423
 424    def get_configured_catalog(
 425        self,
 426        streams: Literal["*"] | list[str] | None = None,
 427        *,
 428        force_full_refresh: bool = False,
 429    ) -> ConfiguredAirbyteCatalog:
 430        """Get a configured catalog for the given streams.
 431
 432        If no streams are provided, the selected streams will be used. If no streams are selected,
 433        all available streams will be used.
 434
 435        If '*' is provided, all available streams will be used.
 436
 437        If force_full_refresh is True, streams will be configured with full_refresh sync mode
 438        when supported by the stream. Otherwise, incremental sync mode is used when supported.
 439        """
 440        selected_streams: list[str] = []
 441        if streams is None:
 442            selected_streams = self._selected_stream_names or self.get_available_streams()
 443        elif streams == "*":
 444            selected_streams = self.get_available_streams()
 445        elif isinstance(streams, list):
 446            selected_streams = streams
 447        else:
 448            raise exc.PyAirbyteInputError(
 449                message="Invalid streams argument.",
 450                input_value=streams,
 451            )
 452
 453        def _get_sync_mode(stream: AirbyteStream) -> SyncMode:
 454            """Determine the sync mode for a stream based on force_full_refresh and support."""
 455            # Use getattr to handle mocks or streams without supported_sync_modes attribute
 456            supported_modes = getattr(stream, "supported_sync_modes", None)
 457
 458            if force_full_refresh:
 459                # When force_full_refresh is True, prefer full_refresh if supported
 460                if supported_modes and SyncMode.full_refresh in supported_modes:
 461                    return SyncMode.full_refresh
 462                # Fall back to incremental if full_refresh is not supported
 463                return SyncMode.incremental
 464
 465            # Default behavior: preserve previous semantics (always incremental)
 466            return SyncMode.incremental
 467
 468        return ConfiguredAirbyteCatalog(
 469            streams=[
 470                ConfiguredAirbyteStream(
 471                    stream=stream,
 472                    destination_sync_mode=DestinationSyncMode.overwrite,
 473                    sync_mode=_get_sync_mode(stream),
 474                    primary_key=(
 475                        [self._primary_key_overrides[stream.name.lower()]]
 476                        if stream.name.lower() in self._primary_key_overrides
 477                        else stream.source_defined_primary_key
 478                    ),
 479                    cursor_field=(
 480                        [self._cursor_key_overrides[stream.name.lower()]]
 481                        if stream.name.lower() in self._cursor_key_overrides
 482                        else stream.default_cursor_field
 483                    ),
 484                    # These are unused in the current implementation:
 485                    generation_id=None,
 486                    minimum_generation_id=None,
 487                    sync_id=None,
 488                )
 489                for stream in self.discovered_catalog.streams
 490                if stream.name in selected_streams
 491            ],
 492        )
 493
 494    def get_stream_json_schema(self, stream_name: str) -> dict[str, Any]:
 495        """Return the JSON Schema spec for the specified stream name."""
 496        catalog: AirbyteCatalog = self.discovered_catalog
 497        found: list[AirbyteStream] = [
 498            stream for stream in catalog.streams if stream.name == stream_name
 499        ]
 500
 501        if len(found) == 0:
 502            raise exc.PyAirbyteInputError(
 503                message="Stream name does not exist in catalog.",
 504                input_value=stream_name,
 505            )
 506
 507        if len(found) > 1:
 508            raise exc.PyAirbyteInternalError(
 509                message="Duplicate streams found with the same name.",
 510                context={
 511                    "found_streams": found,
 512                },
 513            )
 514
 515        return found[0].json_schema
 516
 517    def get_records(
 518        self,
 519        stream: str,
 520        *,
 521        limit: int | None = None,
 522        stop_event: threading.Event | None = None,
 523        normalize_field_names: bool = False,
 524        prune_undeclared_fields: bool = True,
 525    ) -> LazyDataset:
 526        """Read a stream from the connector.
 527
 528        Args:
 529            stream: The name of the stream to read.
 530            limit: The maximum number of records to read. If None, all records will be read.
 531            stop_event: If set, the event can be triggered by the caller to stop reading records
 532                and terminate the process.
 533            normalize_field_names: When `True`, field names will be normalized to lower case, with
 534                special characters removed. This matches the behavior of PyAirbyte caches and most
 535                Airbyte destinations.
 536            prune_undeclared_fields: When `True`, undeclared fields will be pruned from the records,
 537                which generally matches the behavior of PyAirbyte caches and most Airbyte
 538                destinations, specifically when you expect the catalog may be stale. You can disable
 539                this to keep all fields in the records.
 540
 541        This involves the following steps:
 542        * Call discover to get the catalog
 543        * Generate a configured catalog that syncs the given stream in full_refresh mode
 544        * Write the configured catalog and the config to a temporary file
 545        * execute the connector with read --config <config_file> --catalog <catalog_file>
 546        * Listen to the messages and return the first AirbyteRecordMessages that come along.
 547        * Make sure the subprocess is killed when the function returns.
 548        """
 549        stop_event = stop_event or threading.Event()
 550        configured_catalog = self.get_configured_catalog(streams=[stream])
 551        if len(configured_catalog.streams) == 0:
 552            raise exc.PyAirbyteInputError(
 553                message="Requested stream does not exist.",
 554                context={
 555                    "stream": stream,
 556                    "available_streams": self.get_available_streams(),
 557                    "connector_name": self.name,
 558                },
 559            ) from KeyError(stream)
 560
 561        configured_stream = configured_catalog.streams[0]
 562
 563        def _with_logging(records: Iterable[dict[str, Any]]) -> Iterator[dict[str, Any]]:
 564            yield from records
 565
 566        stream_record_handler = StreamRecordHandler(
 567            json_schema=self.get_stream_json_schema(stream),
 568            prune_extra_fields=prune_undeclared_fields,
 569            normalize_keys=normalize_field_names,
 570        )
 571
 572        # This method is non-blocking, so we use "PLAIN" to avoid a live progress display
 573        progress_tracker = ProgressTracker(
 574            ProgressStyle.PLAIN,
 575            source=self,
 576            cache=None,
 577            destination=None,
 578            expected_streams=[stream],
 579        )
 580
 581        iterator: Iterator[dict[str, Any]] = (
 582            StreamRecord.from_record_message(
 583                record_message=record.record,
 584                stream_record_handler=stream_record_handler,
 585            )
 586            for record in self._read_with_catalog(
 587                catalog=configured_catalog,
 588                progress_tracker=progress_tracker,
 589                stop_event=stop_event,
 590            )
 591            if record.record
 592        )
 593        if limit is not None:
 594            # Stop the iterator after the limit is reached
 595            iterator = islice(iterator, limit)
 596
 597        return LazyDataset(
 598            iterator,
 599            stream_metadata=configured_stream,
 600            stop_event=stop_event,
 601            progress_tracker=progress_tracker,
 602        )
 603
 604    def get_documents(
 605        self,
 606        stream: str,
 607        title_property: str | None = None,
 608        content_properties: list[str] | None = None,
 609        metadata_properties: list[str] | None = None,
 610        *,
 611        render_metadata: bool = False,
 612    ) -> Iterable[Document]:
 613        """Read a stream from the connector and return the records as documents.
 614
 615        If metadata_properties is not set, all properties that are not content will be added to
 616        the metadata.
 617
 618        If render_metadata is True, metadata will be rendered in the document, as well as the
 619        the main content.
 620        """
 621        return self.get_records(stream).to_documents(
 622            title_property=title_property,
 623            content_properties=content_properties,
 624            metadata_properties=metadata_properties,
 625            render_metadata=render_metadata,
 626        )
 627
 628    def get_samples(
 629        self,
 630        streams: list[str] | Literal["*"] | None = None,
 631        *,
 632        limit: int = 5,
 633        on_error: Literal["raise", "ignore", "log"] = "raise",
 634    ) -> dict[str, InMemoryDataset | None]:
 635        """Get a sample of records from the given streams."""
 636        if streams == "*":
 637            streams = self.get_available_streams()
 638        elif streams is None:
 639            streams = self.get_selected_streams()
 640
 641        results: dict[str, InMemoryDataset | None] = {}
 642        for stream in streams:
 643            stop_event = threading.Event()
 644            try:
 645                results[stream] = self.get_records(
 646                    stream,
 647                    limit=limit,
 648                    stop_event=stop_event,
 649                ).fetch_all()
 650                stop_event.set()
 651            except Exception as ex:
 652                results[stream] = None
 653                if on_error == "ignore":
 654                    continue
 655
 656                if on_error == "raise":
 657                    raise ex from None
 658
 659                if on_error == "log":
 660                    print(f"Error fetching sample for stream '{stream}': {ex}")
 661
 662        return results
 663
 664    def print_samples(
 665        self,
 666        streams: list[str] | Literal["*"] | None = None,
 667        *,
 668        limit: int = 5,
 669        on_error: Literal["raise", "ignore", "log"] = "log",
 670    ) -> None:
 671        """Print a sample of records from the given streams."""
 672        internal_cols: list[str] = [
 673            AB_EXTRACTED_AT_COLUMN,
 674            AB_META_COLUMN,
 675            AB_RAW_ID_COLUMN,
 676        ]
 677        col_limit = 10
 678        if streams == "*":
 679            streams = self.get_available_streams()
 680        elif streams is None:
 681            streams = self.get_selected_streams()
 682
 683        console = Console()
 684
 685        console.print(
 686            Markdown(
 687                f"# Sample Records from `{self.name}` ({len(streams)} selected streams)",
 688                justify="left",
 689            )
 690        )
 691
 692        for stream in streams:
 693            console.print(Markdown(f"## `{stream}` Stream Sample", justify="left"))
 694            samples = self.get_samples(
 695                streams=[stream],
 696                limit=limit,
 697                on_error=on_error,
 698            )
 699            dataset = samples[stream]
 700
 701            table = Table(
 702                show_header=True,
 703                show_lines=True,
 704            )
 705            if dataset is None:
 706                console.print(
 707                    Markdown("**⚠️ `Error fetching sample records.` ⚠️**"),
 708                )
 709                continue
 710
 711            if len(dataset.column_names) > col_limit:
 712                # We'll pivot the columns so each column is its own row
 713                table.add_column("Column Name")
 714                for _ in range(len(dataset)):
 715                    table.add_column(overflow="fold")
 716                for col in dataset.column_names:
 717                    table.add_row(
 718                        Markdown(f"**`{col}`**"),
 719                        *[escape(str(record[col])) for record in dataset],
 720                    )
 721            else:
 722                for col in dataset.column_names:
 723                    table.add_column(
 724                        Markdown(f"**`{col}`**"),
 725                        overflow="fold",
 726                    )
 727
 728                for record in dataset:
 729                    table.add_row(
 730                        *[
 731                            escape(str(val))
 732                            for key, val in record.items()
 733                            # Exclude internal Airbyte columns.
 734                            if key not in internal_cols
 735                        ]
 736                    )
 737
 738            console.print(table)
 739
 740        console.print(Markdown("--------------"))
 741
 742    def _get_airbyte_message_iterator(
 743        self,
 744        *,
 745        streams: Literal["*"] | list[str] | None = None,
 746        state_provider: StateProviderBase | None = None,
 747        progress_tracker: ProgressTracker,
 748        force_full_refresh: bool = False,
 749    ) -> AirbyteMessageIterator:
 750        """Get an AirbyteMessageIterator for this source."""
 751        return AirbyteMessageIterator(
 752            self._read_with_catalog(
 753                catalog=self.get_configured_catalog(
 754                    streams=streams,
 755                    force_full_refresh=force_full_refresh,
 756                ),
 757                state=state_provider if not force_full_refresh else None,
 758                progress_tracker=progress_tracker,
 759            )
 760        )
 761
 762    def _read_with_catalog(
 763        self,
 764        catalog: ConfiguredAirbyteCatalog,
 765        progress_tracker: ProgressTracker,
 766        *,
 767        state: StateProviderBase | None = None,
 768        stop_event: threading.Event | None = None,
 769    ) -> Generator[AirbyteMessage, None, None]:
 770        """Call read on the connector.
 771
 772        This involves the following steps:
 773        * Write the config to a temporary file
 774        * execute the connector with read --config <config_file> --catalog <catalog_file>
 775        * Listen to the messages and return the AirbyteRecordMessages that come along.
 776        * Send out telemetry on the performed sync (with information about which source was used and
 777          the type of the cache)
 778        """
 779        with as_temp_files(
 780            [
 781                self._hydrated_config,
 782                catalog.model_dump_json(exclude_none=True),
 783                state.to_state_input_file_text() if state else "[]",
 784            ]
 785        ) as [
 786            config_file,
 787            catalog_file,
 788            state_file,
 789        ]:
 790            message_generator = self._execute(
 791                [
 792                    "read",
 793                    "--config",
 794                    config_file,
 795                    "--catalog",
 796                    catalog_file,
 797                    "--state",
 798                    state_file,
 799                ],
 800                progress_tracker=progress_tracker,
 801            )
 802            for message in progress_tracker.tally_records_read(message_generator):
 803                if stop_event and stop_event.is_set():
 804                    progress_tracker._log_sync_cancel()  # noqa: SLF001
 805                    time.sleep(0.1)
 806                    return
 807
 808                yield message
 809
 810        progress_tracker.log_read_complete()
 811
 812    def _peek_airbyte_message(
 813        self,
 814        message: AirbyteMessage,
 815        *,
 816        raise_on_error: bool = True,
 817    ) -> None:
 818        """Process an Airbyte message.
 819
 820        This method handles reading Airbyte messages and taking action, if needed, based on the
 821        message type. For instance, log messages are logged, records are tallied, and errors are
 822        raised as exceptions if `raise_on_error` is True.
 823
 824        Raises:
 825            AirbyteConnectorFailedError: If a TRACE message of type ERROR is emitted.
 826        """
 827        super()._peek_airbyte_message(message, raise_on_error=raise_on_error)
 828
 829    def _log_incremental_streams(
 830        self,
 831        *,
 832        incremental_streams: set[str] | None = None,
 833    ) -> None:
 834        """Log the streams which are using incremental sync mode."""
 835        log_message = (
 836            "The following streams are currently using incremental sync:\n"
 837            f"{incremental_streams}\n"
 838            "To perform a full refresh, set 'force_full_refresh=True' in 'airbyte.read()' method."
 839        )
 840        print(log_message, file=sys.stderr)
 841
 842    def read(
 843        self,
 844        cache: CacheBase | None = None,
 845        *,
 846        streams: str | list[str] | None = None,
 847        write_strategy: str | WriteStrategy = WriteStrategy.AUTO,
 848        force_full_refresh: bool = False,
 849        skip_validation: bool = False,
 850    ) -> ReadResult:
 851        """Read from the connector and write to the cache.
 852
 853        Args:
 854            cache: The cache to write to. If not set, a default cache will be used.
 855            streams: Optional if already set. A list of stream names to select for reading. If set
 856                to "*", all streams will be selected.
 857            write_strategy: The strategy to use when writing to the cache. If a string, it must be
 858                one of "append", "merge", "replace", or "auto". If a WriteStrategy, it must be one
 859                of WriteStrategy.APPEND, WriteStrategy.MERGE, WriteStrategy.REPLACE, or
 860                WriteStrategy.AUTO.
 861            force_full_refresh: If True, the source will operate in full refresh mode. Otherwise,
 862                streams will be read in incremental mode if supported by the connector. This option
 863                must be True when using the "replace" strategy.
 864            skip_validation: If True, PyAirbyte will not pre-validate the input configuration before
 865                running the connector. This can be helpful in debugging, when you want to send
 866                configurations to the connector that otherwise might be rejected by JSON Schema
 867                validation rules.
 868        """
 869        cache = cache or get_default_cache()
 870        progress_tracker = ProgressTracker(
 871            source=self,
 872            cache=cache,
 873            destination=None,
 874            expected_streams=None,  # Will be set later
 875        )
 876
 877        # Set up state provider if not in full refresh mode
 878        if force_full_refresh:
 879            state_provider: StateProviderBase | None = None
 880        else:
 881            state_provider = cache.get_state_provider(
 882                source_name=self._name,
 883            )
 884        state_writer = cache.get_state_writer(source_name=self._name)
 885
 886        if streams:
 887            self.select_streams(streams)
 888
 889        if not self._selected_stream_names:
 890            raise exc.PyAirbyteNoStreamsSelectedError(
 891                connector_name=self.name,
 892                available_streams=self.get_available_streams(),
 893            )
 894
 895        try:
 896            result = self._read_to_cache(
 897                cache=cache,
 898                catalog_provider=CatalogProvider(
 899                    self.get_configured_catalog(force_full_refresh=force_full_refresh)
 900                ),
 901                stream_names=self._selected_stream_names,
 902                state_provider=state_provider,
 903                state_writer=state_writer,
 904                write_strategy=write_strategy,
 905                force_full_refresh=force_full_refresh,
 906                skip_validation=skip_validation,
 907                progress_tracker=progress_tracker,
 908            )
 909        except exc.PyAirbyteInternalError as ex:
 910            progress_tracker.log_failure(exception=ex)
 911            raise exc.AirbyteConnectorFailedError(
 912                connector_name=self.name,
 913                log_text=self._last_log_messages,
 914            ) from ex
 915        except Exception as ex:
 916            progress_tracker.log_failure(exception=ex)
 917            raise
 918
 919        progress_tracker.log_success()
 920        return result
 921
 922    def _read_to_cache(  # noqa: PLR0913  # Too many arguments
 923        self,
 924        cache: CacheBase,
 925        *,
 926        catalog_provider: CatalogProvider,
 927        stream_names: list[str],
 928        state_provider: StateProviderBase | None,
 929        state_writer: StateWriterBase | None,
 930        write_strategy: str | WriteStrategy = WriteStrategy.AUTO,
 931        force_full_refresh: bool = False,
 932        skip_validation: bool = False,
 933        progress_tracker: ProgressTracker,
 934    ) -> ReadResult:
 935        """Internal read method."""
 936        if write_strategy == WriteStrategy.REPLACE and not force_full_refresh:
 937            warnings.warn(
 938                message=(
 939                    "Using `REPLACE` strategy without also setting `force_full_refresh=True` "
 940                    "could result in data loss. "
 941                    "To silence this warning, use the following: "
 942                    'warnings.filterwarnings("ignore", '
 943                    'category="airbyte.warnings.PyAirbyteDataLossWarning")`'
 944                ),
 945                category=exc.PyAirbyteDataLossWarning,
 946                stacklevel=1,
 947            )
 948        if isinstance(write_strategy, str):
 949            try:
 950                write_strategy = WriteStrategy(write_strategy)
 951            except ValueError:
 952                raise exc.PyAirbyteInputError(
 953                    message="Invalid strategy",
 954                    context={
 955                        "write_strategy": write_strategy,
 956                        "available_strategies": [
 957                            s.value
 958                            for s in WriteStrategy  # pyrefly: ignore[not-iterable]
 959                        ],
 960                    },
 961                ) from None
 962
 963        # Run optional validation step
 964        if not skip_validation:
 965            self.validate_config()
 966
 967        # Log incremental stream if incremental streams are known
 968        if state_provider and state_provider.known_stream_names:
 969            # Retrieve set of the known streams support which support incremental sync
 970            incremental_streams = (
 971                set(self._get_incremental_stream_names())
 972                & state_provider.known_stream_names
 973                & set(self.get_selected_streams())
 974            )
 975            if incremental_streams:
 976                self._log_incremental_streams(incremental_streams=incremental_streams)
 977
 978        airbyte_message_iterator = AirbyteMessageIterator(
 979            self._read_with_catalog(
 980                catalog=catalog_provider.configured_catalog,
 981                state=state_provider,
 982                progress_tracker=progress_tracker,
 983            )
 984        )
 985        cache._write_airbyte_message_stream(  # noqa: SLF001  # Non-public API
 986            stdin=airbyte_message_iterator,
 987            catalog_provider=catalog_provider,
 988            write_strategy=write_strategy,
 989            state_writer=state_writer,
 990            progress_tracker=progress_tracker,
 991        )
 992
 993        # Flush the WAL, if applicable
 994        cache.processor._do_checkpoint()  # noqa: SLF001  # Non-public API
 995
 996        return ReadResult(
 997            source_name=self.name,
 998            progress_tracker=progress_tracker,
 999            processed_streams=stream_names,
1000            cache=cache,
1001        )
1002
1003
1004__all__ = [
1005    "Source",
1006]