Enforce supervisor schema class name matches its type literal#66899
Enforce supervisor schema class name matches its type literal#66899jason810496 wants to merge 2 commits into
type literal#66899Conversation
c37f6da to
45c5d86
Compare
type literal
| @@ -869,7 +869,7 @@ class GetXComCount(BaseModel): | |||
| dag_id: str | |||
| run_id: str | |||
| task_id: str | |||
| type: Literal["GetNumberXComs"] = "GetNumberXComs" | |||
| type: Literal["GetXComCount"] = "GetXComCount" | |||
There was a problem hiding this comment.
I made the type literal field here in sync with class name.
Additionally, there isn't any Lang-SDK references the supervisor schema so far, so this should not impact anything IMO.
There was a problem hiding this comment.
Yeah I think we are safe here, but just to be sure, can you check if go-sdk is fine reference wise?
There was a problem hiding this comment.
Yes, I just double check that the current Go-SDK go with the Edge Worker API direction. Only the upcoming Coordinator interface will be impacted (but we haven't released any coordinator yet, so we are good!)
| class XComCountResponse(BaseModel): | ||
| len: int | ||
| type: Literal["XComLengthResponse"] = "XComLengthResponse" | ||
| type: Literal["XComCountResponse"] = "XComCountResponse" |
There was a problem hiding this comment.
Ah symmetrically wrong!
| @@ -869,7 +869,7 @@ class GetXComCount(BaseModel): | |||
| dag_id: str | |||
| run_id: str | |||
| task_id: str | |||
| type: Literal["GetNumberXComs"] = "GetNumberXComs" | |||
| type: Literal["GetXComCount"] = "GetXComCount" | |||
There was a problem hiding this comment.
Yeah I think we are safe here, but just to be sure, can you check if go-sdk is fine reference wise?
type literaltype literal
af83d31 to
e513938
Compare
|
Not sure if considered but besides making it static check at coding time, have we considered making this a runtime property handled by Pydantic? |
Why
The six supervisor discriminated unions (
ToTask,ToSupervisor,ToManager,ToDagProcessor,ToTriggerRunner,ToTriggerSupervisor) need every member's class__name__to equal itstype: Literal[...]value soCommsDecoderroutes wire frames to the right class — but nothing enforced that invariant, and two members had silently drifted onmain.What
task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py— a parametrized unit test (one case per union) that walks every member and asserts class__name__equals its single-valuetypeLiteral. Catches drift, a missingtypefield, and multi-value Literals. Runs as part of the task-sdk test suite.task-sdk/src/airflow/sdk/execution_time/comms.py:XComCountResponse:Literal["XComLengthResponse"]→Literal["XComCountResponse"]GetXComCount:Literal["GetNumberXComs"]→Literal["GetXComCount"]task-sdk/tests/task_sdk/execution_time/test_supervisor.pythat pinned the old wire string.Verfication
Was generative AI tooling used to co-author this PR?