airbyte-cdk 6.62.0.dev2__py3-none-any.whl → 6.62.0.dev4__py3-none-any.whl

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -593,32 +593,6 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
-  CustomIncrementalSync:
-    title: Custom Incremental Sync
-    description: Incremental component whose behavior is derived from a custom code implementation of the connector.
-    type: object
-    additionalProperties: true
-    required:
-      - type
-      - class_name
-      - cursor_field
-    properties:
-      type:
-        type: string
-        enum: [CustomIncrementalSync]
-      class_name:
-        title: Class Name
-        description: Fully-qualified name of the class that will be implementing the custom incremental sync. The format is `source_<name>.<package>.<class_name>`.
-        type: string
-        additionalProperties: true
-        examples:
-          - "source_railz.components.MyCustomIncrementalSync"
-      cursor_field:
-        description: The location of the value on a record that will be used as a bookmark during sync.
-        type: string
-      $parameters:
-        type: object
-        additionalProperties: true
   CustomPaginationStrategy:
     title: Custom Pagination Strategy
     description: Pagination strategy component whose behavior is derived from a custom code implementation of the connector.
@@ -1551,7 +1525,6 @@ definitions:
         anyOf:
           - "$ref": "#/definitions/DatetimeBasedCursor"
           - "$ref": "#/definitions/IncrementingCountCursor"
-          - "$ref": "#/definitions/CustomIncrementalSync"
       primary_key:
         title: Primary Key
         "$ref": "#/definitions/PrimaryKey"

airbyte_cdk/sources/declarative/incremental/__init__.py

@@ -2,36 +2,12 @@
 # Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 #
 
-from airbyte_cdk.legacy.sources.declarative.incremental.per_partition_cursor import (
-    CursorFactory,
-    PerPartitionCursor,
-)
 from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
     ConcurrentCursorFactory,
     ConcurrentPerPartitionCursor,
 )
-from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
-from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
-from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
-    GlobalSubstreamCursor,
-)
-from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
-    PerPartitionWithGlobalCursor,
-)
-from airbyte_cdk.sources.declarative.incremental.resumable_full_refresh_cursor import (
-    ChildPartitionResumableFullRefreshCursor,
-    ResumableFullRefreshCursor,
-)
 
 __all__ = [
-    "CursorFactory",
     "ConcurrentCursorFactory",
     "ConcurrentPerPartitionCursor",
-    "DatetimeBasedCursor",
-    "DeclarativeCursor",
-    "GlobalSubstreamCursor",
-    "PerPartitionCursor",
-    "PerPartitionWithGlobalCursor",
-    "ResumableFullRefreshCursor",
-    "ChildPartitionResumableFullRefreshCursor",
 ]

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -9,7 +9,7 @@ import time
 from collections import OrderedDict
 from copy import deepcopy
 from datetime import timedelta
