🎉 CDK: support loading spec from yaml file (#12104)

* support loading spec from yaml file

* formatting

* remove commented code

* update comment

* remove unused file

* raise correct exception types

* bump version, update changelog

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.

airbyte-cdk/python/airbyte_cdk/connector.py

@@ -10,9 +10,18 @@ import pkgutil
 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 @@ class Connector(ABC):
     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:

airbyte-cdk/python/setup.py

@@ -15,7 +15,7 @@ README = (HERE / "README.md").read_text()
 
 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",

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)