You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jason810496 (via GitHub)" <gi...@apache.org> on 2026/05/14 05:23:46 UTC
[PR] Add prek hook to keep supervisor schema class name in sync with its `… [airflow]
jason810496 opened a new pull request, #66899:
URL: https://github.com/apache/airflow/pull/66899
### Why
The six supervisor discriminated unions (`ToTask`, `ToSupervisor`, `ToManager`, `ToDagProcessor`, `ToTriggerRunner`, `ToTriggerSupervisor`) rely on every member's class `__name__` matching its `type: Literal[...]` value so `CommsDecoder` routes wire frames to the right class — but nothing enforced that invariant, and two members had silently drifted on `main`.
### What
- Add prek hook `check-supervisor-schemas-name-type-sync` that walks all six unions and fails when a member's class `__name__` differs from its `type` literal (or the literal is missing / multi-valued).
- Entrypoint runs in prek's rich-only env; the actual introspection runs under the airflow-core venv via a sibling worker script.
- Scoped to `comms.py`, `processor.py`, `triggerer_job_runner.py` so it only fires when relevant.
- Fix the two pre-existing mismatches the new hook surfaces:
- `XComCountResponse`: `Literal["XComLengthResponse"]` → `Literal["XComCountResponse"]`
- `GetXComCount`: `Literal["GetNumberXComs"]` → `Literal["GetXComCount"]`
- Update the one test assertion that pinned the old wire string in `test_supervisor.py`.
- Add pytest coverage for the hook's classifier (pass / rename-drift / missing `type` / multi-value Literal / cross-union dedup) and for the `.pre-commit-config.yaml` wiring.
### Verfication
<img width="1540" height="221" alt="Screenshot 2026-05-14 at 12 10 48 PM" src="https://github.com/user-attachments/assets/89e580e8-b1a8-4406-9d1c-930c6792d90c" />
---
##### Was generative AI tooling used to co-author this PR?
- [x] Yes, with help of Claude Code Opus 4.7 following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enforce supervisor schema class name matches its `type` literal [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on PR #66899:
URL: https://github.com/apache/airflow/pull/66899#issuecomment-4449419221
Not sure if considered but besides making it static check at coding time, have we considered making this a runtime property handled by Pydantic?
```
class BaseCommClass(BaseModel):
"""Base type class."""
type: str = "__undefined__"
def __init__(self):
super().__init__()
self.type = self.__class__.__name__
class MyType1(BaseCommClass):
"""Another type."""
pass
class MyType2(BaseCommClass):
"""Another type."""
pass
print(MyType1().type)
print(MyType2().type)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enforce supervisor schema class name matches its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 merged PR #66899:
URL: https://github.com/apache/airflow/pull/66899
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add prek hook to keep supervisor schema class name in sync with its `type` literal [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66899:
URL: https://github.com/apache/airflow/pull/66899#discussion_r3239356906
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -509,7 +509,7 @@ def from_xcom_response(cls, xcom_response: XComResponse) -> XComResult:
class XComCountResponse(BaseModel):
len: int
- type: Literal["XComLengthResponse"] = "XComLengthResponse"
+ type: Literal["XComCountResponse"] = "XComCountResponse"
Review Comment:
Ah symmetrically wrong!
##########
dev/breeze/doc/images/output_pr_auto-triage.svg:
##########
Review Comment:
Unrelated change?
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -869,7 +869,7 @@ class GetXComCount(BaseModel):
dag_id: str
run_id: str
task_id: str
- type: Literal["GetNumberXComs"] = "GetNumberXComs"
+ type: Literal["GetXComCount"] = "GetXComCount"
Review Comment:
Yeah I think we are safe here, but just to be sure, can you check if go-sdk is fine reference wise?
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
Review Comment:
Just wondering if you thought of a unit test covering this would be simpler to maintain compared to a prek hook?
Something like:
```python
def test_supervisor_schema_type_literals_match_class_names():
from typing import Annotated, get_args, get_origin
from airflow.sdk.execution_time.comms import ToTask, ToSupervisor
for union_name, union in [("ToTask", ToTask), ("ToSupervisor", ToSupervisor)]:
union_type = get_args(union)[0] if get_origin(union) is Annotated else union
for member in get_args(union_type):
if not isinstance(member, type):
continue
field = member.model_fields.get("type")
assert field is not None, f"{member.__name__} in {union_name} has no `type` field"
args = get_args(field.annotation)
assert len(args) == 1 and args[0] == member.__name__, (
f"{member.__name__}.type literal is {args[0]!r}, expected {member.__name__!r}"
)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enforce supervisor schema class name matches its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on PR #66899:
URL: https://github.com/apache/airflow/pull/66899#issuecomment-4456837127
Tried it. Runtime-binding changes the schema and also breaks the tagged union itself, because Pydantic v2 needs the discriminator field to be a `Literal`.
```python
from typing import Annotated, Literal, Union
from pydantic import BaseModel, Field, TypeAdapter
# --- Current pattern: Literal type field ---
class A1(BaseModel):
type: Literal["A1"] = "A1"
class A2(BaseModel):
type: Literal["A2"] = "A2"
print(A1.model_json_schema())
#. {'properties': {'type': {'const': 'A1', 'default': 'A1', 'title': 'Type', 'type': 'string'}}, 'title': 'A1', 'type': 'object'}
TypeAdapter(Annotated[Union[A1, A2], Field(discriminator="type")]) # OK
# --- Proposed pattern: runtime-bind in __init__ ---
class B(BaseModel):
type: str = "__undefined__"
def __init__(self, **kw):
super().__init__(**kw)
self.type = type(self).__name__
class B1(B): pass
class B2(B): pass
print(B1.model_json_schema())
# {'properties': {'type': {'default': '__undefined__', 'title': 'Type', 'type': 'string'}}, 'title': 'B1', 'type': 'object'}
print(B1().type) # 'B1' (instance state only; not reflected in schema)
TypeAdapter(Annotated[Union[B1, B2], Field(discriminator="type")])
# PydanticUserError: Model 'B1' needs field 'type' to be of type `Literal`
```
So `CommsDecoder`'s `TypeAdapter`-based routing relies on the `Literal` annotation, and any downstream schema consumer relies on the `const` it produces.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enforce supervisor schema class name matches its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on PR #66899:
URL: https://github.com/apache/airflow/pull/66899#issuecomment-4456555941
> Not sure if considered but besides making it static check at coding time, have we considered making this a runtime property handled by Pydantic?
I will follow-up on this one. BTW, not sure if there's any reason to explictly specify the `typ` literal field (e.g. binding the field at runtime might change the generated schema?) I will check it out.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add prek hook to keep supervisor schema class name in sync with its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66899:
URL: https://github.com/apache/airflow/pull/66899#discussion_r3239491745
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
Review Comment:
Good idea! Having it as unit tests will be much simpler than the prek hook. (CI will raise the error anyway even the drift will not be caught in the first place when commit the changes)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add prek hook to keep supervisor schema class name in sync with its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66899:
URL: https://github.com/apache/airflow/pull/66899#discussion_r3239274858
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -869,7 +869,7 @@ class GetXComCount(BaseModel):
dag_id: str
run_id: str
task_id: str
- type: Literal["GetNumberXComs"] = "GetNumberXComs"
+ type: Literal["GetXComCount"] = "GetXComCount"
Review 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.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add prek hook to keep supervisor schema class name in sync with its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66899:
URL: https://github.com/apache/airflow/pull/66899#discussion_r3239274858
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -869,7 +869,7 @@ class GetXComCount(BaseModel):
dag_id: str
run_id: str
task_id: str
- type: Literal["GetNumberXComs"] = "GetNumberXComs"
+ type: Literal["GetXComCount"] = "GetXComCount"
Review Comment:
I made the `type` 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.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enforce supervisor schema class name matches its `type` literal [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66899:
URL: https://github.com/apache/airflow/pull/66899#discussion_r3239568120
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -869,7 +869,7 @@ class GetXComCount(BaseModel):
dag_id: str
run_id: str
task_id: str
- type: Literal["GetNumberXComs"] = "GetNumberXComs"
+ type: Literal["GetXComCount"] = "GetXComCount"
Review 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!)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org