clean up validator implementation

Author: aaronsteers

airbyte_cdk/sources/utils/schema_helpers.py

@@ -13,7 +13,7 @@
 from jsonschema import validate
 from jsonschema.exceptions import ValidationError
 from pydantic.v1 import BaseModel, Field
-from referencing import Registry, Resource
+from referencing import Registry, Resource, Resolver
 from referencing.jsonschema import DRAFT7
 
 from airbyte_cdk.models import ConnectorSpecification, FailureType
@@ -65,21 +65,27 @@ def resolve_ref_links(obj: Any) -> Any:
         return obj
 
 
-def _expand_refs(schema: Any, ref_resolver: Optional[Registry] = None) -> None:
+def _expand_refs(schema: Any, ref_resolver: Optional[Resolver] = None) -> 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(schema, default_specification=DRAFT7)
-        ref_resolver = Registry().with_resource("", resource)
-    resolver = ref_resolver.resolver()
+        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")
-            definition = resolver.lookup(ref_url).contents
+            definition = ref_resolver.lookup(ref_url).contents
             _expand_refs(
                 definition, ref_resolver=ref_resolver
             )  # expand refs in definitions as well

airbyte_cdk/sources/utils/transform.py

@@ -6,8 +6,8 @@
 from enum import Flag, auto
 from typing import Any, Callable, Dict, Generator, Mapping, Optional, cast
 
-from jsonschema import Draft7Validator, RefResolver, ValidationError, Validator, validators
-from referencing import Registry, Resource
+from jsonschema import Draft7Validator, ValidationError, Validator, validators
+from referencing import Registry, Resource, Resolver
 from referencing.jsonschema import DRAFT7
 
 MAX_NESTING_DEPTH = 3
@@ -195,33 +195,15 @@ def normalizator(
             """
 
             def resolve(subschema: dict[str, Any]) -> dict[str, Any]:
-                if "$ref" in subschema:
-                    ref_url = subschema["$ref"]
-
-                    try:
-                        root_schema = validator_instance.schema
-                        resource = Resource.from_contents(root_schema, default_specification=DRAFT7)
-                        registry = Registry().with_resource("", resource)
-                        resolver = registry.resolver()
-                        resolved = resolver.lookup(ref_url).contents
-                        return cast(dict[str, Any], resolved)
-                    except Exception:
-                        try:
-                            if (
-                                hasattr(validator_instance, "resolver")
-                                and validator_instance.resolver is not None
-                            ):
-                                _, resolved = cast(
-                                    RefResolver, validator_instance.resolver
-                                ).resolve(ref_url)
-                                return cast(dict[str, Any], resolved)
-                        except Exception:
-                            # If both fail, we'll return original subschema, below.
-                            # If both fail, we'll return original subschema, below.
-                            pass
-
+                if "$ref" not in subschema:
+                    # Nothing to resolve
                     return subschema
-                return subschema
+
+                # Else, we need to resolve "$ref":
+                ref_url = subschema["$ref"]
+                resolver: Resolver = validator_instance.resolver
+                resolved_contents = resolver.lookup(ref_url).contents
+                return cast(dict[str, Any], resolved_contents)
 
             # Transform object and array values before running json schema type checking for each element.
             # Recursively normalize every value of the "instance" sub-object,