Skip to content

Enforce supervisor schema class name matches its type literal#66899

Open
jason810496 wants to merge 2 commits into
apache:mainfrom
jason810496:ci/guard-supervisor-schema-name-and-type-in-sync
Open

Enforce supervisor schema class name matches its type literal#66899
jason810496 wants to merge 2 commits into
apache:mainfrom
jason810496:ci/guard-supervisor-schema-name-and-type-in-sync

Conversation

@jason810496
Copy link
Copy Markdown
Member

@jason810496 jason810496 commented May 14, 2026

Why

The six supervisor discriminated unions (ToTask, ToSupervisor, ToManager, ToDagProcessor, ToTriggerRunner, ToTriggerSupervisor) need every member's class __name__ to equal 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 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-value type Literal. Catches drift, a missing type field, and multi-value Literals. Runs as part of the task-sdk test suite.
  • Fix the two pre-existing mismatches the new test surfaces in task-sdk/src/airflow/sdk/execution_time/comms.py:
    • XComCountResponse: Literal["XComLengthResponse"]Literal["XComCountResponse"]
    • GetXComCount: Literal["GetNumberXComs"]Literal["GetXComCount"]
  • Update the one assertion in task-sdk/tests/task_sdk/execution_time/test_supervisor.py that pinned the old wire string.

Verfication

Screenshot 2026-05-14 at 2 33 09 PM
Was generative AI tooling used to co-author this PR?

@jason810496 jason810496 force-pushed the ci/guard-supervisor-schema-name-and-type-in-sync branch from c37f6da to 45c5d86 Compare May 14, 2026 05:25
@jason810496 jason810496 changed the title Add prek hook to keep supervisor schema class name in sync with its `… Add prek hook to keep supervisor schema class name in sync with its type literal May 14, 2026
@jason810496 jason810496 self-assigned this May 14, 2026
Comment on lines 510 to +872
@@ -869,7 +869,7 @@ class GetXComCount(BaseModel):
dag_id: str
run_id: str
task_id: str
type: Literal["GetNumberXComs"] = "GetNumberXComs"
type: Literal["GetXComCount"] = "GetXComCount"
Copy link
Copy Markdown
Member Author

@jason810496 jason810496 May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I think we are safe here, but just to be sure, can you check if go-sdk is fine reference wise?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!)

Comment thread dev/breeze/doc/images/output_pr_auto-triage.svg Outdated
class XComCountResponse(BaseModel):
len: int
type: Literal["XComLengthResponse"] = "XComLengthResponse"
type: Literal["XComCountResponse"] = "XComCountResponse"
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah symmetrically wrong!

Comment on lines 510 to +872
@@ -869,7 +869,7 @@ class GetXComCount(BaseModel):
dag_id: str
run_id: str
task_id: str
type: Literal["GetNumberXComs"] = "GetNumberXComs"
type: Literal["GetXComCount"] = "GetXComCount"
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I think we are safe here, but just to be sure, can you check if go-sdk is fine reference wise?

Comment thread task-sdk/src/airflow/sdk/execution_time/comms.py
@jason810496 jason810496 marked this pull request as draft May 14, 2026 06:24
@jason810496 jason810496 changed the title Add prek hook to keep supervisor schema class name in sync with its type literal Enforce supervisor schema class name matches its type literal May 14, 2026
@jason810496 jason810496 force-pushed the ci/guard-supervisor-schema-name-and-type-in-sync branch from af83d31 to e513938 Compare May 14, 2026 06:40
@jason810496 jason810496 marked this pull request as ready for review May 14, 2026 06:41
@jscheffl
Copy link
Copy Markdown
Contributor

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)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:dev-tools area:task-sdk backport-to-v3-2-test Mark PR with this label to backport to v3-2-test branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants