feat(low-code cdk): add dynamic schema loader (#104)

Co-authored-by: octavia-squidington-iii <[email protected]>

Author: lazebnyi

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1218,6 +1218,7 @@ definitions:
         title: Schema Loader
         description: Component used to retrieve the schema for the current stream.
         anyOf:
+          - "$ref": "#/definitions/DynamicSchemaLoader"
           - "$ref": "#/definitions/InlineSchemaLoader"
           - "$ref": "#/definitions/JsonFileSchemaLoader"
           - "$ref": "#/definitions/CustomSchemaLoader"
@@ -1684,6 +1685,92 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  TypesMap:
+    title: Types Map
+    description: (This component is experimental. Use at your own risk.) Represents a mapping between a current type and its corresponding target type.
+    type: object
+    required:
+      - target_type
+      - current_type
+    properties:
+      target_type:
+        anyOf:
+          - type: string
+          - type: array
+            items:
+              type: string
+      current_type:
+        anyOf:
+          - type: string
+          - type: array
+            items:
+              type: string
+  SchemaTypeIdentifier:
+    title: Schema Type Identifier
+    description: (This component is experimental. Use at your own risk.) Identifies schema details for dynamic schema extraction and processing.
+    type: object
+    required:
+      - key_pointer
+    properties:
+      type:
+        type: string
+        enum: [SchemaTypeIdentifier]
+      schema_pointer:
+        title: Schema Path
+        description: List of nested fields defining the schema field path to extract. Defaults to [].
+        type: array
+        default: []
+        items:
+          - type: string
+        interpolation_context:
+          - config
+      key_pointer:
+        title: Key Path
+        description: List of potentially nested fields describing the full path of the field key to extract.
+        type: array
+        items:
+          - type: string
+        interpolation_context:
+          - config
+      type_pointer:
+        title: Type Path
+        description: List of potentially nested fields describing the full path of the field type to extract.
+        type: array
+        items:
+          - type: string
+        interpolation_context:
+          - config
+      types_mapping:
+        type: array
+        items:
+          - "$ref": "#/definitions/TypesMap"
+      $parameters:
+        type: object
+        additionalProperties: true
+  DynamicSchemaLoader:
+    title: Dynamic Schema Loader
+    description: (This component is experimental. Use at your own risk.) Loads a schema by extracting data from retrieved records.
+    type: object
+    required:
+      - type
+      - retriever
+      - schema_type_identifier
+    properties:
+      type:
+        type: string
+        enum: [DynamicSchemaLoader]
+      retriever:
+        title: Retriever
+        description: Component used to coordinate how records are extracted across stream slices and request pages.
+        anyOf:
+          - "$ref": "#/definitions/AsyncRetriever"
+          - "$ref": "#/definitions/CustomRetriever"
+          - "$ref": "#/definitions/SimpleRetriever"
+      schema_type_identifier:
+        "$ref": "#/definitions/SchemaTypeIdentifier"
+      $parameters:
+        type: object
+        additionalProperties: true
   InlineSchemaLoader:
     title: Inline Schema Loader
     description: Loads a schema that is defined directly in the manifest file.

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -650,6 +650,32 @@ class HttpResponseFilter(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class TypesMap(BaseModel):
+    target_type: Union[str, List[str]]
+    current_type: Union[str, List[str]]
+
+
+class SchemaTypeIdentifier(BaseModel):
+    type: Optional[Literal["SchemaTypeIdentifier"]] = None
+    schema_pointer: Optional[List[str]] = Field(
+        [],
+        description="List of nested fields defining the schema field path to extract. Defaults to [].",
+        title="Schema Path",
+    )
+    key_pointer: List[str] = Field(
+        ...,
+        description="List of potentially nested fields describing the full path of the field key to extract.",
+        title="Key Path",
+    )
+    type_pointer: Optional[List[str]] = Field(
+        None,
+        description="List of potentially nested fields describing the full path of the field type to extract.",
+        title="Type Path",
+    )
+    types_mapping: Optional[List[TypesMap]] = None
+    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
+
+
 class InlineSchemaLoader(BaseModel):
     type: Literal["InlineSchemaLoader"]
     schema_: Optional[Dict[str, Any]] = Field(
@@ -822,13 +848,13 @@ class Config:
     )
     extract_output: List[str] = Field(
         ...,
-        description="The DeclarativeOAuth Specific list of strings to indicate which keys should be extracted and returned back to the input config.                ",
+        description="The DeclarativeOAuth Specific list of strings to indicate which keys should be extracted and returned back to the input config.",
         examples=[{"extract_output": ["access_token", "refresh_token", "other_field"]}],
         title="DeclarativeOAuth Extract Output",
     )
     state: Optional[State] = Field(
         None,
-        description="The DeclarativeOAuth Specific object to provide the criteria of how the `state` query param should be constructed,\nincluding length and complexity.                ",
+        description="The DeclarativeOAuth Specific object to provide the criteria of how the `state` query param should be constructed,\nincluding length and complexity.",
         examples=[{"state": {"min": 7, "max": 128}}],
         title="(Optional) DeclarativeOAuth Configurable State Query Param",
     )
@@ -852,13 +878,13 @@ class Config:
     )
     state_key: Optional[str] = Field(
         None,
-        description="The DeclarativeOAuth Specific optional override to provide the custom `state` key name, if required by data-provider.                ",
+        description="The DeclarativeOAuth Specific optional override to provide the custom `state` key name, if required by data-provider.",
         examples=[{"state_key": "my_custom_state_key_key_name"}],
         title="(Optional) DeclarativeOAuth State Key Override",
     )
     auth_code_key: Optional[str] = Field(
         None,
-        description="The DeclarativeOAuth Specific optional override to provide the custom `code` key name to something like `auth_code` or `custom_auth_code`, if required by data-provider.                ",
+        description="The DeclarativeOAuth Specific optional override to provide the custom `code` key name to something like `auth_code` or `custom_auth_code`, if required by data-provider.",
         examples=[{"auth_code_key": "my_custom_auth_code_key_name"}],
         title="(Optional) DeclarativeOAuth Auth Code Key Override",
     )
@@ -1609,12 +1635,17 @@ class Config:
     primary_key: Optional[PrimaryKey] = Field(
         "", description="The primary key of the stream.", title="Primary Key"
     )
-    schema_loader: Optional[Union[InlineSchemaLoader, JsonFileSchemaLoader, CustomSchemaLoader]] = (
-        Field(
-            None,
-            description="Component used to retrieve the schema for the current stream.",
-            title="Schema Loader",
-        )
+    schema_loader: Optional[
+        Union[
+            DynamicSchemaLoader,
+            InlineSchemaLoader,
+            JsonFileSchemaLoader,
+            CustomSchemaLoader,
+        ]
+    ] = Field(
+        None,
+        description="Component used to retrieve the schema for the current stream.",
+        title="Schema Loader",
     )
     transformations: Optional[
         List[Union[AddFields, CustomTransformation, RemoveFields, KeysToLower]]
@@ -1774,6 +1805,17 @@ class HttpRequester(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class DynamicSchemaLoader(BaseModel):
+    type: Literal["DynamicSchemaLoader"]
+    retriever: Union[AsyncRetriever, CustomRetriever, SimpleRetriever] = Field(
+        ...,
+        description="Component used to coordinate how records are extracted across stream slices and request pages.",
+        title="Retriever",
+    )
+    schema_type_identifier: SchemaTypeIdentifier
+    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
+
+
 class ParentStreamConfig(BaseModel):
     type: Literal["ParentStreamConfig"]
     parent_key: str = Field(
@@ -1981,5 +2023,6 @@ class DynamicDeclarativeStream(BaseModel):
 SelectiveAuthenticator.update_forward_refs()
 DeclarativeStream.update_forward_refs()
 SessionTokenAuthenticator.update_forward_refs()
+DynamicSchemaLoader.update_forward_refs()
 SimpleRetriever.update_forward_refs()
 AsyncRetriever.update_forward_refs()

airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py

@@ -64,6 +64,10 @@
     "AddFields.fields": "AddedFieldDefinition",
     # CustomPartitionRouter
     "CustomPartitionRouter.parent_stream_configs": "ParentStreamConfig",
+    # DynamicSchemaLoader
+    "DynamicSchemaLoader.retriever": "SimpleRetriever",
+    # SchemaTypeIdentifier
+    "SchemaTypeIdentifier.types_map": "TypesMap",
 }
 
 # We retain a separate registry for custom components to automatically insert the type if it is missing. This is intended to

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -188,6 +188,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     DpathExtractor as DpathExtractorModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    DynamicSchemaLoader as DynamicSchemaLoaderModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ExponentialBackoffStrategy as ExponentialBackoffStrategyModel,
 )
@@ -278,6 +281,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ResponseToFileExtractor as ResponseToFileExtractorModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    SchemaTypeIdentifier as SchemaTypeIdentifierModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     SelectiveAuthenticator as SelectiveAuthenticatorModel,
 )
@@ -291,6 +297,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     SubstreamPartitionRouter as SubstreamPartitionRouterModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    TypesMap as TypesMapModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import ValueType
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     WaitTimeFromHeader as WaitTimeFromHeaderModel,
@@ -356,8 +365,11 @@
 )
 from airbyte_cdk.sources.declarative.schema import (
     DefaultSchemaLoader,
+    DynamicSchemaLoader,
     InlineSchemaLoader,
     JsonFileSchemaLoader,
+    SchemaTypeIdentifier,
+    TypesMap,
 )
 from airbyte_cdk.sources.declarative.spec import Spec
 from airbyte_cdk.sources.declarative.stream_slicers import StreamSlicer
@@ -455,6 +467,9 @@ def _init_mappings(self) -> None:
             IterableDecoderModel: self.create_iterable_decoder,
             XmlDecoderModel: self.create_xml_decoder,
             JsonFileSchemaLoaderModel: self.create_json_file_schema_loader,
+            DynamicSchemaLoaderModel: self.create_dynamic_schema_loader,
+            SchemaTypeIdentifierModel: self.create_schema_type_identifier,
+            TypesMapModel: self.create_types_map,
             JwtAuthenticatorModel: self.create_jwt_authenticator,
             LegacyToPerPartitionStateMigrationModel: self.create_legacy_to_per_partition_state_migration,
             ListPartitionRouterModel: self.create_list_partition_router,
@@ -1574,6 +1589,63 @@ def create_inline_schema_loader(
     ) -> InlineSchemaLoader:
         return InlineSchemaLoader(schema=model.schema_ or {}, parameters={})
 
+    @staticmethod
+    def create_types_map(model: TypesMapModel, **kwargs: Any) -> TypesMap:
+        return TypesMap(target_type=model.target_type, current_type=model.current_type)
+
+    def create_schema_type_identifier(
+        self, model: SchemaTypeIdentifierModel, config: Config, **kwargs: Any
+    ) -> SchemaTypeIdentifier:
+        types_mapping = []
+        if model.types_mapping:
+            types_mapping.extend(
+                [
+                    self._create_component_from_model(types_map, config=config)
+                    for types_map in model.types_mapping
+                ]
+            )
+        model_schema_pointer: List[Union[InterpolatedString, str]] = (
+            [x for x in model.schema_pointer] if model.schema_pointer else []
+        )
+        model_key_pointer: List[Union[InterpolatedString, str]] = [x for x in model.key_pointer]
+        model_type_pointer: Optional[List[Union[InterpolatedString, str]]] = (
+            [x for x in model.type_pointer] if model.type_pointer else None
+        )
+
+        return SchemaTypeIdentifier(
+            schema_pointer=model_schema_pointer,
+            key_pointer=model_key_pointer,
+            type_pointer=model_type_pointer,
+            types_mapping=types_mapping,
+            parameters=model.parameters or {},
+        )
+
+    def create_dynamic_schema_loader(
+        self, model: DynamicSchemaLoaderModel, config: Config, **kwargs: Any
+    ) -> DynamicSchemaLoader:
+        stream_slicer = self._build_stream_slicer_from_partition_router(model.retriever, config)
+        combined_slicers = self._build_resumable_cursor_from_paginator(
+            model.retriever, stream_slicer
+        )
+
+        retriever = self._create_component_from_model(
+            model=model.retriever,
+            config=config,
+            name="",
+            primary_key=None,
+            stream_slicer=combined_slicers,
+            transformations=[],
+        )
+        schema_type_identifier = self._create_component_from_model(
+            model.schema_type_identifier, config=config, parameters=model.parameters or {}
+        )
+        return DynamicSchemaLoader(
+            retriever=retriever,
+            config=config,
+            schema_type_identifier=schema_type_identifier,
+            parameters=model.parameters or {},
+        )
+
     @staticmethod
     def create_json_decoder(model: JsonDecoderModel, config: Config, **kwargs: Any) -> JsonDecoder:
         return JsonDecoder(parameters={})

airbyte_cdk/sources/declarative/schema/__init__.py

@@ -6,5 +6,6 @@
 from airbyte_cdk.sources.declarative.schema.inline_schema_loader import InlineSchemaLoader
 from airbyte_cdk.sources.declarative.schema.json_file_schema_loader import JsonFileSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
+from airbyte_cdk.sources.declarative.schema.dynamic_schema_loader import DynamicSchemaLoader, TypesMap, SchemaTypeIdentifier
 
-__all__ = ["JsonFileSchemaLoader", "DefaultSchemaLoader", "SchemaLoader", "InlineSchemaLoader"]
+__all__ = ["JsonFileSchemaLoader", "DefaultSchemaLoader", "SchemaLoader", "InlineSchemaLoader", "DynamicSchemaLoader", "TypesMap", "SchemaTypeIdentifier"]