airbyte_cdk.sources.declarative.schema

33@dataclass
34class JsonFileSchemaLoader(ResourceSchemaLoader, SchemaLoader):
35    """
36    Loads the schema from a json file
37
38    Attributes:
39        file_path (Union[InterpolatedString, str]): The path to the json file describing the schema
40        name (str): The stream's name
41        config (Config): The user-provided configuration as specified by the source's spec
42        parameters (Mapping[str, Any]): Additional arguments to pass to the string interpolation if needed
43    """
44
45    config: Config
46    parameters: InitVar[Mapping[str, Any]]
47    file_path: Union[InterpolatedString, str] = field(default="")
48
49    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
50        if not self.file_path:
51            self.file_path = _default_file_path()
52        self.file_path = InterpolatedString.create(self.file_path, parameters=parameters)
53
54    def get_json_schema(self) -> Mapping[str, Any]:
55        # todo: It is worth revisiting if we can replace file_path with just file_name if every schema is in the /schemas directory
56        # this would require that we find a creative solution to store or retrieve source_name in here since the files are mounted there
57        json_schema_path = self._get_json_filepath()
58        resource, schema_path = self.extract_resource_and_schema_path(json_schema_path)
59        raw_json_file = pkgutil.get_data(resource, schema_path)
60
61        if not raw_json_file:
62            raise IOError(f"Cannot find file {json_schema_path}")
63        try:
64            raw_schema = json.loads(raw_json_file)
65        except ValueError as err:
66            raise RuntimeError(f"Invalid JSON file format for file {json_schema_path}") from err
67        self.package_name = resource
68        return self._resolve_schema_references(raw_schema)
69
70    def _get_json_filepath(self) -> Any:
71        return self.file_path.eval(self.config)  # type: ignore # file_path is always cast to an interpolated string
72
73    @staticmethod
74    def extract_resource_and_schema_path(json_schema_path: str) -> Tuple[str, str]:
75        """
76        When the connector is running on a docker container, package_data is accessible from the resource (source_<name>), so we extract
77        the resource from the first part of the schema path and the remaining path is used to find the schema file. This is a slight
78        hack to identify the source name while we are in the airbyte_cdk module.
79        :param json_schema_path: The path to the schema JSON file
80        :return: Tuple of the resource name and the path to the schema file
81        """
82        split_path = json_schema_path.split("/")
83
84        if split_path[0] == "" or split_path[0] == ".":
85            split_path = split_path[1:]
86
87        if len(split_path) == 0:
88            return "", ""
89
90        if len(split_path) == 1:
91            return "", split_path[0]
92
93        return split_path[0], "/".join(split_path[1:])

15@dataclass
16class DefaultSchemaLoader(SchemaLoader):
17    """
18    Loads a schema from the default location or returns an empty schema for streams that have not defined their schema file yet.
19
20    Attributes:
21        config (Config): The user-provided configuration as specified by the source's spec
22        parameters (Mapping[str, Any]): Additional arguments to pass to the string interpolation if needed
23    """
24
25    config: Config
26    parameters: InitVar[Mapping[str, Any]]
27
28    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
29        self._parameters = parameters
30        self.default_loader = JsonFileSchemaLoader(parameters=parameters, config=self.config)
31
32    def get_json_schema(self) -> Mapping[str, Any]:
33        """
34        Attempts to retrieve a schema from the default filepath location or returns the empty schema if a schema cannot be found.
35
36        :return: The empty schema
37        """
38
39        try:
40            return self.default_loader.get_json_schema()
41        except (OSError, ValueError):
42            # A slight hack since we don't directly have the stream name. However, when building the default filepath we assume the
43            # runtime options stores stream name 'name' so we'll do the same here
44            stream_name = self._parameters.get("name", "")
45            logging.info(
46                f"Could not find schema for stream {stream_name}, defaulting to the empty schema"
47            )
48            return {}

11@dataclass
12class SchemaLoader:
13    """Describes a stream's schema"""
14
15    @abstractmethod
16    def get_json_schema(self) -> Mapping[str, Any]:
17        """Returns a mapping describing the stream's schema"""
18        pass

12@dataclass
13class InlineSchemaLoader(SchemaLoader):
14    """Describes a stream's schema"""
15
16    schema: Dict[str, Any]
17    parameters: InitVar[Mapping[str, Any]]
18
19    def get_json_schema(self) -> Mapping[str, Any]:
20        return self.schema

