Skip to content

Commit

Permalink
fix: fix schema missing indexes/constraints if schemas reference each…
Browse files Browse the repository at this point in the history
… other
  • Loading branch information
matmoncon committed Jan 17, 2024
1 parent 6c7553d commit eb4d666
Showing 1 changed file with 16 additions and 1 deletion.
17 changes: 16 additions & 1 deletion pyneo4j_ogm/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@

if IS_PYDANTIC_V2:
from pydantic import SerializationInfo, model_serializer, model_validator
from pydantic.errors import PydanticSchemaGenerationError
from pydantic.json_schema import GenerateJsonSchema
else:
from pydantic.class_validators import root_validator
Expand Down Expand Up @@ -160,7 +161,21 @@ class CustomGenerateJsonSchema(GenerateJsonSchema):
"""

def generate(self, *args, **kwargs):
model_cls = cast(Type[BaseModel], args[0]["schema"]["cls"])
model_cls: Optional[Type[BaseModel]] = None

if "definitions" in args[0]:
schema_ref = args[0]["schema"]["schema_ref"]

for definition in args[0]["definitions"]:
if definition["ref"] == schema_ref:
model_cls = cast(Type[BaseModel], definition["schema"]["cls"])
break
else:
model_cls = cast(Type[BaseModel], args[0]["schema"]["cls"]) if "cls" in args[0]["schema"] else None

if model_cls is None:
raise PydanticSchemaGenerationError("Could not find model class in definitions")

generated_schema = super().generate(*args, **kwargs)

for field_name, field in get_model_fields(model_cls).items():
Expand Down

0 comments on commit eb4d666

Please sign in to comment.