refactored KeyTransformation: added prefix and suffix

Author: darynaishchenko

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -2307,6 +2307,33 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  KeyTransformation:
+    title: Transformation to apply for extracted object keys by Dpath Flatten Fields
+    type: object
+    required:
+      - type
+    properties:
+      type:
+        type: string
+        enum: [ KeyTransformation ]
+      prefix:
+        title: Key Prefix
+        description: Prefix to add for object keys. If not provided original keys remain unchanged.
+        type: string
+        examples:
+          - flattened_
+      suffix:
+        title: Key Suffix
+        description: Suffix to add for object keys. If not provided original keys remain unchanged.
+        type: string
+        examples:
+          - _flattened
+      custom:
+        title: Custom Transformation for keys
+        description: Custom transformation. Can be used with {{ key }} as a original value for key name. If not provided original keys remain unchanged.
+        type: string
+        examples:
+          - flattened_{{ key }}
   DpathFlattenFields:
     title: Dpath Flatten Fields
     description: A transformation that flatten field values to the to top of the record.
@@ -2338,9 +2365,8 @@ definitions:
       key_transformation:
         title: Key transformation
         description: Transformation for object keys. If not provided, original key will be used.
-        type: string
-        examples:
-          - "flattened_{{ key }}"
+        type: object
+        "$ref": "#/definitions/KeyTransformation"
       $parameters:
         type: object
         additionalProperties: true

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -879,6 +879,33 @@ class FlattenFields(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class KeyTransformation(BaseModel):
+    prefix: Optional[Union[str, None]] = Field(
+        None,
+        description="Prefix to add for object keys. If not provided original keys remain unchanged.",
+        examples=[
+            "flattened_",
+        ],
+        title="Key Prefix",
+    )
+    suffix: Optional[Union[str, None]] = Field(
+        None,
+        description="Suffix to add for object keys. If not provided original keys remain unchanged.",
+        examples=[
+            "_flattened",
+        ],
+        title="Key Suffix",
+    )
+    custom: Optional[Union[str, None]] = Field(
+        None,
+        description="Custom transformation. Can be used with {{ key }} as a original value for key name. If not provided original keys remain unchanged.",
+        examples=[
+            "flattened_{{ key }}",
+        ],
+        title="Custom Transformation for keys",
+    )
+
+
 class DpathFlattenFields(BaseModel):
     type: Literal["DpathFlattenFields"]
     field_path: List[str] = Field(
@@ -897,12 +924,9 @@ class DpathFlattenFields(BaseModel):
         description="Whether to replace the origin record or not. Default is False.",
         title="Replace Origin Record",
     )
-    key_transformation: Optional[str] = Field(
+    key_transformation: Optional[Union[KeyTransformation, None]] = Field(
         None,
         description="Transformation for object keys. If not provided, original key will be used.",
-        examples=[
-            "flattened_{{ key }}",
-        ],
         title="Key transformation",
     )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -498,6 +498,7 @@
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.transformations.dpath_flatten_fields import (
     DpathFlattenFields,
+    KeyTransformation,
 )
 from airbyte_cdk.sources.declarative.transformations.flatten_fields import (
     FlattenFields,
@@ -790,14 +791,23 @@ def create_dpath_flatten_fields(
         self, model: DpathFlattenFieldsModel, config: Config, **kwargs: Any
     ) -> DpathFlattenFields:
         model_field_path: List[Union[InterpolatedString, str]] = [x for x in model.field_path]
+        key_transformation = (
+            KeyTransformation(
+                prefix=model.key_transformation.prefix,
+                suffix=model.key_transformation.suffix,
+                custom=model.key_transformation.custom,
+            )
+            if model.key_transformation is not None
+            else None
+        )
         return DpathFlattenFields(
             config=config,
             field_path=model_field_path,
             delete_origin_value=model.delete_origin_value
             if model.delete_origin_value is not None
             else False,
             replace_record=model.replace_record if model.replace_record is not None else False,
-            key_transformation=model.key_transformation,
+            key_transformation=key_transformation,
             parameters=model.parameters or {},
         )
 

airbyte_cdk/sources/declarative/transformations/dpath_flatten_fields.py

@@ -8,6 +8,13 @@
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 
 
+@dataclass(frozen=True)
+class KeyTransformation:
+    prefix: Union[InterpolatedString, str, None] = None
+    suffix: Union[InterpolatedString, str, None] = None
+    custom: Union[InterpolatedString, str, None] = None
+
+
 @dataclass
 class DpathFlattenFields(RecordTransformation):
     """
@@ -16,7 +23,7 @@ class DpathFlattenFields(RecordTransformation):
     field_path: List[Union[InterpolatedString, str]] path to the field to flatten.
     delete_origin_value: bool = False whether to delete origin field or keep it. Default is False.
     replace_record: bool = False whether to replace origin record or not. Default is False.
-    key_transformation: string = None how to transform extracted object keys
+    key_transformation: KeyTransformation = None how to transform extracted object keys
 
     """
 
@@ -25,7 +32,7 @@ class DpathFlattenFields(RecordTransformation):
     parameters: InitVar[Mapping[str, Any]]
     delete_origin_value: bool = False
     replace_record: bool = False
-    key_transformation: Union[InterpolatedString, str, None] = None
+    key_transformation: Union[KeyTransformation, None] = None
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
@@ -38,6 +45,30 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
                     self.field_path[path_index], parameters=self._parameters
                 )
 
+    def _apply_key_transformation(self, extracted: Mapping[str, Any]) -> Mapping[str, Any]:
+        if self.key_transformation:
+            if self.key_transformation.prefix:
+                prefix = InterpolatedString.create(
+                    self.key_transformation.prefix, parameters=self._parameters
+                ).eval(config=self.config)
+                extracted = {f"{prefix}{key}": value for key, value in extracted.items()}
+
+            if self.key_transformation.suffix:
+                suffix = InterpolatedString.create(
+                    self.key_transformation.suffix, parameters=self._parameters
+                ).eval(config=self.config)
+                extracted = {f"{key}{suffix}": value for key, value in extracted.items()}
+
+            if self.key_transformation.custom:
+                updated_extracted = {}
+                for key, value in extracted.items():
+                    updated_key = InterpolatedString.create(
+                        self.key_transformation.custom, parameters=self._parameters
+                    ).eval(key=key, config=self.config)
+                    updated_extracted[updated_key] = value
+                extracted = updated_extracted
+        return extracted
+
     def transform(
         self,
         record: Dict[str, Any],
@@ -53,14 +84,7 @@ def transform(
             extracted = dpath.get(record, path, default=[])
 
         if isinstance(extracted, dict):
-            if self.key_transformation:
-                updated_extracted = {}
-                for key, value in extracted.items():
-                    updated_key = InterpolatedString.create(
-                        self.key_transformation, parameters=self._parameters
-                    ).eval(key=key, config=self.config)
-                    updated_extracted[updated_key] = value
-                extracted = updated_extracted
+            extracted = self._apply_key_transformation(extracted)
 
             if self.replace_record and extracted:
                 dpath.delete(record, "**")

unit_tests/sources/declarative/transformations/test_dpath_flatten_fields.py

@@ -1,12 +1,18 @@
 import pytest
 
-from airbyte_cdk.sources.declarative.transformations.dpath_flatten_fields import DpathFlattenFields
+from airbyte_cdk.sources.declarative.transformations.dpath_flatten_fields import (
+    DpathFlattenFields,
+    KeyTransformation,
+)
 
 _ANY_VALUE = -1
 _DELETE_ORIGIN_VALUE = True
 _REPLACE_WITH_VALUE = True
 _DO_NOT_DELETE_ORIGIN_VALUE = False
 _DO_NOT_REPLACE_WITH_VALUE = False
+_NO_KEY_PREFIX = None
+_NO_KEY_SUFFIX = None
+_NO_CUSTOM_TRANSFORMATION = None
 _NO_KEY_TRANSFORMATIONS = None
 
 
@@ -147,29 +153,69 @@
             ["field2"],
             _DO_NOT_DELETE_ORIGIN_VALUE,
             _DO_NOT_REPLACE_WITH_VALUE,
-            "field2_{{ key }}",
+            (_NO_KEY_PREFIX, _NO_KEY_SUFFIX, "field2_{{ key }}"),
             {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}, "field2_field3": _ANY_VALUE},
-            id="flatten by dpath, not delete origin value, add keys transformation",
+            id="flatten by dpath, not delete origin value, add keys custom transformation",
         ),
         pytest.param(
             {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}},
             {},
             ["field2"],
             _DELETE_ORIGIN_VALUE,
             _DO_NOT_REPLACE_WITH_VALUE,
-            "field2_{{ key }}",
+            (_NO_KEY_PREFIX, _NO_KEY_SUFFIX, "field2_{{ key }}"),
             {"field1": _ANY_VALUE, "field2_field3": _ANY_VALUE},
-            id="flatten by dpath, delete origin value, add keys transformation",
+            id="flatten by dpath, delete origin value, add keys custom transformation",
         ),
         pytest.param(
             {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}},
             {},
             ["field2"],
             _DO_NOT_DELETE_ORIGIN_VALUE,
             _REPLACE_WITH_VALUE,
-            "field2_{{ key }}",
+            (_NO_KEY_PREFIX, _NO_KEY_SUFFIX, "field2_{{ key }}"),
             {"field2_field3": _ANY_VALUE},
-            id="flatten by dpath, not delete origin value, replace record, add keys transformation",
+            id="flatten by dpath, not delete origin value, replace record, add keys custom transformation",
+        ),
+        pytest.param(
+            {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}},
+            {},
+            ["field2"],
+            _DO_NOT_DELETE_ORIGIN_VALUE,
+            _REPLACE_WITH_VALUE,
+            ("prefix_", _NO_KEY_SUFFIX, _NO_CUSTOM_TRANSFORMATION),
+            {"prefix_field3": _ANY_VALUE},
+            id="flatten by dpath, not delete origin value, replace record, add keys prefix",
+        ),
+        pytest.param(
+            {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}},
+            {},
+            ["field2"],
+            _DO_NOT_DELETE_ORIGIN_VALUE,
+            _REPLACE_WITH_VALUE,
+            (_NO_KEY_PREFIX, "_suffix", _NO_CUSTOM_TRANSFORMATION),
+            {"field3_suffix": _ANY_VALUE},
+            id="flatten by dpath, not delete origin value, replace record, add keys suffix",
+        ),
+        pytest.param(
+            {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}},
+            {},
+            ["field2"],
+            _DO_NOT_DELETE_ORIGIN_VALUE,
+            _REPLACE_WITH_VALUE,
+            ("prefix_", "_suffix", _NO_CUSTOM_TRANSFORMATION),
+            {"prefix_field3_suffix": _ANY_VALUE},
+            id="flatten by dpath, not delete origin value, replace record, add keys prefix and suffix",
+        ),
+        pytest.param(
+            {"field1": _ANY_VALUE, "field2": {"field3": _ANY_VALUE}},
+            {},
+            ["field2"],
+            _DO_NOT_DELETE_ORIGIN_VALUE,
+            _REPLACE_WITH_VALUE,
+            ("prefix_", "_suffix", "{{ key|upper }}"),
+            {"PREFIX_FIELD3_SUFFIX": _ANY_VALUE},
+            id="flatten by dpath, not delete origin value, replace record, add keys prefix and suffix with custom to upper case",
         ),
     ],
 )
@@ -182,6 +228,9 @@ def test_dpath_flatten_lists(
     key_transformation,
     expected_record,
 ):
+    if key_transformation:
+        key_transformation = KeyTransformation(*key_transformation)
+
     flattener = DpathFlattenFields(
         field_path=field_path,
         parameters={},