airbyte.sources

 48def get_source(  # noqa: PLR0913 # Too many arguments
 49    name: str,
 50    config: dict[str, Any] | None = None,
 51    *,
 52    config_change_callback: ConfigChangeCallback | None = None,
 53    streams: str | list[str] | None = None,
 54    version: str | None = None,
 55    use_python: bool | Path | str | None = None,
 56    pip_url: str | None = None,
 57    local_executable: Path | str | None = None,
 58    docker_image: bool | str | None = None,
 59    use_host_network: bool = False,
 60    source_manifest: bool | dict | Path | str | None = None,
 61    install_if_missing: bool = True,
 62    install_root: Path | None = None,
 63    no_executor: bool = False,
 64) -> Source:
 65    """Get a connector by name and version.
 66
 67    If an explicit install or execution method is requested (e.g. `local_executable`,
 68    `docker_image`, `pip_url`, `source_manifest`), the connector will be executed using this method.
 69
 70    Otherwise, an appropriate method will be selected based on the available connector metadata:
 71    1. If the connector is registered and has a YAML source manifest is available, the YAML manifest
 72       will be downloaded and used to to execute the connector.
 73    2. Else, if the connector is registered and has a PyPI package, it will be installed via pip.
 74    3. Else, if the connector is registered and has a Docker image, and if Docker is available, it
 75       will be executed using Docker.
 76
 77    Args:
 78        name: connector name
 79        config: connector config - if not provided, you need to set it later via the set_config
 80            method.
 81        config_change_callback: callback function to be called when the connector config changes.
 82        streams: list of stream names to select for reading. If set to "*", all streams will be
 83            selected. If not provided, you can set it later via the `select_streams()` or
 84            `select_all_streams()` method.
 85        version: connector version - if not provided, the currently installed version will be used.
 86            If no version is installed, the latest available version will be used. The version can
 87            also be set to "latest" to force the use of the latest available version.
 88        use_python: (Optional.) Python interpreter specification:
 89            - True: Use current Python interpreter. (Inferred if `pip_url` is set.)
 90            - False: Use Docker instead.
 91            - Path: Use interpreter at this path.
 92            - str: Use specific Python version. E.g. "3.11" or "3.11.10". If the version is not yet
 93                installed, it will be installed by uv. (This generally adds less than 3 seconds
 94                to install times.)
 95        pip_url: connector pip URL - if not provided, the pip url will be inferred from the
 96            connector name.
 97        local_executable: If set, the connector will be assumed to already be installed and will be
 98            executed using this path or executable name. Otherwise, the connector will be installed
 99            automatically in a virtual environment.
100        docker_image: If set, the connector will be executed using Docker. You can specify `True`
101            to use the default image for the connector, or you can specify a custom image name.
102            If `version` is specified and your image name does not already contain a tag
103            (e.g. `my-image:latest`), the version will be appended as a tag (e.g. `my-image:0.1.0`).
104        use_host_network: If set, along with docker_image, the connector will be executed using
105            the host network. This is useful for connectors that need to access resources on
106            the host machine, such as a local database. This parameter is ignored when
107            `docker_image` is not set.
108        source_manifest: If set, the connector will be executed based on a declarative YAML
109            source definition. This input can be `True` to attempt to auto-download a YAML spec,
110            `dict` to accept a Python dictionary as the manifest, `Path` to pull a manifest from
111            the local file system, or `str` to pull the definition from a web URL.
112        install_if_missing: Whether to install the connector if it is not available locally. This
113            parameter is ignored when `local_executable` or `source_manifest` are set.
114        install_root: (Optional.) The root directory where the virtual environment will be
115            created. If not provided, the current working directory will be used.
116        no_executor: If True, use NoOpExecutor which fetches specs from the registry without
117            local installation. This is useful for scenarios where you need to validate
118            configurations but don't need to run the connector locally (e.g., deploying to Cloud).
119    """
120    executor = get_connector_executor(
121        name=name,
122        version=version,
123        use_python=use_python,
124        pip_url=pip_url,
125        local_executable=local_executable,
126        docker_image=docker_image,
127        use_host_network=use_host_network,
128        source_manifest=source_manifest,
129        install_if_missing=install_if_missing,
130        install_root=install_root,
131        no_executor=no_executor,
132    )
133
134    return Source(
135        name=name,
136        config=config,
137        config_change_callback=config_change_callback,
138        streams=streams,
139        executor=executor,
140    )

