[WIP] destination discover PoC

Author: maxi297

airbyte_cdk/destinations/destination.py

@@ -14,6 +14,7 @@
 from airbyte_cdk.connector import Connector
 from airbyte_cdk.exception_handler import init_uncaught_exception_handler
 from airbyte_cdk.models import (
+    AirbyteCatalog,
     AirbyteMessage,
     AirbyteMessageSerializer,
     ConfiguredAirbyteCatalog,
@@ -26,8 +27,70 @@
 logger = logging.getLogger("airbyte")
 
 
+def parse_args(args: List[str]) -> argparse.Namespace:
+    """
+    :param args: commandline arguments
+    :return:
+    """
+
+    parent_parser = argparse.ArgumentParser(add_help=False)
+    parent_parser.add_argument("--debug", action="store_true", help="enables detailed debug logs related to the sync")
+    main_parser = argparse.ArgumentParser()
+    subparsers = main_parser.add_subparsers(title="commands", dest="command")
+
+    # spec
+    subparsers.add_parser(
+        "spec", help="outputs the json configuration specification", parents=[parent_parser]
+    )
+
+    # check
+    check_parser = subparsers.add_parser(
+        "check", help="checks the config can be used to connect", parents=[parent_parser]
+    )
+    required_check_parser = check_parser.add_argument_group("required named arguments")
+    required_check_parser.add_argument(
+        "--config", type=str, required=True, help="path to the json configuration file"
+    )
+
+    # discover
+    discover_parser = subparsers.add_parser(
+        "discover", help="discover the objects available in the destination", parents=[parent_parser]
+    )
+    required_discover_parser = discover_parser.add_argument_group("required named arguments")
+    required_discover_parser.add_argument(
+        "--config", type=str, required=True, help="path to the json configuration file"
+    )
+
+    # write
+    write_parser = subparsers.add_parser(
+        "write", help="Writes data to the destination", parents=[parent_parser]
+    )
+    write_required = write_parser.add_argument_group("required named arguments")
+    write_required.add_argument(
+        "--config", type=str, required=True, help="path to the JSON configuration file"
+    )
+    write_required.add_argument(
+        "--catalog", type=str, required=True, help="path to the configured catalog JSON file"
+    )
+
+    parsed_args = main_parser.parse_args(args)
+    cmd = parsed_args.command
+    if not cmd:
+        raise Exception("No command entered. ")
+    elif cmd not in ["spec", "check", "discover", "write"]:
+        # This is technically dead code since parse_args() would fail if this was the case
+        # But it's non-obvious enough to warrant placing it here anyways
+        raise Exception(f"Unknown command entered: {cmd}")
+
+    return parsed_args
+
+
 class Destination(Connector, ABC):
-    VALID_CMDS = {"spec", "check", "write"}
+    VALID_CMDS = {"spec", "check", "discover", "write"}
+
+    def discover(self) -> AirbyteCatalog:
+        """Implement to define what objects are available in the destination"""
+        raise NotImplementedError("Discover method is not implemented")
 
     @abstractmethod
     def write(
@@ -68,52 +131,9 @@ def _run_write(
         )
         logger.info("Writing complete.")
 
-    def parse_args(self, args: List[str]) -> argparse.Namespace:
-        """
-        :param args: commandline arguments
-        :return:
-        """
-
-        parent_parser = argparse.ArgumentParser(add_help=False)
-        main_parser = argparse.ArgumentParser()
-        subparsers = main_parser.add_subparsers(title="commands", dest="command")
-
-        # spec
-        subparsers.add_parser(
-            "spec", help="outputs the json configuration specification", parents=[parent_parser]
-        )
-
-        # check
-        check_parser = subparsers.add_parser(
-            "check", help="checks the config can be used to connect", parents=[parent_parser]
-        )
-        required_check_parser = check_parser.add_argument_group("required named arguments")
-        required_check_parser.add_argument(
-            "--config", type=str, required=True, help="path to the json configuration file"
-        )
-
-        # write
-        write_parser = subparsers.add_parser(
-            "write", help="Writes data to the destination", parents=[parent_parser]
-        )
-        write_required = write_parser.add_argument_group("required named arguments")
-        write_required.add_argument(
-            "--config", type=str, required=True, help="path to the JSON configuration file"
-        )
-        write_required.add_argument(
-            "--catalog", type=str, required=True, help="path to the configured catalog JSON file"
-        )
-
-        parsed_args = main_parser.parse_args(args)
-        cmd = parsed_args.command
-        if not cmd:
-            raise Exception("No command entered. ")
-        elif cmd not in ["spec", "check", "write"]:
-            # This is technically dead code since parse_args() would fail if this was the case
-            # But it's non-obvious enough to warrant placing it here anyways
-            raise Exception(f"Unknown command entered: {cmd}")
-
-        return parsed_args
+    @staticmethod
+    def parse_args(args: List[str]) -> argparse.Namespace:
+        return parse_args(args)
 
     def run_cmd(self, parsed_args: argparse.Namespace) -> Iterable[AirbyteMessage]:
         cmd = parsed_args.command
@@ -137,6 +157,8 @@ def run_cmd(self, parsed_args: argparse.Namespace) -> Iterable[AirbyteMessage]:
 
         if cmd == "check":
             yield self._run_check(config=config)
+        elif cmd == "discover":
+            yield AirbyteMessage(type=Type.CATALOG, catalog=self.discover())
         elif cmd == "write":
             # Wrap in UTF-8 to override any other input encodings
             wrapped_stdin = io.TextIOWrapper(sys.stdin.buffer, encoding="utf-8")

airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py

@@ -1,11 +1,10 @@
 #
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 #
-
-
+from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import InitVar, dataclass, field
-from typing import Any, List, Mapping, MutableMapping, Optional, Union
+from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Union
 
 import dpath
 from typing_extensions import deprecated
@@ -114,6 +113,16 @@ def _update_pointer(
         )
 
 
+@deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
+class AdditionalPropertyFieldsInferer(ABC):
+    @abstractmethod
+    def infer(self, property_definition: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
+        """
+        Infers additional property fields from the given property definition.
+        """
+        pass
+
+
 @deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
 @dataclass
 class DynamicSchemaLoader(SchemaLoader):
@@ -126,6 +135,7 @@ class DynamicSchemaLoader(SchemaLoader):
     parameters: InitVar[Mapping[str, Any]]
     schema_type_identifier: SchemaTypeIdentifier
     schema_transformations: List[RecordTransformation] = field(default_factory=lambda: [])
+    additional_property_fields_inferrer: Optional[AdditionalPropertyFieldsInferer] = None
 
     def get_json_schema(self) -> Mapping[str, Any]:
         """
@@ -149,9 +159,11 @@ def get_json_schema(self) -> Mapping[str, Any]:
                 property_definition,
                 self.schema_type_identifier.type_pointer,
             )
+
+            value.update(self.additional_property_fields_inferrer.infer(property_definition) if self.additional_property_fields_inferrer else {})
             properties[key] = value
 
-        transformed_properties = self._transform(properties, {})
+        transformed_properties = self._transform(properties)
 
         return {
             "$schema": "https://json-schema.org/draft-07/schema#",
@@ -163,8 +175,6 @@ def get_json_schema(self) -> Mapping[str, Any]:
     def _transform(
         self,
         properties: Mapping[str, Any],
-        stream_state: StreamState,
-        stream_slice: Optional[StreamSlice] = None,
     ) -> Mapping[str, Any]:
         for transformation in self.schema_transformations:
             transformation.transform(
@@ -190,7 +200,7 @@ def _get_type(
         self,
         raw_schema: MutableMapping[str, Any],
         field_type_path: Optional[List[Union[InterpolatedString, str]]],
-    ) -> Union[Mapping[str, Any], List[Mapping[str, Any]]]:
+    ) -> Dict[str, Any]:
         """
         Determines the JSON Schema type for a field, supporting nullable and combined types.
         """