-from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional
+from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional, TypeVar
 
 from airbyte_cdk.models import (
     AirbyteStateBlob,
@@ -19,10 +19,6 @@ from airbyte_cdk.models import (
     StreamDescriptor,
 )
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
-from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
-    Timer,
-    iterate_with_last_flag_and_state,
-)
 from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams.checkpoint.per_partition_key_serializer import (
@@ -38,6 +34,63 @@ from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
 logger = logging.getLogger("airbyte")
 
 
+T = TypeVar("T")
+
+
+def iterate_with_last_flag_and_state(
+    generator: Iterable[T], get_stream_state_func: Callable[[], Optional[Mapping[str, StreamState]]]
+) -> Iterable[tuple[T, bool, Any]]:
+    """
+    Iterates over the given generator, yielding tuples containing the element, a flag
+    indicating whether it's the last element in the generator, and the result of
+    `get_stream_state_func` applied to the element.
+
+    Args:
+        generator: The iterable to iterate over.
+        get_stream_state_func: A function that takes an element from the generator and
+            returns its state.
+
+    Returns:
+        An iterator that yields tuples of the form (element, is_last, state).
+    """
+
+    iterator = iter(generator)
+
+    try:
+        current = next(iterator)
+        state = get_stream_state_func()
+    except StopIteration:
+        return  # Return an empty iterator
+
+    for next_item in iterator:
+        yield current, False, state
+        current = next_item
+        state = get_stream_state_func()
+
+    yield current, True, state
+
+
+class Timer:
+    """
+    A simple timer class that measures elapsed time in seconds using a high-resolution performance counter.
+    """
+
+    def __init__(self) -> None:
+        self._start: Optional[int] = None
+
+    def start(self) -> None:
+        self._start = time.perf_counter_ns()
+
+    def finish(self) -> int:
+        if self._start:
+            return ((time.perf_counter_ns() - self._start) / 1e9).__ceil__()
+        else:
+            raise RuntimeError("Global substream cursor timer not started")
+
+    def is_running(self) -> bool:
+        return self._start is not None
+
+
 class ConcurrentCursorFactory:
     def __init__(self, create_function: Callable[..., ConcurrentCursor]):
         self._create_function = create_function

airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py

@@ -5,7 +5,6 @@ from typing import Any, Mapping
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
 from airbyte_cdk.sources.declarative.models import (
-    CustomIncrementalSync,
     DatetimeBasedCursor,
     SubstreamPartitionRouter,
 )
@@ -36,7 +35,7 @@ class LegacyToPerPartitionStateMigration(StateMigration):
     def __init__(
         self,
         partition_router: SubstreamPartitionRouter,
-        cursor: CustomIncrementalSync | DatetimeBasedCursor,
+        cursor: DatetimeBasedCursor,
         config: Mapping[str, Any],
         parameters: Mapping[str, Any],
     ):

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -174,24 +174,6 @@ class CustomErrorHandler(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class CustomIncrementalSync(BaseModel):
-    class Config:
-        extra = Extra.allow
-
-    type: Literal["CustomIncrementalSync"]
-    class_name: str = Field(
-        ...,
-        description="Fully-qualified name of the class that will be implementing the custom incremental sync. The format is `source_<name>.<package>.<class_name>`.",
-        examples=["source_railz.components.MyCustomIncrementalSync"],
-        title="Class Name",
-    )
-    cursor_field: str = Field(
-        ...,
-        description="The location of the value on a record that will be used as a bookmark during sync.",
-    )
-    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
-
-
 class CustomPaginationStrategy(BaseModel):
     class Config:
         extra = Extra.allow
@@ -2432,9 +2414,7 @@ class DeclarativeStream(BaseModel):
         description="Component used to coordinate how records are extracted across stream slices and request pages.",
         title="Retriever",
     )
-    incremental_sync: Optional[
-        Union[DatetimeBasedCursor, IncrementingCountCursor, CustomIncrementalSync]
-    ] = Field(
+    incremental_sync: Optional[Union[DatetimeBasedCursor, IncrementingCountCursor]] = Field(
         None,
         description="Component used to fetch data incrementally based on a time field in the data.",
         title="Incremental Sync",

airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py

@@ -19,11 +19,6 @@ DEFAULT_MODEL_TYPES: Mapping[str, str] = {
     "DatetimeBasedCursor.end_time_option": "RequestOption",
     "DatetimeBasedCursor.start_datetime": "MinMaxDatetime",
     "DatetimeBasedCursor.start_time_option": "RequestOption",
-    # CustomIncrementalSync
-    "CustomIncrementalSync.end_datetime": "MinMaxDatetime",
-    "CustomIncrementalSync.end_time_option": "RequestOption",
-    "CustomIncrementalSync.start_datetime": "MinMaxDatetime",
-    "CustomIncrementalSync.start_time_option": "RequestOption",
     # DeclarativeSource
     "DeclarativeSource.check": "CheckStream",
     "DeclarativeSource.spec": "Spec",

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -33,6 +33,10 @@ from requests import Response
 from airbyte_cdk.connector_builder.models import (
     LogMessage as ConnectorBuilderLogMessage,
 )
+from airbyte_cdk.legacy.sources.declarative.declarative_stream import DeclarativeStream
+from airbyte_cdk.legacy.sources.declarative.incremental import (
+    DatetimeBasedCursor,
+)
 from airbyte_cdk.models import (
     AirbyteStateBlob,
     AirbyteStateMessage,
@@ -75,7 +79,6 @@ from airbyte_cdk.sources.declarative.checks import (
 )
 from airbyte_cdk.sources.declarative.concurrency_level import ConcurrencyLevel
 from airbyte_cdk.sources.declarative.datetime.min_max_datetime import MinMaxDatetime
-from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.decoders import (
     Decoder,
     IterableDecoder,
@@ -105,10 +108,6 @@ from airbyte_cdk.sources.declarative.extractors.record_filter import (
 from airbyte_cdk.sources.declarative.incremental import (
     ConcurrentCursorFactory,
     ConcurrentPerPartitionCursor,
-    CursorFactory,
-    DatetimeBasedCursor,
-    GlobalSubstreamCursor,
-    PerPartitionWithGlobalCursor,
 )
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
@@ -200,9 +199,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     CustomErrorHandler as CustomErrorHandlerModel,
 )
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    CustomIncrementalSync as CustomIncrementalSyncModel,
-)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     CustomPaginationStrategy as CustomPaginationStrategyModel,
 )
@@ -701,7 +697,6 @@ class ModelToComponentFactory:
             CustomBackoffStrategyModel: self.create_custom_component,
             CustomDecoderModel: self.create_custom_component,
             CustomErrorHandlerModel: self.create_custom_component,
-            CustomIncrementalSyncModel: self.create_custom_component,
             CustomRecordExtractorModel: self.create_custom_component,
             CustomRecordFilterModel: self.create_custom_component,
             CustomRequesterModel: self.create_custom_component,
@@ -752,7 +747,7 @@ class ModelToComponentFactory:
             OAuthAuthenticatorModel: self.create_oauth_authenticator,
             OffsetIncrementModel: self.create_offset_increment,
             PageIncrementModel: self.create_page_increment,
-            ParentStreamConfigModel: self._create_message_repository_substream_wrapper,
+            ParentStreamConfigModel: self.create_parent_stream_config_with_substream_wrapper,
             PredicateValidatorModel: self.create_predicate_validator,
             PropertiesFromEndpointModel: self.create_properties_from_endpoint,
             PropertyChunkingModel: self.create_property_chunking,
@@ -1748,7 +1743,11 @@ class ModelToComponentFactory:
 
             if self._is_component(model_value):
                 model_args[model_field] = self._create_nested_component(
-                    model, model_field, model_value, config, **kwargs,
+                    model,
+                    model_field,
+                    model_value,
+                    config,
+                    **kwargs,
                 )
             elif isinstance(model_value, list):
                 vals = []
@@ -1760,7 +1759,15 @@ class ModelToComponentFactory:
                         if derived_type:
                             v["type"] = derived_type
                     if self._is_component(v):
-                        vals.append(self._create_nested_component(model, model_field, v, config, **kwargs,))
+                        vals.append(
+                            self._create_nested_component(
+                                model,
+                                model_field,
+                                v,
+                                config,
+                                **kwargs,
+                            )
+                        )
                     else:
                         vals.append(v)
                 model_args[model_field] = vals
@@ -1965,7 +1972,7 @@ class ModelToComponentFactory:
 
     def create_default_stream(
         self, model: DeclarativeStreamModel, config: Config, is_parent: bool = False, **kwargs: Any
-    ) -> Union[DeclarativeStream, AbstractStream]:
+    ) -> AbstractStream:
         primary_key = model.primary_key.__root__ if model.primary_key else None
 
         partition_router = self._build_stream_slicer_from_partition_router(
@@ -2062,6 +2069,7 @@ class ModelToComponentFactory:
             primary_key=primary_key,
             request_options_provider=request_options_provider,
             stream_slicer=stream_slicer,
+            partition_router=partition_router,
             stop_condition_cursor=concurrent_cursor
             if self._is_stop_condition_on_cursor(model)
             else None,
@@ -2524,7 +2532,9 @@ class ModelToComponentFactory:
             config=config,
             name=name,
             primary_key=None,
-            stream_slicer=self._build_stream_slicer_from_partition_router(model.retriever, config),
+            partition_router=self._build_stream_slicer_from_partition_router(
+                model.retriever, config
+            ),
             transformations=[],
             use_cache=True,
             log_formatter=(
@@ -2603,6 +2613,8 @@ class ModelToComponentFactory:
             fallback_parser=gzip_parser.inner_parser,
         )
 
+    # todo: This method should be removed once we deprecate the SimpleRetriever.cursor field and the various
+    #  state methods
     @staticmethod
     def create_incrementing_count_cursor(
         model: IncrementingCountCursorModel, config: Config, **kwargs: Any
@@ -3133,12 +3145,11 @@ class ModelToComponentFactory:
         transformations: List[RecordTransformation],
         file_uploader: Optional[DefaultFileUploader] = None,
         incremental_sync: Optional[
-            Union[
-                IncrementingCountCursorModel, DatetimeBasedCursorModel, CustomIncrementalSyncModel
-            ]
+            Union[IncrementingCountCursorModel, DatetimeBasedCursorModel]
         ] = None,
         use_cache: Optional[bool] = None,
         log_formatter: Optional[Callable[[Response], Any]] = None,
+        partition_router: Optional[PartitionRouter] = None,
         **kwargs: Any,
     ) -> SimpleRetriever:
         def _get_url(req: Requester) -> str:
@@ -3236,6 +3247,10 @@ class ModelToComponentFactory:
 
         if not request_options_provider:
             request_options_provider = DefaultRequestOptionsProvider(parameters={})
+        if isinstance(request_options_provider, DefaultRequestOptionsProvider) and isinstance(
+            partition_router, PartitionRouter
+        ):
+            request_options_provider = partition_router
 
         paginator = (
             self._create_component_from_model(
@@ -3684,7 +3699,7 @@ class ModelToComponentFactory:
         if model.parent_stream_configs:
             parent_stream_configs.extend(
                 [
-                    self._create_message_repository_substream_wrapper(
+                    self.create_parent_stream_config_with_substream_wrapper(
                         model=parent_stream_config, config=config, **kwargs
                     )
                     for parent_stream_config in model.parent_stream_configs
@@ -3697,13 +3712,11 @@ class ModelToComponentFactory:
             config=config,
         )
 
-    def _create_message_repository_substream_wrapper(
+    def create_parent_stream_config_with_substream_wrapper(
         self, model: ParentStreamConfigModel, config: Config, *, stream_name: str, **kwargs: Any
     ) -> Any:
         # getting the parent state
-        child_state = self._connector_state_manager.get_stream_state(
-            stream_name, None
-        )
+        child_state = self._connector_state_manager.get_stream_state(stream_name, None)
 
         # This flag will be used exclusively for StateDelegatingStream when a parent stream is created
         has_parent_state = bool(
@@ -3771,7 +3784,6 @@ class ModelToComponentFactory:
                         incremental_sync_model: Union[
                             DatetimeBasedCursorModel,
                             IncrementingCountCursorModel,
-                            CustomIncrementalSyncModel,
                         ] = (
                             model.stream.incremental_sync  # type: ignore  # if we are there, it is because there is incremental_dependency and therefore there is an incremental_sync on the parent stream
                             if isinstance(model.stream, DeclarativeStreamModel)

airbyte_cdk/sources/declarative/requesters/http_job_repository.py

@@ -1,4 +1,5 @@
-# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+
 import logging
 import uuid
 from dataclasses import dataclass, field
@@ -8,9 +9,8 @@ from typing import Any, Dict, Iterable, Mapping, Optional
 import requests
 from requests import Response
 
-from airbyte_cdk import AirbyteMessage
 from airbyte_cdk.logger import lazy_log
-from airbyte_cdk.models import FailureType, Type
+from airbyte_cdk.models import AirbyteMessage, FailureType, Type
 from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
 from airbyte_cdk.sources.declarative.async_job.repository import AsyncJobRepository
 from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -23,10 +23,10 @@ from typing import (
 import requests
 from typing_extensions import deprecated
 
+from airbyte_cdk.legacy.sources.declarative.incremental import ResumableFullRefreshCursor
+from airbyte_cdk.legacy.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
 from airbyte_cdk.models import AirbyteMessage
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
-from airbyte_cdk.sources.declarative.incremental import ResumableFullRefreshCursor
-from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
     SinglePartitionRouter,

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -14,10 +14,21 @@ from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import Stre
 from airbyte_cdk.sources.types import Record, StreamSlice
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
+
 # For Connector Builder test read operations, we track the total number of records
-# read for the stream at the global level so that we can stop reading early if we
-# exceed the record limit
-total_record_counter = 0
+# read for the stream so that we can stop reading early if we exceed the record limit.
+class RecordCounter:
+    def __init__(self) -> None:
+        self.total_record_counter = 0
+
+    def increment(self) -> None:
+        self.total_record_counter += 1
+
+    def reset(self) -> None:
+        self.total_record_counter = 0
+
+    def get_total_records(self) -> int:
+        return self.total_record_counter
 
 
 class SchemaLoaderCachingDecorator(SchemaLoader):
@@ -51,6 +62,7 @@ class DeclarativePartitionFactory:
         self._retriever = retriever
         self._message_repository = message_repository
         self._max_records_limit = max_records_limit
+        self._record_counter = RecordCounter()
 
     def create(self, stream_slice: StreamSlice) -> Partition:
         return DeclarativePartition(
@@ -60,6 +72,7 @@ class DeclarativePartitionFactory:
             message_repository=self._message_repository,
             max_records_limit=self._max_records_limit,
             stream_slice=stream_slice,
+            record_counter=self._record_counter,
         )
 
 
@@ -72,6 +85,7 @@ class DeclarativePartition(Partition):
         message_repository: MessageRepository,
         max_records_limit: Optional[int],
         stream_slice: StreamSlice,
+        record_counter: RecordCounter,
     ):
         self._stream_name = stream_name
         self._schema_loader = schema_loader
@@ -80,17 +94,17 @@ class DeclarativePartition(Partition):
         self._max_records_limit = max_records_limit
         self._stream_slice = stream_slice
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
+        self._record_counter = record_counter
 
     def read(self) -> Iterable[Record]:
         if self._max_records_limit is not None:
-            global total_record_counter
-            if total_record_counter >= self._max_records_limit:
+            if self._record_counter.get_total_records() >= self._max_records_limit:
                 return
         for stream_data in self._retriever.read_records(
             self._schema_loader.get_json_schema(), self._stream_slice
         ):
             if self._max_records_limit is not None:
-                if total_record_counter >= self._max_records_limit:
+                if self._record_counter.get_total_records() >= self._max_records_limit:
                     break
 
             if isinstance(stream_data, Mapping):
@@ -108,7 +122,7 @@ class DeclarativePartition(Partition):
                 self._message_repository.emit_message(stream_data)
 
             if self._max_records_limit is not None:
-                total_record_counter += 1
+                self._record_counter.increment()
 
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         return self._stream_slice

airbyte_cdk/sources/declarative/transformations/keys_replace_transformation.py

@@ -5,7 +5,7 @@
 from dataclasses import InitVar, dataclass
 from typing import Any, Dict, Mapping, Optional
 
-from airbyte_cdk import InterpolatedString
+from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 

airbyte_cdk/sources/declarative/yaml_declarative_source.py

@@ -14,7 +14,7 @@ from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
 from airbyte_cdk.sources.types import ConnectionDefinition
 
 
-class YamlDeclarativeSource(ConcurrentDeclarativeSource[List[AirbyteStateMessage]]):
+class YamlDeclarativeSource(ConcurrentDeclarativeSource):
     """Declarative source defined by a yaml file"""
 
     def __init__(

airbyte_cdk/sources/streams/concurrent/abstract_stream.py

@@ -14,10 +14,6 @@ from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 
 
-@deprecated(
-    "This class is experimental. Use at your own risk.",
-    category=ExperimentalClassWarning,
-)
 class AbstractStream(ABC):
     """
     AbstractStream is an experimental interface for streams developed as part of the Concurrent CDK.

airbyte_cdk/sources/utils/schema_helpers.py

@@ -7,12 +7,16 @@ import importlib
 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: