uranusjr commented on code in PR #67596:
URL: https://github.com/apache/airflow/pull/67596#discussion_r3310553942
##########
scripts/ci/prek/dump_supervisor_schemas.py:
##########
@@ -31,25 +31,49 @@
import json
import os
import sys
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
- from pydantic import BaseModel
+from pydantic import json_schema
os.environ["_AIRFLOW__AS_LIBRARY"] = "1"
from airflow.sdk.execution_time.schema import bundle, registered_models_by_name
-
-def _registered_models_sorted() -> tuple[type[BaseModel], ...]:
- """Return registered head models sorted by class name for stable snapshot
diffs."""
- by_name = registered_models_by_name()
- return tuple(by_name[name] for name in sorted(by_name))
-
-
+# Models are sorted by name before being passed to generation for stable diffs.
+_, snapshot = json_schema.models_json_schema(
+ [(m, "serialization") for _, m in
sorted(registered_models_by_name().items())],
+ schema_generator=json_schema.GenerateJsonSchema,
+)
snapshot = {
+ "$schema": json_schema.GenerateJsonSchema.schema_dialect,
"api_version": str(bundle.versions[0].value),
- "schemas": {cls.__name__: cls.model_json_schema() for cls in
_registered_models_sorted()},
+ "description": "Apache Airflow SDK Supervisor Schema",
+ "$defs": snapshot["$defs"],
}
-json.dump(snapshot, sys.stdout, indent=2, sort_keys=True)
+
+# We currently still have references to non-SDK models from SDK. This cause
Review Comment:
I don’t remember which ones from the top of my head, but the problem seems
to be some SDK models indirectly reference models defined in
`airflow.api_fastapi` instead of `airflow.sdk`. Since the latter is generated
from the former, you get some defs with the same name. Pydantic resolves this
by renaming those to use something like
`airflow__sdk__execution_time__AssetResult` (not a real example), which is just
aweful and causes more problems since unlike in Python, you can’t just convince
a static-typed language to treat an object to be a different class with the
same shape.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
