fixes, simplification

Author: aaronsteers

airbyte_cdk/sources/utils/schema_helpers.py

@@ -7,7 +7,18 @@
 import json
 import os
 import pkgutil
-from typing import Any, ClassVar, Dict, List, Mapping, MutableMapping, Optional, Tuple
+from copy import deepcopy
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+)
 
 import jsonref
 from jsonschema import validate
@@ -66,23 +77,23 @@ def resolve_ref_links(obj: Any) -> Any:
         return obj
 
 
-def _expand_refs(schema: Any, ref_resolver: Optional[Resolver] = None) -> None:
+def get_ref_resolver_registry(schema: dict[str, Any]) -> Registry:
+    """Get a reference resolver registry for the given schema."""
+    resource: Resource = Resource.from_contents(
+        contents=deepcopy(schema),
+        default_specification=DRAFT7,
+    )
+    return Registry().with_resource(
+        uri="",
+        resource=resource,
+    )
+
+
+def _expand_refs(schema: Any, ref_resolver: Resolver) -> None:
     """Internal function to iterate over schema and replace all occurrences of $ref with their definitions. Recursive.
 
     :param schema: schema that will be patched
-    :param ref_resolver: resolver to get definition from $ref, if None pass it will be instantiated
     """
-    if ref_resolver is None:
-        resource = Resource.from_contents(
-            contents=schema,
-            default_specification=DRAFT7,
-        )
-        resolver_registry = Registry().with_resource(
-            uri="",
-            resource=resource,
-        )
-        ref_resolver = resolver_registry.resolver()
-
     if isinstance(schema, MutableMapping):
         if "$ref" in schema:
             ref_url = schema.pop("$ref")
@@ -102,10 +113,14 @@ def _expand_refs(schema: Any, ref_resolver: Optional[Resolver] = None) -> None:
 def expand_refs(schema: Any) -> None:
     """Iterate over schema and replace all occurrences of $ref with their definitions.
 
+    If a "definitions" section is present at the root of the schema, it will be removed
+    after $ref resolution is complete.
+
     :param schema: schema that will be patched
     """
-    _expand_refs(schema)
-    schema.pop("definitions", None)  # remove definitions created by $ref
+    ref_resolver = get_ref_resolver_registry(schema).resolver()
+    _expand_refs(schema, ref_resolver)
+    schema.pop("definitions", None)
 
 
 def rename_key(schema: Any, old_key: str, new_key: str) -> None:

airbyte_cdk/sources/utils/transform.py

@@ -3,14 +3,27 @@
 #
 
 import logging
+from copy import deepcopy
 from enum import Flag, auto
-from typing import Any, Callable, Dict, Generator, Mapping, Optional, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, Mapping, Optional, cast
 
-from jsonschema import Draft7Validator, ValidationError, Validator, validators
+from jsonschema import Draft7Validator, ValidationError, validators
+from jsonschema._typing import SchemaKeywordValidator
 from referencing import Registry, Resource
-from referencing._core import Resolver  # used for type hints
+from referencing._core import Resolver
+from referencing.exceptions import Unresolvable
 from referencing.jsonschema import DRAFT7
 
+from airbyte_cdk.sources.utils.schema_helpers import expand_refs
+
+from .schema_helpers import get_ref_resolver_registry
+
+try:
+    from jsonschema.validators import Validator
+except:
+    from jsonschema import Validator
+
+
 MAX_NESTING_DEPTH = 3
 json_to_python_simple = {
     "string": str,
@@ -194,20 +207,22 @@ def normalizator(
             validators parameter for detailed description.
             :
             """
+            # Very first step is to expand references in the schema itself
+            expand_refs(schema)
+            if isinstance(property_value, dict):
+                expand_refs(property_value)
+            # resolver_registry: Registry = get_ref_resolver_registry(schema)
+            # ref_resolver: Resolver = resolver_registry.resolver()
 
             def resolve(subschema: dict[str, Any]) -> dict[str, Any]:
-                if "$ref" in subschema:
-                    ref_url = subschema["$ref"]
-                    try:
-                        if hasattr(validator_instance.resolver, 'lookup'):
-                            resolved = validator_instance.resolver.lookup(ref_url).contents
-                            return cast(dict[str, Any], resolved)
-                        elif hasattr(validator_instance.resolver, 'resolve'):
-                            _, resolved = validator_instance.resolver.resolve(ref_url)
-                            return cast(dict[str, Any], resolved)
-                    except Exception:
-                        pass
-                    return subschema
+                # if "$ref" in subschema:
+                #     try:
+                #         resolved = ref_resolver.lookup(subschema["$ref"]).contents
+                #     except Unresolvable as e:
+                #         raise ValidationError(
+                #             f"Failed to resolve $ref '{subschema['$ref']}' from {ref_resolver!r} and schema {schema!r}: {e}"
+                #         ) from e
+                #     return cast(dict[str, Any], resolved)
                 return subschema
 
             # Transform object and array values before running json schema type checking for each element.
@@ -216,14 +231,14 @@ def resolve(subschema: dict[str, Any]) -> dict[str, Any]:
             if schema_key == "properties" and isinstance(instance, dict):
                 for k, subschema in property_value.items():
                     if k in instance:
-                        subschema = resolve(subschema)
+                        # subschema = resolve(subschema)
                         instance[k] = self.__normalize(instance[k], subschema)
             # Recursively normalize every item of the "instance" sub-array,
             # if "instance" is an incorrect type - skip recursive normalization of "instance"
             elif schema_key == "items" and isinstance(instance, list):
-                subschema = resolve(property_value)
+                # subschema = resolve(property_value)
                 for index, item in enumerate(instance):
-                    instance[index] = self.__normalize(item, subschema)
+                    instance[index] = self.__normalize(item, property_value)
 
             # Running native jsonschema traverse algorithm after field normalization is done.
             yield from original_validator(

unit_tests/sources/utils/test_transform.py

@@ -20,7 +20,7 @@
         "too_many_types": {"type": ["boolean", "null", "string"]},
         "def": {
             "type": "object",
-            "properties": {"dd": {"$ref": "#/definitions/my_type"}},
+            "properties": {"dd": {"$ref": "#/definitions/my_type"}},  # << Broken (missing?) on purpose?
         },
         "array": {"type": "array", "items": {"$ref": "#/definitions/str_type"}},
         "nested": {"$ref": "#/definitions/nested_type"},
@@ -31,6 +31,7 @@
     },
     "definitions": {
         "str_type": {"type": "string"},
+        "my_type": {"type": "string"},  # << Fixed
         "nested_type": {"type": "object", "properties": {"a": {"type": "string"}}},
     },
 }