deduplication version 2

Author: bazarnov

airbyte_cdk/sources/declarative/parsers/manifest_deduplicator.py

@@ -6,32 +6,30 @@
 import hashlib
 import json
 from collections import defaultdict
-from typing import Any, DefaultDict, Dict, Hashable, List, Optional, Tuple
+from typing import Any, DefaultDict, Dict, List, Optional, Tuple
 
 from airbyte_cdk.sources.declarative.parsers.custom_exceptions import ManifestDeduplicationException
 
 # Type definitions for better readability
 ManifestType = Dict[str, Any]
 DefinitionsType = Dict[str, Any]
-FieldDuplicatesType = DefaultDict[Tuple[str, Any], List[Tuple[List[str], Dict]]]
-ComponentDuplicatesType = DefaultDict[str, List[Tuple[List[str], Dict, Dict]]]
+DuplicatesType = DefaultDict[str, List[Tuple[List[str], Dict, Dict]]]
 
 # Configuration constants
 N_OCCURANCES = 2
 
 DEF_TAG = "definitions"
 SHARED_TAG = "shared"
 
-# SPECIFY COMPONENT TAGS FOR DEDUPLICATION
-COMPONENT_TAGS = [
+# SPECIFY TAGS FOR DEDUPLICATION
+TAGS = [
     "authenticator",
-]
-
-# SPECIFY FIELD TAGS FOR DEDUPLICATION
-FIELD_TAGS = [
     "url_base",
 ]
 
+# the placeholder for collected duplicates
+DUPLICATES: DuplicatesType = defaultdict(list, {})
+
 
 def deduplicate_definitions(resolved_manifest: ManifestType) -> ManifestType:
     """
@@ -41,124 +39,169 @@ def deduplicate_definitions(resolved_manifest: ManifestType) -> ManifestType:
         resolved_manifest: A dictionary representing a JSON structure to be analyzed.
 
     Returns:
-        A refactored JSON structure with common properties extracted.
+        A refactored JSON structure with common properties extracted to `definitions.shared`,
+        the duplicated properties replaced with references
     """
 
     try:
         _manifest = copy.deepcopy(resolved_manifest)
         definitions = _manifest.get(DEF_TAG, {})
-        field_duplicates, component_duplicates = _collect_all_duplicates(definitions)
-        _process_duplicates(definitions, field_duplicates, component_duplicates)
+
+        _collect_duplicates(definitions)
+        _handle_duplicates(definitions)
+
         return _manifest
     except ManifestDeduplicationException:
-        # we don't want to fix every single error which might occur,
-        # due to the varaety of possible manifest configurations,
         # if any arror occurs, we just return the original manifest.
         return resolved_manifest
 
 
-def _process_duplicates(
-    definitions: DefinitionsType,
-    field_duplicates: FieldDuplicatesType,
-    component_duplicates: ComponentDuplicatesType,
-) -> None:
+def _replace_duplicates_with_refs(definitions: ManifestType) -> None:
+    """
+    Process duplicate objects and replace them with references.
+
+    Args:
+        definitions: The definitions dictionary to modify
+    """
+    for _, occurrences in DUPLICATES.items():
+        # Skip non-duplicates
+        if len(occurrences) < N_OCCURANCES:
+            continue
+
+        # Take the value from the first occurrence, as they are the same
+        path, _, value = occurrences[0]
+        # take the component's name as the last part of it's path
+        key = path[-1]
+        # Create a meaningful reference key
+        ref_key = _create_reference_key(definitions, key)
+        # Add to definitions
+        _add_to_shared_definitions(definitions, ref_key, value)
+
+        # Replace all occurrences with references
+        for path, parent_obj, _ in occurrences:
+            if path:  # Make sure the path is valid
+                key = path[-1]
+                parent_obj[key] = _create_ref_object(ref_key)
+
+
+def _handle_duplicates(definitions: DefinitionsType) -> None:
     """
-    Process the duplicates and replace them with references.
+    Process the DUPLICATES and replace them with references.
 
     Args:
-        field_duplicates: Dictionary of duplicate primitive values
-        component_duplicates: Dictionary of duplicate objects
+        DUPLICATES: Dictionary of duplicate objects
     """
     # process duplicates only if there are any
-    if len(field_duplicates) > 0 or len(component_duplicates) > 0:
+    if len(DUPLICATES) > 0:
         if not SHARED_TAG in definitions:
             definitions[SHARED_TAG] = {}
 
         try:
-            _process_component_duplicates(definitions, component_duplicates)
-            _process_field_duplicates(definitions, field_duplicates)
+            _replace_duplicates_with_refs(definitions)
         except Exception as e:
             raise ManifestDeduplicationException(str(e))
 
 
-def _is_allowed_component(key: str) -> bool:
+def _is_allowed_tag(key: str) -> bool:
     """
-    Check if the key is an allowed component tag.
+    Check if the key is an allowed tag for deduplication.
 
     Args:
         key: The key to check
 
     Returns:
         True if the key is allowed, False otherwise
     """
-    return key in COMPONENT_TAGS
+    return key in TAGS
 
 
-def _is_allowed_field(key: str) -> bool:
+def _add_duplicate(
+    current_path: List[str],
+    obj: Dict,
+    value: Any,
+    key: Optional[str] = None,
+) -> None:
     """
-    Check if the key is an allowed field tag.
+    Adds a duplicate record of an observed object by computing a unique hash for the provided value.
 
-    Args:
-        key: The key to check
+    This function computes a hash for the given value (or a dictionary composed of the key and value if a key is provided)
+    and appends a tuple containing the current path, the original object, and the value to the global DUPLICATES
+    dictionary under the corresponding hash.
 
-    Returns:
-        True if the key is allowed, False otherwise
+    Parameters:
+        current_path (List[str]): The list of keys or indices representing the current location in the object hierarchy.
+        obj (Dict): The original dictionary object where the duplicate is observed.
+        value (Any): The value to be hashed and used for identifying duplicates.
+        key (Optional[str]): An optional key that, if provided, wraps the value in a dictionary before hashing.
+    """
+    # create hash for the duplicate observed
+    value_to_hash = value if key is None else {key: value}
+    obj_hash = _hash_object(value_to_hash)
+    if obj_hash:
+        DUPLICATES[obj_hash].append((current_path, obj, value))
+
+
+def _add_to_shared_definitions(
+    definitions: DefinitionsType,
+    key: str,
+    value: Any,
+) -> DefinitionsType:
     """
-    return key in FIELD_TAGS
+    Add a value to the shared definitions under the specified key.
+
+    Args:
+        definitions: The definitions dictionary to modify
+        key: The key to use
+        value: The value to add
+    """
+
+    if key not in definitions[SHARED_TAG]:
+        definitions[SHARED_TAG][key] = value
+
+    return definitions
 
 
-def _collect_all_duplicates(
-    node: ManifestType,
-) -> Tuple[FieldDuplicatesType, ComponentDuplicatesType]:
+def _collect_duplicates(node: ManifestType, path: Optional[List[str]] = None) -> None:
     """
     Traverse the JSON object and collect all potential duplicate values and objects.
 
     Args:
         node: The JSON object to analyze
 
     Returns:
-        A tuple of (field_duplicates, component_duplicates)
+        DUPLICATES: A dictionary of duplicate objects
     """
 
-    field_duplicates: FieldDuplicatesType = defaultdict(list)
-    component_duplicates: ComponentDuplicatesType = defaultdict(list)
-
-    def collect_duplicates(obj: Dict, path: Optional[List[str]] = None) -> None:
-        if not isinstance(obj, dict):
+    try:
+        if not isinstance(node, dict):
             return
 
         path = [] if path is None else path
+
         # Check if the object is empty
-        for key, value in obj.items():
+        for key, value in node.items():
             current_path = path + [key]
 
             if isinstance(value, dict):
                 # First process nested dictionaries
-                collect_duplicates(value, current_path)
-
+                _collect_duplicates(value, current_path)
                 # Process allowed-only component tags
-                if _is_allowed_component(key):
-                    obj_hash = _hash_object(value)
-                    if obj_hash:
-                        component_duplicates[obj_hash].append((current_path, obj, value))
+                if _is_allowed_tag(key):
+                    _add_duplicate(current_path, node, value)
+
+            # handle primitive types
+            elif isinstance(value, (str, int, float, bool)):
+                # Process allowed-only field tags
+                if _is_allowed_tag(key):
+                    _add_duplicate(current_path, node, value, key)
 
-            # handle list[dict] cases
+            # handle list cases
             elif isinstance(value, list):
                 for i, item in enumerate(value):
-                    collect_duplicates(item, current_path + [str(i)])
-
-            # Process allowed-only field tags
-            elif _is_allowed_field(key):
-                hashable_value = _make_hashable(value)
-                field_duplicates[(key, hashable_value)].append((current_path, obj))
-
-    try:
-        collect_duplicates(node)
+                    _collect_duplicates(item, current_path + [str(i)])
     except Exception as e:
         raise ManifestDeduplicationException(str(e))
 
-    return field_duplicates, component_duplicates
-
 
 def _hash_object(node: Dict) -> Optional[str]:
     """
@@ -173,26 +216,10 @@ def _hash_object(node: Dict) -> Optional[str]:
     if isinstance(node, Dict):
         # Sort keys to ensure consistent hash for same content
         return hashlib.md5(json.dumps(node, sort_keys=True).encode()).hexdigest()
-
     return None
 
 
-def _make_hashable(value: Any) -> Any:
-    """
-    Convert a value to a hashable representation.
-
-    Args:
-        value: The value to make hashable
-
-    Returns:
-        A hashable representation of the value
-    """
-    return json.dumps(value) if not isinstance(value, Hashable) else value
-
-
-def _create_reference_key(
-    definitions: DefinitionsType, key: str, value: Optional[Any] = None
-) -> str:
+def _create_reference_key(definitions: DefinitionsType, key: str) -> str:
     """
     Create a unique reference key and handle collisions.
 
@@ -206,9 +233,6 @@ def _create_reference_key(
 
     counter = 1
     while key in definitions[SHARED_TAG]:
-        # If the value is already in shared definitions with this key, no need to rename
-        if value is not None and _is_same_value(definitions[SHARED_TAG].get(key), value):
-            return key
         key = f"{key}_{counter}"
         counter += 1
     return key
@@ -225,97 +249,3 @@ def _create_ref_object(ref_key: str) -> Dict[str, str]:
         A reference object in the proper format
     """
     return {"$ref": f"#/{DEF_TAG}/{SHARED_TAG}/{ref_key}"}
-
-
-def _is_same_value(val1: Any, val2: Any) -> bool:
-    """
-    Check if two values are the same by comparing their JSON representation.
-
-    Args:
-        val1: First value
-        val2: Second value
-
-    Returns:
-        True if the values are the same, False otherwise
-    """
-    return json.dumps(val1, sort_keys=True) == json.dumps(val2, sort_keys=True)
-
-
-def _process_component_duplicates(
-    definitions: ManifestType,
-    component_duplicates: ComponentDuplicatesType,
-) -> None:
-    """
-    Process duplicate objects and replace them with references.
-
-    Args:
-        definitions: The definitions dictionary to modify
-        component_duplicates: Dictionary of duplicate objects
-    """
-    for obj_hash, occurrences in component_duplicates.items():
-        # Skip non-duplicates
-        if len(occurrences) < N_OCCURANCES:
-            continue
-
-        # Take the value from the first occurrence, as they are the same
-        path, _, value = occurrences[0]
-        # take the component's name as the last part of it's path
-        key = path[-1]
-        # Create a meaningful reference key
-        ref_key = _create_reference_key(definitions, key)
-        # Add to definitions
-        _add_to_shared_definitions(definitions, ref_key, value)
-
-        # Replace all occurrences with references
-        for path, parent_obj, _ in occurrences:
-            if path:  # Make sure the path is valid
-                key = path[-1]
-                parent_obj[key] = _create_ref_object(ref_key)
-
-
-def _add_to_shared_definitions(
-    definitions: DefinitionsType,
-    key: str,
-    value: Any,
-) -> DefinitionsType:
-    """
-    Add a value to the shared definitions under the specified key.
-
-    Args:
-        definitions: The definitions dictionary to modify
-        key: The key to use
-        value: The value to add
-    """
-
-    if key not in definitions[SHARED_TAG]:
-        definitions[SHARED_TAG][key] = value
-
-    return definitions
-
-
-def _process_field_duplicates(
-    definitions: ManifestType,
-    field_duplicates: FieldDuplicatesType,
-) -> None:
-    """
-    Process duplicate primitive values and replace them with references.
-
-    Args:
-        definitions: The definitions dictionary to modify
-        field_duplicates: Dictionary of duplicate primitive values
-    """
-
-    for (key, value), occurrences in field_duplicates.items():
-        # Skip non-duplicates
-        if len(occurrences) < N_OCCURANCES:
-            continue
-
-        ref_key = _create_reference_key(definitions, key, value)
-        # Add to definitions if not already there
-        _add_to_shared_definitions(definitions, ref_key, value)
-
-        # Replace all occurrences with references
-        for path, parent_obj in occurrences:
-            if path:  # Make sure the path is valid
-                key = path[-1]
-                parent_obj[key] = _create_ref_object(ref_key)