120@deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
121@dataclass
122class DynamicSchemaLoader(SchemaLoader):
123    """
124    Dynamically loads a JSON Schema by extracting data from retrieved records.
125
126    With an AsyncRetriever, each schema load creates one async job (retried on failure) and stops at
127    the first record, so the job's delete endpoint is not called. During a read the schema loads on
128    the first partition read and its job takes a slot from max_concurrent_async_job_count.
129    """
130
131    retriever: Retriever
132    config: Config
133    parameters: InitVar[Mapping[str, Any]]
134    schema_type_identifier: SchemaTypeIdentifier
135    schema_transformations: List[RecordTransformation] = field(default_factory=lambda: [])
136    schema_filter: Optional[RecordFilter] = None
137
138    def get_json_schema(self) -> Mapping[str, Any]:
139        """
140        Constructs a JSON Schema based on retrieved data.
141        """
142        properties = {}
143        retrieved_record = next(
144            (
145                record
146                for record in self.retriever.read_records({})
147                # AsyncRetriever emits a slice log message before its records
148                if not (isinstance(record, AirbyteMessage) and record.type == Type.LOG)
149            ),
150            None,
151        )
152
153        raw_schema = (
154            self._extract_data(
155                retrieved_record,  # type: ignore[arg-type] # Expected that retrieved_record will be only Mapping[str, Any]
156                self.schema_type_identifier.schema_pointer,
157            )
158            if retrieved_record
159            else []
160        )
161
162        for property_definition in raw_schema:
163            key = self._get_key(property_definition, self.schema_type_identifier.key_pointer)
164            value = self._get_type(
165                property_definition,
166                self.schema_type_identifier.type_pointer,
167            )
168            properties[key] = value
169
170        filtered_transformed_properties = self._transform(self._filter(properties))
171
172        return {
173            "$schema": "https://json-schema.org/draft-07/schema#",
174            "type": "object",
175            "additionalProperties": True,
176            "properties": filtered_transformed_properties,
177        }
178
179    def _transform(
180        self,
181        properties: Mapping[str, Any],
182    ) -> Mapping[str, Any]:
183        for transformation in self.schema_transformations:
184            transformation.transform(
185                properties,  # type: ignore  # properties has type Mapping[str, Any], but Dict[str, Any] expected
186                config=self.config,
187            )
188        return properties
189
190    def _filter(
191        self,
192        properties: Mapping[str, Any],
193    ) -> Mapping[str, Any]:
194        if not self.schema_filter:
195            return properties
196
197        filtered_properties: MutableMapping[str, Any] = {}
198        for item in self.schema_filter.filter_records(
199            ({k: v} for k, v in properties.items()),
200            {},
201        ):
202            filtered_properties.update(item)
203        return filtered_properties
204
205    def _get_key(
206        self,
207        raw_schema: MutableMapping[str, Any],
208        field_key_path: List[Union[InterpolatedString, str]],
209    ) -> str:
210        """
211        Extracts the key field from the schema using the specified path.
212        """
213        field_key = self._extract_data(raw_schema, field_key_path)
214        if not isinstance(field_key, str):
215            raise ValueError(f"Expected key to be a string. Got {field_key}")
216        return field_key
217
218    def _get_type(
219        self,
220        raw_schema: MutableMapping[str, Any],
221        field_type_path: Optional[List[Union[InterpolatedString, str]]],
222    ) -> Union[Mapping[str, Any], List[Mapping[str, Any]]]:
223        """
224        Determines the JSON Schema type for a field, supporting nullable and combined types.
225        """
226        raw_field_type = (
227            self._extract_data(raw_schema, field_type_path, default="string")
228            if field_type_path
229            else "string"
230        )
231        mapped_field_type = self._replace_type_if_not_valid(raw_field_type, raw_schema)
232        if (
233            isinstance(mapped_field_type, list)
234            and len(mapped_field_type) == 2
235            and all(isinstance(item, str) for item in mapped_field_type)
236        ):
237            first_type = self._get_airbyte_type(mapped_field_type[0])
238            second_type = self._get_airbyte_type(mapped_field_type[1])
239            return {"oneOf": [first_type, second_type]}
240
241        elif isinstance(mapped_field_type, str):
242            return self._get_airbyte_type(mapped_field_type)
243
244        elif isinstance(mapped_field_type, ComplexFieldType):
245            return self._resolve_complex_type(mapped_field_type)
246
247        else:
248            raise ValueError(
249                f"Invalid data type. Available string or two items list of string. Got {mapped_field_type}."
250            )
251
252    def _resolve_complex_type(self, complex_type: ComplexFieldType) -> Mapping[str, Any]:
253        if not complex_type.items:
254            return self._get_airbyte_type(complex_type.field_type)
255
256        field_type = self._get_airbyte_type(complex_type.field_type)
257
258        field_type["items"] = (
259            self._get_airbyte_type(complex_type.items)
260            if isinstance(complex_type.items, str)
261            else self._resolve_complex_type(complex_type.items)
262        )
263
264        return field_type
265
266    def _replace_type_if_not_valid(
267        self,
268        field_type: Union[List[str], str],
269        raw_schema: MutableMapping[str, Any],
270    ) -> Union[List[str], str, ComplexFieldType]:
271        """
272        Replaces a field type if it matches a type mapping in `types_map`.
273        """
274        if self.schema_type_identifier.types_mapping:
275            for types_map in self.schema_type_identifier.types_mapping:
276                # conditional is optional param, setting to true if not provided
277                condition = InterpolatedBoolean(
278                    condition=types_map.condition if types_map.condition is not None else "True",
279                    parameters={},
280                ).eval(config=self.config, raw_schema=raw_schema)
281
282                if field_type == types_map.current_type and condition:
283                    return types_map.target_type
284        return field_type
285
286    @staticmethod
287    def _get_airbyte_type(field_type: str) -> MutableMapping[str, Any]:
288        """
289        Maps a field type to its corresponding Airbyte type definition.
290        """
291        if field_type not in AIRBYTE_DATA_TYPES:
292            raise ValueError(f"Invalid Airbyte data type: {field_type}")
293
294        return deepcopy(AIRBYTE_DATA_TYPES[field_type])
295
296    def _extract_data(
297        self,
298        body: Mapping[str, Any],
299        extraction_path: Optional[List[Union[InterpolatedString, str]]] = None,
300        default: Any = None,
301    ) -> Any:
302        """
303        Extracts data from the body based on the provided extraction path.
304        """
305
306        if not extraction_path:
307            return body
308
309        path = [
310            node.eval(self.config) if not isinstance(node, str) else node
311            for node in extraction_path
312        ]
313
314        return dpath.get(body, path, default=default)  # type: ignore # extracted will be a MutableMapping, given input data structure