143def get_benchmark_source(
144    num_records: int | str = "5e5",
145    *,
146    install_if_missing: bool = True,
147) -> Source:
148    """Get a source for benchmarking.
149
150    This source will generate dummy records for performance benchmarking purposes.
151    You can specify the number of records to generate using the `num_records` parameter.
152    The `num_records` parameter can be an integer or a string in scientific notation.
153    For example, `"5e6"` will generate 5 million records. If underscores are providing
154    within a numeric a string, they will be ignored.
155
156    Args:
157        num_records: The number of records to generate. Defaults to "5e5", or
158            500,000 records.
159            Can be an integer (`1000`) or a string in scientific notation.
160            For example, `"5e6"` will generate 5 million records.
161        install_if_missing: Whether to install the source if it is not available locally.
162
163    Returns:
164        Source: The source object for benchmarking.
165    """
166    if isinstance(num_records, str):
167        try:
168            num_records = int(Decimal(num_records.replace("_", "")))
169        except InvalidOperation as ex:
170            raise PyAirbyteInputError(
171                message="Invalid number format.",
172                original_exception=ex,
173                input_value=str(num_records),
174            ) from None
175
176    return get_source(
177        name="source-e2e-test",
178        docker_image=True,
179        # docker_image="airbyte/source-e2e-test:latest",
180        config={
181            "type": "BENCHMARK",
182            "schema": "FIVE_STRING_COLUMNS",
183            "terminationCondition": {
184                "type": "MAX_RECORDS",
185                "max": num_records,
186            },
187        },
188        streams="*",
189        install_if_missing=install_if_missing,
190    )

293def get_available_connectors(
294    install_type: InstallType | str | None = InstallType.INSTALLABLE,
295) -> list[str]:
296    """Return a list of all available connectors.
297
298    Connectors will be returned in alphabetical order, with the standard prefix "source-".
299
300    Args:
301        install_type: The type of installation for the connector.
302            Defaults to `InstallType.INSTALLABLE`.
303    """
304    if install_type is None or install_type == InstallType.INSTALLABLE:
305        # Filter for installable connectors (default behavior).
306        if is_docker_installed():
307            logger.info("Docker is detected. Returning all connectors.")
308            return sorted(_get_registry_cache().keys())
309
310        logger.info("Docker was not detected. Returning only Python and Manifest-only connectors.")
311        return sorted(
312            [
313                connector_name
314                for connector_name, conn_info in _get_registry_cache().items()
315                if conn_info.language in {Language.PYTHON, Language.MANIFEST_ONLY}
316            ]
317        )
318
319    if not isinstance(install_type, InstallType):
320        install_type = InstallType(install_type)
321
322    if install_type == InstallType.PYTHON:
323        return sorted(
324            connector_name
325            for connector_name, conn_info in _get_registry_cache().items()
326            if conn_info.pypi_package_name is not None
327        )
328
329    if install_type == InstallType.JAVA:
330        warnings.warn(
331            message="Java connectors are not yet supported.",
332            stacklevel=2,
333        )
334        return sorted(
335            connector_name
336            for connector_name, conn_info in _get_registry_cache().items()
337            if conn_info.language == Language.JAVA
338        )
339
340    if install_type in {InstallType.DOCKER, InstallType.ANY}:
341        return sorted(_get_registry_cache().keys())
342
343    if install_type == InstallType.YAML:
344        return sorted(
345            conn.name
346            for conn in _get_registry_cache().values()
347            if InstallType.YAML in conn.install_types
348        )
349
350    # pragma: no cover  # Should never be reached.
351    raise exc.PyAirbyteInputError(
352        message="Invalid install type.",
353        context={
354            "install_type": install_type,
355        },
356    )

263def get_connector_metadata(name: str) -> ConnectorMetadata | None:
264    """Check the cache for the connector.
265
266    If the cache is empty, populate by calling update_cache.
267    """
268    registry_url = _get_registry_url()
269
270    if _is_registry_disabled(registry_url):
271        return None
272
273    cache = copy(_get_registry_cache())
274
275    if not cache:
276        raise exc.PyAirbyteInternalError(
277            message="Connector registry could not be loaded.",
278            context={
279                "registry_url": _get_registry_url(),
280            },
281        )
282    if name not in cache:
283        raise exc.AirbyteConnectorNotRegisteredError(
284            connector_name=name,
285            context={
286                "registry_url": _get_registry_url(),
287                "available_connectors": get_available_connectors(),
288            },
289        )
290    return cache[name]

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