-
Notifications
You must be signed in to change notification settings - Fork 8
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(low-code): add check dynamic stream (#223)
Co-authored-by: octavia-squidington-iii <[email protected]>
- Loading branch information
Showing
8 changed files
with
272 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,8 +1,24 @@ | ||
| # | ||
| # Copyright (c) 2023 Airbyte, Inc., all rights reserved. | ||
| # Copyright (c) 2025 Airbyte, Inc., all rights reserved. | ||
| # | ||
|
|
||
| from typing import Mapping | ||
|
|
||
| from pydantic.v1 import BaseModel | ||
|
|
||
| from airbyte_cdk.sources.declarative.checks.check_dynamic_stream import CheckDynamicStream | ||
| from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream | ||
| from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker | ||
| from airbyte_cdk.sources.declarative.models import ( | ||
| CheckDynamicStream as CheckDynamicStreamModel, | ||
| ) | ||
| from airbyte_cdk.sources.declarative.models import ( | ||
| CheckStream as CheckStreamModel, | ||
| ) | ||
|
|
||
| COMPONENTS_CHECKER_TYPE_MAPPING: Mapping[str, type[BaseModel]] = { | ||
| "CheckStream": CheckStreamModel, | ||
| "CheckDynamicStream": CheckDynamicStreamModel, | ||
| } | ||
|
|
||
| __all__ = ["CheckStream", "ConnectionChecker"] | ||
| __all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker"] |
51 changes: 51 additions & 0 deletions
51
airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| # | ||
| # Copyright (c) 2025 Airbyte, Inc., all rights reserved. | ||
| # | ||
|
|
||
| import logging | ||
| import traceback | ||
| from dataclasses import InitVar, dataclass | ||
| from typing import Any, List, Mapping, Tuple | ||
|
|
||
| from airbyte_cdk import AbstractSource | ||
| from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker | ||
| from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy | ||
|
|
||
|
|
||
| @dataclass | ||
| class CheckDynamicStream(ConnectionChecker): | ||
| """ | ||
| Checks the connections by checking availability of one or many dynamic streams | ||
| Attributes: | ||
| stream_count (int): numbers of streams to check | ||
| """ | ||
|
|
||
| stream_count: int | ||
| parameters: InitVar[Mapping[str, Any]] | ||
|
|
||
| def __post_init__(self, parameters: Mapping[str, Any]) -> None: | ||
| self._parameters = parameters | ||
|
|
||
| def check_connection( | ||
| self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any] | ||
| ) -> Tuple[bool, Any]: | ||
| streams = source.streams(config=config) | ||
| if len(streams) == 0: | ||
| return False, f"No streams to connect to from source {source}" | ||
|
|
||
| for stream_index in range(min(self.stream_count, len(streams))): | ||
| stream = streams[stream_index] | ||
| availability_strategy = HttpAvailabilityStrategy() | ||
| try: | ||
| stream_is_available, reason = availability_strategy.check_availability( | ||
| stream, logger | ||
| ) | ||
| if not stream_is_available: | ||
| return False, reason | ||
| except Exception as error: | ||
| logger.error( | ||
| f"Encountered an error trying to connect to stream {stream.name}. Error: \n {traceback.format_exc()}" | ||
| ) | ||
| return False, f"Unable to connect to stream {stream.name} - {error}" | ||
| return True, None |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
159 changes: 159 additions & 0 deletions
159
unit_tests/sources/declarative/checks/test_check_dynamic_stream.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,159 @@ | ||
| # | ||
| # Copyright (c) 2025 Airbyte, Inc., all rights reserved. | ||
| # | ||
|
|
||
| import json | ||
| import logging | ||
|
|
||
| import pytest | ||
|
|
||
| from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( | ||
| ConcurrentDeclarativeSource, | ||
| ) | ||
| from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse | ||
|
|
||
| logger = logging.getLogger("test") | ||
|
|
||
| _CONFIG = {"start_date": "2024-07-01T00:00:00.000Z"} | ||
|
|
||
| _MANIFEST = { | ||
| "version": "6.7.0", | ||
| "type": "DeclarativeSource", | ||
| "check": {"type": "CheckDynamicStream", "stream_count": 1}, | ||
| "dynamic_streams": [ | ||
| { | ||
| "type": "DynamicDeclarativeStream", | ||
| "stream_template": { | ||
| "type": "DeclarativeStream", | ||
| "name": "", | ||
| "primary_key": [], | ||
| "schema_loader": { | ||
| "type": "InlineSchemaLoader", | ||
| "schema": { | ||
| "$schema": "http://json-schema.org/schema#", | ||
| "properties": { | ||
| "ABC": {"type": "number"}, | ||
| "AED": {"type": "number"}, | ||
| }, | ||
| "type": "object", | ||
| }, | ||
| }, | ||
| "retriever": { | ||
| "type": "SimpleRetriever", | ||
| "requester": { | ||
| "type": "HttpRequester", | ||
| "$parameters": {"item_id": ""}, | ||
| "url_base": "https://api.test.com", | ||
| "path": "/items/{{parameters['item_id']}}", | ||
| "http_method": "GET", | ||
| "authenticator": { | ||
| "type": "ApiKeyAuthenticator", | ||
| "header": "apikey", | ||
| "api_token": "{{ config['api_key'] }}", | ||
| }, | ||
| }, | ||
| "record_selector": { | ||
| "type": "RecordSelector", | ||
| "extractor": {"type": "DpathExtractor", "field_path": []}, | ||
| }, | ||
| "paginator": {"type": "NoPagination"}, | ||
| }, | ||
| }, | ||
| "components_resolver": { | ||
| "type": "HttpComponentsResolver", | ||
| "retriever": { | ||
| "type": "SimpleRetriever", | ||
| "requester": { | ||
| "type": "HttpRequester", | ||
| "url_base": "https://api.test.com", | ||
| "path": "items", | ||
| "http_method": "GET", | ||
| "authenticator": { | ||
| "type": "ApiKeyAuthenticator", | ||
| "header": "apikey", | ||
| "api_token": "{{ config['api_key'] }}", | ||
| }, | ||
| }, | ||
| "record_selector": { | ||
| "type": "RecordSelector", | ||
| "extractor": {"type": "DpathExtractor", "field_path": []}, | ||
| }, | ||
| "paginator": {"type": "NoPagination"}, | ||
| }, | ||
| "components_mapping": [ | ||
| { | ||
| "type": "ComponentMappingDefinition", | ||
| "field_path": ["name"], | ||
| "value": "{{components_values['name']}}", | ||
| }, | ||
| { | ||
| "type": "ComponentMappingDefinition", | ||
| "field_path": [ | ||
| "retriever", | ||
| "requester", | ||
| "$parameters", | ||
| "item_id", | ||
| ], | ||
| "value": "{{components_values['id']}}", | ||
| }, | ||
| ], | ||
| }, | ||
| } | ||
| ], | ||
| } | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| "response_code, available_expectation, expected_messages", | ||
| [ | ||
| pytest.param( | ||
| 404, | ||
| False, | ||
| ["Not found. The requested resource was not found on the server."], | ||
| id="test_stream_unavailable_unhandled_error", | ||
| ), | ||
| pytest.param( | ||
| 403, | ||
| False, | ||
| ["Forbidden. You don't have permission to access this resource."], | ||
| id="test_stream_unavailable_handled_error", | ||
| ), | ||
| pytest.param(200, True, [], id="test_stream_available"), | ||
| pytest.param( | ||
| 401, | ||
| False, | ||
| ["Unauthorized. Please ensure you are authenticated correctly."], | ||
| id="test_stream_unauthorized_error", | ||
| ), | ||
| ], | ||
| ) | ||
| def test_check_dynamic_stream(response_code, available_expectation, expected_messages): | ||
| with HttpMocker() as http_mocker: | ||
| http_mocker.get( | ||
| HttpRequest(url="https://api.test.com/items"), | ||
| HttpResponse( | ||
| body=json.dumps( | ||
| [ | ||
| {"id": 1, "name": "item_1"}, | ||
| {"id": 2, "name": "item_2"}, | ||
| ] | ||
| ) | ||
| ), | ||
| ) | ||
| http_mocker.get( | ||
| HttpRequest(url="https://api.test.com/items/1"), | ||
| HttpResponse(body=json.dumps(expected_messages), status_code=response_code), | ||
| ) | ||
|
|
||
| source = ConcurrentDeclarativeSource( | ||
| source_config=_MANIFEST, | ||
| config=_CONFIG, | ||
| catalog=None, | ||
| state=None, | ||
| ) | ||
|
|
||
| stream_is_available, reason = source.check_connection(logger, _CONFIG) | ||
|
|
||
| assert stream_is_available == available_expectation | ||
| for message in expected_messages: | ||
| assert message in reason |