-
Notifications
You must be signed in to change notification settings - Fork 4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
CDK: Add base pydantic model for connector config and schemas (#8485)
* add base spec model * fix usage of state_checkpoint_interval in case it is dynamic * add schema base models, fix spelling, signatures and polishing Co-authored-by: Eugene Kulak <kulak.eugene@gmail.com>
- Loading branch information
1 parent
161a1ee
commit aa67604
Showing
12 changed files
with
336 additions
and
54 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,5 +1,8 @@ | ||
# Changelog | ||
|
||
## 0.1.42 | ||
Add base pydantic model for connector config and schemas. | ||
|
||
## 0.1.41 | ||
Fix build error | ||
|
||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,65 @@ | ||
# | ||
# Copyright (c) 2021 Airbyte, Inc., all rights reserved. | ||
# | ||
|
||
from typing import Any, Dict, List, MutableMapping, Optional | ||
|
||
from jsonschema import RefResolver | ||
from pydantic import BaseModel | ||
|
||
|
||
class BaseConfig(BaseModel): | ||
"""Base class for connector spec, adds the following behaviour: | ||
- resolve $ref and replace it with definition | ||
- replace all occurrences of anyOf with oneOf | ||
- drop description | ||
""" | ||
|
||
@classmethod | ||
def _rename_key(cls, schema: Any, old_key: str, new_key: str) -> None: | ||
"""Iterate over nested dictionary and replace one key with another. Used to replace anyOf with oneOf. Recursive." | ||
:param schema: schema that will be patched | ||
:param old_key: name of the key to replace | ||
:param new_key: new name of the key | ||
""" | ||
if not isinstance(schema, MutableMapping): | ||
return | ||
|
||
for key, value in schema.items(): | ||
cls._rename_key(value, old_key, new_key) | ||
if old_key in schema: | ||
schema[new_key] = schema.pop(old_key) | ||
|
||
@classmethod | ||
def _expand_refs(cls, schema: Any, ref_resolver: Optional[RefResolver] = None) -> None: | ||
"""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) | ||
cls._expand_refs(definition, ref_resolver=ref_resolver) # expand refs in definitions as well | ||
schema.update(definition) | ||
else: | ||
for key, value in schema.items(): | ||
cls._expand_refs(value, ref_resolver=ref_resolver) | ||
elif isinstance(schema, List): | ||
for value in schema: | ||
cls._expand_refs(value, ref_resolver=ref_resolver) | ||
|
||
@classmethod | ||
def schema(cls, **kwargs) -> Dict[str, Any]: | ||
"""We're overriding the schema classmethod to enable some post-processing""" | ||
schema = super().schema(**kwargs) | ||
cls._rename_key(schema, old_key="anyOf", new_key="oneOf") # UI supports only oneOf | ||
cls._expand_refs(schema) # UI and destination doesn't support $ref's | ||
schema.pop("definitions", None) # remove definitions created by $ref | ||
schema.pop("description", None) # description added from the docstring | ||
return schema |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
76 changes: 76 additions & 0 deletions
76
airbyte-cdk/python/airbyte_cdk/sources/utils/schema_models.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,76 @@ | ||
# | ||
# Copyright (c) 2021 Airbyte, Inc., all rights reserved. | ||
# | ||
|
||
from typing import Any, Dict, Optional, Type | ||
|
||
from pydantic import BaseModel, Extra | ||
from pydantic.main import ModelMetaclass | ||
from pydantic.typing import resolve_annotations | ||
|
||
|
||
class AllOptional(ModelMetaclass): | ||
""" | ||
Metaclass for marking all Pydantic model fields as Optional | ||
Here is example of declaring model using this metaclass like: | ||
''' | ||
class MyModel(BaseModel, metaclass=AllOptional): | ||
a: str | ||
b: str | ||
''' | ||
it is an equivalent of: | ||
''' | ||
class MyModel(BaseModel): | ||
a: Optional[str] | ||
b: Optional[str] | ||
''' | ||
It would make code more clear and eliminate a lot of manual work. | ||
""" | ||
|
||
def __new__(mcs, name, bases, namespaces, **kwargs): | ||
""" | ||
Iterate through fields and wrap then with typing.Optional type. | ||
""" | ||
annotations = resolve_annotations(namespaces.get("__annotations__", {}), namespaces.get("__module__", None)) | ||
for base in bases: | ||
annotations = {**annotations, **getattr(base, "__annotations__", {})} | ||
for field in annotations: | ||
if not field.startswith("__"): | ||
annotations[field] = Optional[annotations[field]] | ||
namespaces["__annotations__"] = annotations | ||
return super().__new__(mcs, name, bases, namespaces, **kwargs) | ||
|
||
|
||
class BaseSchemaModel(BaseModel): | ||
""" | ||
Base class for all schema models. It has some extra schema postprocessing. | ||
Can be used in combination with AllOptional metaclass | ||
""" | ||
|
||
class Config: | ||
extra = Extra.allow | ||
|
||
@classmethod | ||
def schema_extra(cls, schema: Dict[str, Any], model: Type[BaseModel]) -> None: | ||
"""Modify generated jsonschema, remove "title", "description" and "required" fields. | ||
Pydantic doesn't treat Union[None, Any] type correctly when generate jsonschema, | ||
so we can't set field as nullable (i.e. field that can have either null and non-null values), | ||
We generate this jsonschema value manually. | ||
:param schema: generated jsonschema | ||
:param model: | ||
""" | ||
schema.pop("title", None) | ||
schema.pop("description", None) | ||
schema.pop("required", None) | ||
for name, prop in schema.get("properties", {}).items(): | ||
prop.pop("title", None) | ||
prop.pop("description", None) | ||
allow_none = model.__fields__[name].allow_none | ||
if allow_none: | ||
if "type" in prop: | ||
prop["type"] = ["null", prop["type"]] | ||
elif "$ref" in prop: | ||
ref = prop.pop("$ref") | ||
prop["oneOf"] = [{"type": "null"}, {"$ref": ref}] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.