fix: fix schema missing indexes/constraints if schemas reference each other

Author: matmoncon

pyneo4j_ogm/core/base.py

@@ -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
@@ -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():