jason810496 commented on code in PR #67596:
URL: https://github.com/apache/airflow/pull/67596#discussion_r3310370508
##########
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:
Does the "non-SDK" modules mean the data models define in airflow-core like
`DagFileParseRequest`?
https://github.com/apache/airflow/blob/1bb4276e84849f7a266c9f30c0f5c8492bd259c7/airflow-core/src/airflow/dag_processing/processor.py#L109-L127
--
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]