51@deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
52@dataclass(frozen=True)
53class ComplexFieldType:
54    """
55    Identifies complex field type
56    """
57
58    field_type: str
59    items: Optional[Union[str, "ComplexFieldType"]] = None
60
61    def __post_init__(self) -> None:
62        """
63        Enforces that `items` is only used when `field_type` is a array
64        """
65        # `items_type` is valid only for array target types
66        if self.items and self.field_type != "array":
67            raise ValueError("'items' can only be used when 'field_type' is an array.")

70@deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
71@dataclass(frozen=True)
72class TypesMap:
73    """
74    Represents a mapping between a current type and its corresponding target type.
75    """
76
77    target_type: Union[List[str], str, ComplexFieldType]
78    current_type: Union[List[str], str]
79    condition: Optional[str]

 82@deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
 83@dataclass
 84class SchemaTypeIdentifier:
 85    """
 86    Identifies schema details for dynamic schema extraction and processing.
 87    """
 88
 89    key_pointer: List[Union[InterpolatedString, str]]
 90    parameters: InitVar[Mapping[str, Any]]
 91    type_pointer: Optional[List[Union[InterpolatedString, str]]] = None
 92    types_mapping: Optional[List[TypesMap]] = None
 93    schema_pointer: Optional[List[Union[InterpolatedString, str]]] = None
 94
 95    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 96        self.schema_pointer = (
 97            self._update_pointer(self.schema_pointer, parameters) if self.schema_pointer else []
 98        )  # type: ignore[assignment]  # This is reqired field in model
 99        self.key_pointer = self._update_pointer(self.key_pointer, parameters)  # type: ignore[assignment]  # This is reqired field in model
100        self.type_pointer = (
101            self._update_pointer(self.type_pointer, parameters) if self.type_pointer else None
102        )
103
104    @staticmethod
105    def _update_pointer(
106        pointer: Optional[List[Union[InterpolatedString, str]]], parameters: Mapping[str, Any]
107    ) -> Optional[List[Union[InterpolatedString, str]]]:
108        return (
109            [
110                InterpolatedString.create(path, parameters=parameters)
111                if isinstance(path, str)
112                else path
113                for path in pointer
114            ]
115            if pointer
116            else None
117        )