CDK: private configuration option _limit and _page_size (#5617)

* CDK: private configuration option _limit and _page_size

Author: avida

airbyte-cdk/python/CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.1.15
+Add \_limit and \_page_size as internal config parameters for SAT
+
 ## 0.1.14
 If the input config file does not comply with spec schema, raise an exception instead of `system.exit`.
 

airbyte-cdk/python/airbyte_cdk/entrypoint.py

@@ -33,7 +33,7 @@
 from airbyte_cdk.logger import AirbyteLogger
 from airbyte_cdk.models import AirbyteMessage, Status, Type
 from airbyte_cdk.sources import Source
-from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit
+from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit, split_config
 
 logger = AirbyteLogger()
 
@@ -90,7 +90,12 @@ def run(self, parsed_args: argparse.Namespace) -> Iterable[str]:
  else:
  raw_config = self.source.read_config(parsed_args.config)
  config = self.source.configure(raw_config, temp_dir)
+ # Remove internal flags from config before validating so
+ # jsonschema's additionalProperties flag wont fail the validation
+ config, internal_config = split_config(config)
  check_config_against_spec_or_exit(config, source_spec, logger)
+ # Put internal flags back to config dict
+ config.update(internal_config.dict())
 
  if cmd == "check":
  check_result = self.source.check(logger, config)

airbyte-cdk/python/airbyte_cdk/sources/abstract_source.py

@@ -43,6 +43,8 @@
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.source import Source
 from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.sources.streams.http.http import HttpStream
+from airbyte_cdk.sources.utils.schema_helpers import InternalConfig, split_config
 
 
 class AbstractSource(Source, ABC):
@@ -95,6 +97,7 @@ def read(
  """Implements the Read operation from the Airbyte Specification. See https://docs.airbyte.io/architecture/airbyte-specification."""
  connector_state = copy.deepcopy(state or {})
  logger.info(f"Starting syncing {self.name}")
+ config, internal_config = split_config(config)
  # TODO assert all streams exist in the connector
  # get the streams once in case the connector needs to make any queries to generate them
  stream_instances = {s.name: s for s in self.streams(config)}
@@ -107,7 +110,11 @@ def read(
 
  try:
  yield from self._read_stream(
- logger=logger, stream_instance=stream_instance, configured_stream=configured_stream, connector_state=connector_state
+ logger=logger,
+ stream_instance=stream_instance,
+ configured_stream=configured_stream,
+ connector_state=connector_state,
+ internal_config=internal_config,
  )
  except Exception as e:
  logger.exception(f"Encountered an exception while reading stream {self.name}")
@@ -121,8 +128,13 @@ def _read_stream(
  stream_instance: Stream,
  configured_stream: ConfiguredAirbyteStream,
  connector_state: MutableMapping[str, Any],
+ internal_config: InternalConfig,
  ) -> Iterator[AirbyteMessage]:
 
+ if internal_config.page_size and isinstance(stream_instance, HttpStream):
+ logger.info(f"Setting page size for {stream_instance.name} to {internal_config.page_size}")
+ stream_instance.page_size = internal_config.page_size
+
  use_incremental = configured_stream.sync_mode == SyncMode.incremental and stream_instance.supports_incremental
  if use_incremental:
  record_iterator = self._read_incremental(logger, stream_instance, configured_stream, connector_state)
@@ -135,6 +147,9 @@ def _read_stream(
  for record in record_iterator:
  if record.type == MessageType.RECORD:
  record_counter += 1
+ if internal_config.limit and record_counter > internal_config.limit:
+ logger.info(f"Reached limit defined by internal config ({internal_config.limit}), stop reading")
+ break
  yield record
 
  logger.info(f"Read {record_counter} records from {stream_name} stream")

airbyte-cdk/python/airbyte_cdk/sources/streams/http/http.py

@@ -44,6 +44,7 @@ class HttpStream(Stream, ABC):
  """
 
  source_defined_cursor = True # Most HTTP streams use a source defined cursor (i.e: the user can't configure it like on a SQL table)
+ page_size = None # Use this variable to define page size for API http requests with pagination support
 
  def __init__(self, authenticator: HttpAuthenticator = NoAuth()):
  self._authenticator = authenticator

airbyte-cdk/python/airbyte_cdk/sources/utils/schema_helpers.py

@@ -26,13 +26,14 @@
 import json
 import os
 import pkgutil
-from typing import Any, Dict, Mapping
+from typing import Any, ClassVar, Dict, Mapping, Tuple
 
 import pkg_resources
 from airbyte_cdk.logger import AirbyteLogger
 from airbyte_cdk.models import ConnectorSpecification
 from jsonschema import RefResolver, validate
 from jsonschema.exceptions import ValidationError
+from pydantic import BaseModel, Field
 
 
 class JsonSchemaResolver:
@@ -142,3 +143,31 @@ def check_config_against_spec_or_exit(config: Mapping[str, Any], spec: Connector
  validate(instance=config, schema=spec_schema)
  except ValidationError as validation_error:
  raise Exception("Config validation error: " + validation_error.message)
+
+
+class InternalConfig(BaseModel):
+ KEYWORDS: ClassVar[set] = {"_limit", "_page_size"}
+ limit: int = Field(None, alias="_limit")
+ page_size: int = Field(None, alias="_page_size")
+
+ def dict(self):
+ return super().dict(by_alias=True, exclude_unset=True)
+
+
+def split_config(config: Mapping[str, Any]) -> Tuple[dict, InternalConfig]:
+ """
+ Break config map object into 2 instances: first is a dict with user defined
+ configuration and second is internal config that contains private keys for
+ acceptance test configuration.
+ :param config - Dict object that has been loaded from config file.
+ :return tuple of user defined config dict with filtered out internal
+ parameters and SAT internal config object.
+ """
+ main_config = {}
+ internal_config = {}
+ for k, v in config.items():
+ if k in InternalConfig.KEYWORDS:
+ internal_config[k] = v
+ else:
+ main_config[k] = v
+ return main_config, InternalConfig.parse_obj(internal_config)

airbyte-cdk/python/docs/tutorials/cdk-tutorial-python-http/3-define-inputs.md

@@ -35,6 +35,11 @@ Given that we'll pulling currency data for our example source, we'll define the
  }
 }
 ```
+Beside regular parameter there is intenal CDK config that started with '_' character and used mainly for testing purposes:
+
+* _limit - set maximum number of records being read for each stream
+* _page_size - for http based streams set number of records for each page. Depends on stream implementation.
+
 
 In addition to metadata, we define two inputs:
 

airbyte-cdk/python/setup.py

@@ -35,7 +35,7 @@
 
 setup(
  name="airbyte-cdk",
- version="0.1.14",
+ version="0.1.15",
  description="A framework for writing Airbyte Connectors.",
  long_description=README,
  long_description_content_type="text/markdown",

airbyte-cdk/python/unit_tests/sources/test_source.py

@@ -26,11 +26,14 @@
 import json
 import tempfile
 from typing import Any, Mapping, MutableMapping
+from unittest.mock import MagicMock
 
 import pytest
 from airbyte_cdk.logger import AirbyteLogger
 from airbyte_cdk.models import ConfiguredAirbyteCatalog
-from airbyte_cdk.sources import Source
+from airbyte_cdk.sources import AbstractSource, Source
+from airbyte_cdk.sources.streams.core import Stream
+from airbyte_cdk.sources.streams.http.http import HttpStream
 
 
 class MockSource(Source):
@@ -51,6 +54,40 @@ def source():
  return MockSource()
 
 
+@pytest.fixture
+def abstract_source(mocker):
+ mocker.patch.multiple(HttpStream, __abstractmethods__=set())
+ mocker.patch.multiple(Stream, __abstractmethods__=set())
+
+ class MockHttpStream(MagicMock, HttpStream):
+ url_base = "http://example.com"
+ path = "/dummy/path"
+
+ def __init__(self, *args, **kvargs):
+ MagicMock.__init__(self)
+ HttpStream.__init__(self, *args, kvargs)
+ self.read_records = MagicMock()
+
+ class MockStream(MagicMock, Stream):
+ page_size = None
+
+ def __init__(self, *args, **kvargs):
+ MagicMock.__init__(self)
+ self.read_records = MagicMock()
+
+ streams = [MockHttpStream(), MockStream()]
+
+ class MockAbstractSource(AbstractSource):
+ def check_connection(self):
+ return True, None
+
+ def streams(self, config):
+ self.streams_config = config
+ return streams
+
+ return MockAbstractSource()
+
+
 def test_read_state(source):
  state = {"updated_at": "yesterday"}
 
@@ -81,3 +118,60 @@ def test_read_catalog(source):
  catalog_file.flush()
  actual = source.read_catalog(catalog_file.name)
  assert actual == expected
+
+
+def test_internal_config(abstract_source):
+ configured_catalog = {
+ "streams": [
+ {
+ "stream": {"name": "mock_http_stream", "json_schema": {}},
+ "destination_sync_mode": "overwrite",
+ "sync_mode": "full_refresh",
+ },
+ {
+ "stream": {"name": "mock_stream", "json_schema": {}},
+ "destination_sync_mode": "overwrite",
+ "sync_mode": "full_refresh",
+ },
+ ]
+ }
+ catalog = ConfiguredAirbyteCatalog.parse_obj(configured_catalog)
+ streams = abstract_source.streams(None)
+ assert len(streams) == 2
+ http_stream = streams[0]
+ non_http_stream = streams[1]
+ assert isinstance(http_stream, HttpStream)
+ assert not isinstance(non_http_stream, HttpStream)
+ http_stream.read_records.return_value = [{}] * 3
+ non_http_stream.read_records.return_value = [{}] * 3
+
+ # Test with empty config
+ records = [r for r in abstract_source.read(logger=MagicMock(), config={}, catalog=catalog, state={})]
+ # 3 for http stream and 3 for non http stream
+ assert len(records) == 3 + 3
+ assert http_stream.read_records.called
+ assert non_http_stream.read_records.called
+ # Make sure page_size havent been set
+ assert not http_stream.page_size
+ assert not non_http_stream.page_size
+ # Test with records limit set to 1
+ internal_config = {"some_config": 100, "_limit": 1}
+ records = [r for r in abstract_source.read(logger=MagicMock(), config=internal_config, catalog=catalog, state={})]
+ # 1 from http stream + 1 from non http stream
+ assert len(records) == 1 + 1
+ assert "_limit" not in abstract_source.streams_config
+ assert "some_config" in abstract_source.streams_config
+ # Test with records limit set to number that exceeds expceted records
+ internal_config = {"some_config": 100, "_limit": 20}
+ records = [r for r in abstract_source.read(logger=MagicMock(), config=internal_config, catalog=catalog, state={})]
+ assert len(records) == 3 + 3
+
+ # Check if page_size paramter is set to http instance only
+ internal_config = {"some_config": 100, "_page_size": 2}
+ records = [r for r in abstract_source.read(logger=MagicMock(), config=internal_config, catalog=catalog, state={})]
+ assert "_page_size" not in abstract_source.streams_config
+ assert "some_config" in abstract_source.streams_config
+ assert len(records) == 3 + 3
+ assert http_stream.page_size == 2
+ # Make sure page_size havent been set for non http streams
+ assert not non_http_stream.page_size

airbyte-cdk/python/unit_tests/test_entrypoint.py

@@ -144,6 +144,11 @@ def config_mock(mocker, request):
  ({"username": "fake"}, {"type": "object", "properties": {"name": {"type": "string"}}, "additionalProperties": False}, False),
  ({"username": "fake"}, {"type": "object", "properties": {"username": {"type": "string"}}, "additionalProperties": False}, True),
  ({"username": "fake"}, {"type": "object", "properties": {"user": {"type": "string"}}}, True),
+ (
+ {"username": "fake", "_limit": 22},
+ {"type": "object", "properties": {"username": {"type": "string"}}, "additionalProperties": False},
+ True,
+ ),
  ],
  indirect=["config_mock"],
 )