do not merge: vector db work (split)

airbyte_cdk/__init__.py

@@ -193,7 +193,6 @@
 from .utils.constants import ENV_REQUEST_CACHE_PATH
 from .utils.event_timing import create_timer
 from .utils.oneof_option_config import OneOfOptionConfig
-from .utils.spec_schema_transformations import resolve_refs
 from .utils.stream_status_utils import as_airbyte_message
 
 __all__ = [
@@ -340,7 +339,6 @@
     "ENV_REQUEST_CACHE_PATH",
     "create_timer",
     "OneOfOptionConfig",
-    "resolve_refs",
     "as_airbyte_message",
     # Types
     "Config",

airbyte_cdk/destinations/vector_db_based/config.py

@@ -7,8 +7,8 @@
 import dpath
 from pydantic.v1 import BaseModel, Field
 
+from airbyte_cdk.sources.utils.schema_helpers import expand_refs
 from airbyte_cdk.utils.oneof_option_config import OneOfOptionConfig
-from airbyte_cdk.utils.spec_schema_transformations import resolve_refs
 
 
 class SeparatorSplitterConfigModel(BaseModel):
@@ -293,6 +293,6 @@ def remove_discriminator(schema: Dict[str, Any]) -> None:
     def schema(cls, by_alias: bool = True, ref_template: str = "") -> Dict[str, Any]:
         """we're overriding the schema classmethod to enable some post-processing"""
         schema: Dict[str, Any] = super().schema()
-        schema = resolve_refs(schema)
+        expand_refs(schema)
         cls.remove_discriminator(schema)
         return schema

airbyte_cdk/manifest_server/Dockerfile

@@ -11,7 +11,7 @@ FROM python:3.12-slim-bookworm
 RUN apt-get update && \
     apt-get install -y git && \
     rm -rf /var/lib/apt/lists/* && \
-    pip install poetry==1.8.3
+    pip install poetry==2.0.1
 
 # Configure poetry to not create virtual environments and disable interactive mode
 ENV POETRY_NO_INTERACTION=1 \
@@ -42,4 +42,4 @@ USER airbyte:airbyte
 
 EXPOSE 8080
 
-CMD ["uvicorn", "airbyte_cdk.manifest_server.app:app", "--host", "0.0.0.0", "--port", "8080"]
+CMD ["uvicorn", "airbyte_cdk.manifest_server.app:app", "--host", "0.0.0.0", "--port", "8080"]

airbyte_cdk/sources/utils/schema_helpers.py

@@ -7,12 +7,16 @@
 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, Tuple, cast
 
 import jsonref
-from jsonschema import RefResolver, validate
+from jsonschema import validate
 from jsonschema.exceptions import ValidationError
 from pydantic.v1 import BaseModel, Field
+from referencing import Registry, Resource
+from referencing._core import Resolver  # used for type hints
+from referencing.jsonschema import DRAFT7
 
 from airbyte_cdk.models import ConnectorSpecification, FailureType
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
@@ -63,18 +67,30 @@ def resolve_ref_links(obj: Any) -> Any:
         return obj
 
 
-def _expand_refs(schema: Any, ref_resolver: Optional[RefResolver] = 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=schema,
+        default_specification=DRAFT7,
+    )
+    return cast(  # Mypy has a hard time detecting this return type.
+        "Registry",
+        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
     """
-    ref_resolver = ref_resolver or RefResolver.from_schema(schema)
-
     if isinstance(schema, MutableMapping):
         if "$ref" in schema:
             ref_url = schema.pop("$ref")
-            _, definition = ref_resolver.resolve(ref_url)
+            definition = ref_resolver.lookup(ref_url).contents
             _expand_refs(
                 definition, ref_resolver=ref_resolver
             )  # expand refs in definitions as well
@@ -90,10 +106,14 @@ def _expand_refs(schema: Any, ref_resolver: Optional[RefResolver] = None) -> Non
 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,10 +3,25 @@
 #
 
 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, validators
+from referencing import Registry, Resource
+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
 
-from jsonschema import Draft7Validator, RefResolver, ValidationError, Validator, validators
 
 MAX_NESTING_DEPTH = 3
 json_to_python_simple = {
@@ -191,30 +206,27 @@ def normalizator(
             validators parameter for detailed description.
             :
             """
+            # Very first step is to expand $refs in the schema itself:
+            expand_refs(schema)
+
+            # Now we can expand $refs in the property value:
+            if isinstance(property_value, dict):
+                expand_refs(property_value)
 
-            def resolve(subschema: dict[str, Any]) -> dict[str, Any]:
-                if "$ref" in subschema:
-                    _, resolved = cast(
-                        RefResolver,
-                        validator_instance.resolver,
-                    ).resolve(subschema["$ref"])
-                    return cast(dict[str, Any], resolved)
-                return subschema
+            # Now we can validate and normalize the values:
 
             # Transform object and array values before running json schema type checking for each element.
             # Recursively normalize every value of the "instance" sub-object,
             # if "instance" is an incorrect type - skip recursive normalization of "instance"
             if schema_key == "properties" and isinstance(instance, dict):
                 for k, subschema in property_value.items():
                     if k in instance:
-                        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)
                 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(