diff --git a/airbyte-cdk/python/CHANGELOG.md b/airbyte-cdk/python/CHANGELOG.md index 729660e35351c3..145ef289cc7c37 100644 --- a/airbyte-cdk/python/CHANGELOG.md +++ b/airbyte-cdk/python/CHANGELOG.md @@ -1,5 +1,8 @@ # Changelog +## 0.1.55 +Add support for reading the spec from a YAML file (`spec.yaml`) + ## 0.1.54 - Add ability to import `IncrementalMixin` from `airbyte_cdk.sources.streams`. - Bumped minimum supported Python version to 3.9. diff --git a/airbyte-cdk/python/airbyte_cdk/connector.py b/airbyte-cdk/python/airbyte_cdk/connector.py index f17c76ab5754eb..f75bf0c717a786 100644 --- a/airbyte-cdk/python/airbyte_cdk/connector.py +++ b/airbyte-cdk/python/airbyte_cdk/connector.py @@ -10,9 +10,18 @@ from abc import ABC, abstractmethod from typing import Any, Mapping, Optional +import yaml from airbyte_cdk.models import AirbyteConnectionStatus, ConnectorSpecification +def load_optional_package_file(package: str, filename: str) -> Optional[bytes]: + """Gets a resource from a package, returning None if it does not exist""" + try: + return pkgutil.get_data(package, filename) + except FileNotFoundError: + return None + + class AirbyteSpec(object): @staticmethod def from_file(file_name: str): @@ -51,12 +60,25 @@ def write_config(config: Mapping[str, Any], config_path: str): def spec(self, logger: logging.Logger) -> ConnectorSpecification: """ Returns the spec for this integration. The spec is a JSON-Schema object describing the required configurations (e.g: username and password) - required to run this integration. + required to run this integration. By default, this will be loaded from a "spec.yaml" or a "spec.json" in the package root. """ - raw_spec: Optional[bytes] = pkgutil.get_data(self.__class__.__module__.split(".")[0], "spec.json") - if not raw_spec: - raise ValueError("Unable to find spec.json.") - return ConnectorSpecification.parse_obj(json.loads(raw_spec)) + + package = self.__class__.__module__.split(".")[0] + + yaml_spec = load_optional_package_file(package, "spec.yaml") + json_spec = load_optional_package_file(package, "spec.json") + + if yaml_spec and json_spec: + raise RuntimeError("Found multiple spec files in the package. Only one of spec.yaml or spec.json should be provided.") + + if yaml_spec: + spec_obj = yaml.load(yaml_spec, Loader=yaml.SafeLoader) + elif json_spec: + spec_obj = json.loads(json_spec) + else: + raise FileNotFoundError("Unable to find spec.yaml or spec.json in the package.") + + return ConnectorSpecification.parse_obj(spec_obj) @abstractmethod def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus: diff --git a/airbyte-cdk/python/setup.py b/airbyte-cdk/python/setup.py index 5be56e0ee84d37..df9eaacc46fb0c 100644 --- a/airbyte-cdk/python/setup.py +++ b/airbyte-cdk/python/setup.py @@ -15,7 +15,7 @@ setup( name="airbyte-cdk", - version="0.1.54", + version="0.1.55", description="A framework for writing Airbyte Connectors.", long_description=README, long_description_content_type="text/markdown", diff --git a/airbyte-cdk/python/unit_tests/test_connector.py b/airbyte-cdk/python/unit_tests/test_connector.py index 463b58290cba81..b1840ef08d528a 100644 --- a/airbyte-cdk/python/unit_tests/test_connector.py +++ b/airbyte-cdk/python/unit_tests/test_connector.py @@ -5,14 +5,23 @@ import json import logging +import os +import sys import tempfile from pathlib import Path from typing import Any, Mapping import pytest +import yaml from airbyte_cdk import AirbyteSpec, Connector from airbyte_cdk.models import AirbyteConnectionStatus +logger = logging.getLogger("airbyte") + +MODULE = sys.modules[__name__] +MODULE_PATH = os.path.abspath(MODULE.__file__) +SPEC_ROOT = os.path.dirname(MODULE_PATH) + class TestAirbyteSpec: VALID_SPEC = { @@ -71,3 +80,53 @@ def test_write_config(integration, mock_config): integration.write_config(mock_config, str(config_path)) with open(config_path, "r") as actual: assert mock_config == json.loads(actual.read()) + + +class TestConnectorSpec: + CONNECTION_SPECIFICATION = { + "type": "object", + "required": ["api_token"], + "additionalProperties": False, + "properties": {"api_token": {"type": "string"}}, + } + + @pytest.fixture + def use_json_spec(self): + spec = { + "documentationUrl": "https://airbyte.com/#json", + "connectionSpecification": self.CONNECTION_SPECIFICATION, + } + + json_path = os.path.join(SPEC_ROOT, "spec.json") + with open(json_path, "w") as f: + f.write(json.dumps(spec)) + yield + os.remove(json_path) + + @pytest.fixture + def use_yaml_spec(self): + spec = {"documentationUrl": "https://airbyte.com/#yaml", "connectionSpecification": self.CONNECTION_SPECIFICATION} + + yaml_path = os.path.join(SPEC_ROOT, "spec.yaml") + with open(yaml_path, "w") as f: + f.write(yaml.dump(spec)) + yield + os.remove(yaml_path) + + def test_spec_from_json_file(self, integration, use_json_spec): + connector_spec = integration.spec(logger) + assert connector_spec.documentationUrl == "https://airbyte.com/#json" + assert connector_spec.connectionSpecification == self.CONNECTION_SPECIFICATION + + def test_spec_from_yaml_file(self, integration, use_yaml_spec): + connector_spec = integration.spec(logger) + assert connector_spec.documentationUrl == "https://airbyte.com/#yaml" + assert connector_spec.connectionSpecification == self.CONNECTION_SPECIFICATION + + def test_multiple_spec_files_raises_exception(self, integration, use_yaml_spec, use_json_spec): + with pytest.raises(RuntimeError, match="spec.yaml or spec.json"): + integration.spec(logger) + + def test_no_spec_file_raises_exception(self, integration): + with pytest.raises(FileNotFoundError, match="Unable to find spec."): + integration.spec(logger)