diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/AGENTS.md b/airflow-core/src/airflow/api_fastapi/execution_api/AGENTS.md index aff91c070d9ef..4cacd2368baf0 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/AGENTS.md +++ b/airflow-core/src/airflow/api_fastapi/execution_api/AGENTS.md @@ -16,7 +16,7 @@ Workers and API servers deploy independently, so backward compatibility is criti 3. Regenerate Task SDK models: ```bash -cd task-sdk && python dev/generate_task_sdk_models.py +uv run --active --group codegen --project apache-airflow-task-sdk --directory task-sdk -s dev/generate_task_sdk_models.py ``` 4. Add tests for both the new and previous API versions. @@ -44,7 +44,7 @@ Adding a new Execution API feature touches multiple packages. All of these must 5. **Task SDK client** — add the client method in `task-sdk/src/airflow/sdk/api/client.py`. 6. **Supervisor** — handle the new message in `task-sdk/src/airflow/sdk/execution_time/supervisor.py`. 7. **Dag processor & triggerer exclusions** — these use `InProcessExecutionAPI` and have explicit message type unions. Add new types to their handler or exclusion lists in `airflow/dag_processing/processor.py` and `airflow/jobs/triggerer_job_runner.py`. -8. **Regenerate models** — `cd task-sdk && python dev/generate_task_sdk_models.py`. +8. **Regenerate models** — `uv run --active --group codegen --project apache-airflow-task-sdk --directory task-sdk -s dev/generate_task_sdk_models.py`. 9. **Tests** — if the new message type requires an API endpoint, add tests in all of these: - `airflow-core/tests/unit/api_fastapi/execution_api/` — endpoint tests - `task-sdk/tests/task_sdk/api/test_client.py` — client method tests diff --git a/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py b/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py index b2c5ad1b14071..ec661a156f259 100644 --- a/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py +++ b/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py @@ -123,7 +123,7 @@ def mock_get_connection(conn_id): host=conn.host, login=conn.login, password=conn.password, - schema_=conn.schema, + schema=conn.schema, port=conn.port, extra=conn.extra, ) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 2da26aab1241b..dcdd614fb7a2a 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2167,6 +2167,10 @@ def test_handle_request_get_connection_masks_password_and_extra(self, proc): conn_type="mysql", password="super-secret-password", extra='{"api_key":"super-secret-extra"}', + host=None, + schema=None, + login=None, + port=None, ) with ( @@ -2197,6 +2201,10 @@ def test_handle_request_get_connection_masks_password_and_extra(self, proc): "conn_type": "mysql", "password": "super-secret-password", "extra": '{"api_key":"super-secret-extra"}', + "host": None, + "schema": None, + "login": None, + "port": None, "type": "ConnectionResult", } diff --git a/airflow-ctl/pyproject.toml b/airflow-ctl/pyproject.toml index 961b4a6c23eef..a29986429c59b 100644 --- a/airflow-ctl/pyproject.toml +++ b/airflow-ctl/pyproject.toml @@ -149,7 +149,7 @@ dev = [ "apache-airflow-devel-common", ] codegen = [ - "datamodel-code-generator[http]==0.33.0", + "datamodel-code-generator[http]>=0.71.0", "apache-airflow-devel-common" ] diff --git a/airflow-ctl/src/airflowctl/api/datamodels/auth_generated.py b/airflow-ctl/src/airflowctl/api/datamodels/auth_generated.py index ec8d9c16b22a1..9204ee0c583b9 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/auth_generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/auth_generated.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: v2-simple-auth-manager-generated.yaml -# version: 0.33.0 +# version: 0.71.0 from __future__ import annotations diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 746ba1e02dabf..60a6cd33c965c 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: v2-rest-api-generated.yaml -# version: 0.33.0 +# version: 0.71.0 from __future__ import annotations @@ -102,7 +102,7 @@ class BaseInfoResponse(BaseModel): Base info serializer for responses. """ - status: Annotated[str | None, Field(title="Status")] = None + status: Annotated[str | None, Field(title="Status")] class BulkActionNotOnExistence(str, Enum): @@ -229,8 +229,8 @@ class ClearPartitionsResponse(BaseModel): dry_run: Annotated[bool, Field(title="Dry Run")] -class TaskIds(RootModel[list]): - root: Annotated[list, Field(max_length=2, min_length=2)] +class TaskIds(RootModel[tuple[str, int]]): + root: Annotated[tuple[str, int], Field(max_length=2, min_length=2)] class ClearTaskInstancesBody(BaseModel): @@ -270,8 +270,8 @@ class ClearTaskInstancesBody(BaseModel): note: Annotated[Note | None, Field(title="Note")] = None -class Value(RootModel[list]): - root: Annotated[list, Field(max_length=2, min_length=2, title="Value")] +class Value(RootModel[tuple[str, str]]): + root: Annotated[tuple[str, str], Field(max_length=2, min_length=2, title="Value")] class ConfigOption(BaseModel): @@ -329,14 +329,14 @@ class ConnectionResponse(BaseModel): connection_id: Annotated[str, Field(title="Connection Id")] conn_type: Annotated[str, Field(title="Conn Type")] - description: Annotated[str | None, Field(title="Description")] = None - host: Annotated[str | None, Field(title="Host")] = None - login: Annotated[str | None, Field(title="Login")] = None - schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None - port: Annotated[int | None, Field(title="Port")] = None - password: Annotated[str | None, Field(title="Password")] = None - extra: Annotated[str | None, Field(title="Extra")] = None - team_name: Annotated[str | None, Field(title="Team Name")] = None + description: Annotated[str | None, Field(title="Description")] + host: Annotated[str | None, Field(title="Host")] + login: Annotated[str | None, Field(title="Login")] + schema_: Annotated[str | None, Field(alias="schema", title="Schema")] + port: Annotated[int | None, Field(title="Port")] + password: Annotated[str | None, Field(title="Password")] + extra: Annotated[str | None, Field(title="Extra")] + team_name: Annotated[str | None, Field(title="Team Name")] class ConnectionTestQueuedResponse(BaseModel): @@ -459,9 +459,9 @@ class DAGSourceResponse(BaseModel): Dag Source serializer for responses. """ - content: Annotated[str | None, Field(title="Content")] = None + content: Annotated[str | None, Field(title="Content")] dag_id: Annotated[str, Field(title="Dag Id")] - version_number: Annotated[int | None, Field(title="Version Number")] = None + version_number: Annotated[int | None, Field(title="Version Number")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] @@ -479,10 +479,8 @@ class DagProcessorInfoResponse(BaseModel): DagProcessor info serializer for responses. """ - status: Annotated[str | None, Field(title="Status")] = None - latest_dag_processor_heartbeat: Annotated[str | None, Field(title="Latest Dag Processor Heartbeat")] = ( - None - ) + status: Annotated[str | None, Field(title="Status")] + latest_dag_processor_heartbeat: Annotated[str | None, Field(title="Latest Dag Processor Heartbeat")] class DagRunAssetReference(BaseModel): @@ -495,13 +493,13 @@ class DagRunAssetReference(BaseModel): ) run_id: Annotated[str, Field(title="Run Id")] dag_id: Annotated[str, Field(title="Dag Id")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] start_date: Annotated[datetime, Field(title="Start Date")] - end_date: Annotated[datetime | None, Field(title="End Date")] = None + end_date: Annotated[datetime | None, Field(title="End Date")] state: Annotated[str, Field(title="State")] - data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] = None - data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] = None - partition_key: Annotated[str | None, Field(title="Partition Key")] = None + data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] + data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] + partition_key: Annotated[str | None, Field(title="Partition Key")] class DagRunMutableStates(str, Enum): @@ -597,11 +595,11 @@ class DagVersionResponse(BaseModel): id: Annotated[UUID, Field(title="Id")] version_number: Annotated[int, Field(title="Version Number")] dag_id: Annotated[str, Field(title="Dag Id")] - bundle_name: Annotated[str | None, Field(title="Bundle Name")] = None - bundle_version: Annotated[str | None, Field(title="Bundle Version")] = None + bundle_name: Annotated[str | None, Field(title="Bundle Name")] + bundle_version: Annotated[str | None, Field(title="Bundle Version")] created_at: Annotated[datetime, Field(title="Created At")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] - bundle_url: Annotated[str | None, Field(title="Bundle Url")] = None + bundle_url: Annotated[str | None, Field(title="Bundle Url")] class DagWarningType(str, Enum): @@ -622,9 +620,9 @@ class DryRunBackfillResponse(BaseModel): Backfill serializer for responses in dry-run mode. """ - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None - partition_key: Annotated[str | None, Field(title="Partition Key")] = None - partition_date: Annotated[datetime | None, Field(title="Partition Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] + partition_key: Annotated[str | None, Field(title="Partition Key")] + partition_date: Annotated[datetime | None, Field(title="Partition Date")] class EventLogResponse(BaseModel): @@ -634,15 +632,15 @@ class EventLogResponse(BaseModel): event_log_id: Annotated[int, Field(title="Event Log Id")] when: Annotated[datetime, Field(title="When")] - dag_id: Annotated[str | None, Field(title="Dag Id")] = None - task_id: Annotated[str | None, Field(title="Task Id")] = None - run_id: Annotated[str | None, Field(title="Run Id")] = None - map_index: Annotated[int | None, Field(title="Map Index")] = None - try_number: Annotated[int | None, Field(title="Try Number")] = None + dag_id: Annotated[str | None, Field(title="Dag Id")] + task_id: Annotated[str | None, Field(title="Task Id")] + run_id: Annotated[str | None, Field(title="Run Id")] + map_index: Annotated[int | None, Field(title="Map Index")] + try_number: Annotated[int | None, Field(title="Try Number")] event: Annotated[str, Field(title="Event")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None - owner: Annotated[str | None, Field(title="Owner")] = None - extra: Annotated[str | None, Field(title="Extra")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] + owner: Annotated[str | None, Field(title="Owner")] + extra: Annotated[str | None, Field(title="Extra")] dag_display_name: Annotated[str | None, Field(title="Dag Display Name")] = None task_display_name: Annotated[str | None, Field(title="Task Display Name")] = None @@ -741,7 +739,7 @@ class ImportErrorResponse(BaseModel): import_error_id: Annotated[int, Field(title="Import Error Id")] timestamp: Annotated[datetime, Field(title="Timestamp")] filename: Annotated[str, Field(title="Filename")] - bundle_name: Annotated[str | None, Field(title="Bundle Name")] = None + bundle_name: Annotated[str | None, Field(title="Bundle Name")] stack_trace: Annotated[str, Field(title="Stack Trace")] @@ -751,15 +749,15 @@ class JobResponse(BaseModel): """ id: Annotated[int, Field(title="Id")] - dag_id: Annotated[str | None, Field(title="Dag Id")] = None - state: Annotated[str | None, Field(title="State")] = None - job_type: Annotated[str | None, Field(title="Job Type")] = None - start_date: Annotated[datetime | None, Field(title="Start Date")] = None - end_date: Annotated[datetime | None, Field(title="End Date")] = None - latest_heartbeat: Annotated[datetime | None, Field(title="Latest Heartbeat")] = None - executor_class: Annotated[str | None, Field(title="Executor Class")] = None - hostname: Annotated[str | None, Field(title="Hostname")] = None - unixname: Annotated[str | None, Field(title="Unixname")] = None + dag_id: Annotated[str | None, Field(title="Dag Id")] + state: Annotated[str | None, Field(title="State")] + job_type: Annotated[str | None, Field(title="Job Type")] + start_date: Annotated[datetime | None, Field(title="Start Date")] + end_date: Annotated[datetime | None, Field(title="End Date")] + latest_heartbeat: Annotated[datetime | None, Field(title="Latest Heartbeat")] + executor_class: Annotated[str | None, Field(title="Executor Class")] + hostname: Annotated[str | None, Field(title="Hostname")] + unixname: Annotated[str | None, Field(title="Unixname")] dag_display_name: Annotated[str | None, Field(title="Dag Display Name")] = None @@ -865,7 +863,7 @@ class PoolResponse(BaseModel): scheduled_slots: Annotated[int, Field(title="Scheduled Slots")] open_slots: Annotated[int, Field(title="Open Slots")] deferred_slots: Annotated[int, Field(title="Deferred Slots")] - team_name: Annotated[str | None, Field(title="Team Name")] = None + team_name: Annotated[str | None, Field(title="Team Name")] class ProviderResponse(BaseModel): @@ -876,7 +874,7 @@ class ProviderResponse(BaseModel): package_name: Annotated[str, Field(title="Package Name")] description: Annotated[str, Field(title="Description")] version: Annotated[str, Field(title="Version")] - documentation_url: Annotated[str | None, Field(title="Documentation Url")] = None + documentation_url: Annotated[str | None, Field(title="Documentation Url")] class QueuedEventResponse(BaseModel): @@ -935,8 +933,8 @@ class SchedulerInfoResponse(BaseModel): Scheduler info serializer for responses. """ - status: Annotated[str | None, Field(title="Status")] = None - latest_scheduler_heartbeat: Annotated[str | None, Field(title="Latest Scheduler Heartbeat")] = None + status: Annotated[str | None, Field(title="Status")] + latest_scheduler_heartbeat: Annotated[str | None, Field(title="Latest Scheduler Heartbeat")] class StructuredLogMessage(BaseModel): @@ -1042,7 +1040,7 @@ class TaskInstancesLogResponse(BaseModel): """ content: Annotated[list[StructuredLogMessage] | list[str], Field(title="Content")] - continuation_token: Annotated[str | None, Field(title="Continuation Token")] = None + continuation_token: Annotated[str | None, Field(title="Continuation Token")] class TaskOutletAssetReference(BaseModel): @@ -1074,7 +1072,7 @@ class TaskStateStoreBody(BaseModel): extra="forbid", ) value: JsonValue - expires_at: Annotated[datetime | str | None, Field(title="Expires At")] = "default" + expires_at: Annotated[datetime | Literal["default"] | None, Field(title="Expires At")] = "default" class TaskStateStorePatchBody(BaseModel): @@ -1096,7 +1094,7 @@ class TaskStateStoreResponse(BaseModel): key: Annotated[str, Field(title="Key")] value: JsonValue updated_at: Annotated[datetime, Field(title="Updated At")] - expires_at: Annotated[datetime | None, Field(title="Expires At")] = None + expires_at: Annotated[datetime | None, Field(title="Expires At")] class TimeDelta(BaseModel): @@ -1121,7 +1119,7 @@ class TriggerDAGRunPostBody(BaseModel): dag_run_id: Annotated[str | None, Field(title="Dag Run Id")] = None data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] = None data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] = None - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] run_after: Annotated[datetime | None, Field(title="Run After")] = None conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None note: Annotated[str | None, Field(title="Note")] = None @@ -1135,10 +1133,10 @@ class TriggerResponse(BaseModel): id: Annotated[int, Field(title="Id")] classpath: Annotated[str, Field(title="Classpath")] - kwargs: Annotated[str, Field(title="Kwargs")] + kwargs: Annotated[str, Field(deprecated=True, title="Kwargs")] created_date: Annotated[datetime, Field(title="Created Date")] - queue: Annotated[str | None, Field(title="Queue")] = None - triggerer_id: Annotated[int | None, Field(title="Triggerer Id")] = None + queue: Annotated[str | None, Field(title="Queue")] + triggerer_id: Annotated[int | None, Field(title="Triggerer Id")] class TriggererInfoResponse(BaseModel): @@ -1146,8 +1144,8 @@ class TriggererInfoResponse(BaseModel): Triggerer info serializer for responses. """ - status: Annotated[str | None, Field(title="Status")] = None - latest_triggerer_heartbeat: Annotated[str | None, Field(title="Latest Triggerer Heartbeat")] = None + status: Annotated[str | None, Field(title="Status")] + latest_triggerer_heartbeat: Annotated[str | None, Field(title="Latest Triggerer Heartbeat")] class UpdateHITLDetailPayload(BaseModel): @@ -1188,9 +1186,9 @@ class VariableResponse(BaseModel): key: Annotated[str, Field(title="Key")] value: Annotated[str | None, Field(title="Value")] = None - description: Annotated[str | None, Field(title="Description")] = None + description: Annotated[str | None, Field(title="Description")] is_encrypted: Annotated[bool, Field(title="Is Encrypted")] - team_name: Annotated[str | None, Field(title="Team Name")] = None + team_name: Annotated[str | None, Field(title="Team Name")] class VersionInfo(BaseModel): @@ -1199,7 +1197,7 @@ class VersionInfo(BaseModel): """ version: Annotated[str, Field(title="Version")] - git_version: Annotated[str | None, Field(title="Git Version")] = None + git_version: Annotated[str | None, Field(title="Git Version")] class XComCreateBody(BaseModel): @@ -1222,7 +1220,7 @@ class XComResponse(BaseModel): key: Annotated[str, Field(title="Key")] timestamp: Annotated[datetime, Field(title="Timestamp")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] map_index: Annotated[int, Field(title="Map Index")] task_id: Annotated[str, Field(title="Task Id")] dag_id: Annotated[str, Field(title="Dag Id")] @@ -1239,7 +1237,7 @@ class XComResponseNative(BaseModel): key: Annotated[str, Field(title="Key")] timestamp: Annotated[datetime, Field(title="Timestamp")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] map_index: Annotated[int, Field(title="Map Index")] task_id: Annotated[str, Field(title="Task Id")] dag_id: Annotated[str, Field(title="Dag Id")] @@ -1257,7 +1255,7 @@ class XComResponseString(BaseModel): key: Annotated[str, Field(title="Key")] timestamp: Annotated[datetime, Field(title="Timestamp")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] map_index: Annotated[int, Field(title="Map Index")] task_id: Annotated[str, Field(title="Task Id")] dag_id: Annotated[str, Field(title="Dag Id")] @@ -1265,7 +1263,7 @@ class XComResponseString(BaseModel): dag_display_name: Annotated[str, Field(title="Dag Display Name")] task_display_name: Annotated[str, Field(title="Task Display Name")] run_after: Annotated[datetime, Field(title="Run After")] - value: Annotated[str | None, Field(title="Value")] = None + value: Annotated[str | None, Field(title="Value")] class XComUpdateBody(BaseModel): @@ -1396,12 +1394,12 @@ class BackfillResponse(BaseModel): dag_id: Annotated[str, Field(title="Dag Id")] from_date: Annotated[datetime, Field(title="From Date")] to_date: Annotated[datetime, Field(title="To Date")] - dag_run_conf: Annotated[dict[str, Any] | None, Field(title="Dag Run Conf")] = None + dag_run_conf: Annotated[dict[str, Any] | None, Field(title="Dag Run Conf")] is_paused: Annotated[bool, Field(title="Is Paused")] reprocess_behavior: ReprocessBehavior max_active_runs: Annotated[int, Field(title="Max Active Runs")] created_at: Annotated[datetime, Field(title="Created At")] - completed_at: Annotated[datetime | None, Field(title="Completed At")] = None + completed_at: Annotated[datetime | None, Field(title="Completed At")] updated_at: Annotated[datetime, Field(title="Updated At")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] @@ -1713,47 +1711,45 @@ class DAGDetailsResponse(BaseModel): dag_display_name: Annotated[str, Field(title="Dag Display Name")] is_paused: Annotated[bool, Field(title="Is Paused")] is_stale: Annotated[bool, Field(title="Is Stale")] - last_parsed_time: Annotated[datetime | None, Field(title="Last Parsed Time")] = None - last_parse_duration: Annotated[float | None, Field(title="Last Parse Duration")] = None - last_expired: Annotated[datetime | None, Field(title="Last Expired")] = None - bundle_name: Annotated[str | None, Field(title="Bundle Name")] = None - bundle_version: Annotated[str | None, Field(title="Bundle Version")] = None - relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] = None + last_parsed_time: Annotated[datetime | None, Field(title="Last Parsed Time")] + last_parse_duration: Annotated[float | None, Field(title="Last Parse Duration")] + last_expired: Annotated[datetime | None, Field(title="Last Expired")] + bundle_name: Annotated[str | None, Field(title="Bundle Name")] + bundle_version: Annotated[str | None, Field(title="Bundle Version")] + relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] fileloc: Annotated[str, Field(title="Fileloc")] - description: Annotated[str | None, Field(title="Description")] = None - timetable_summary: Annotated[str | None, Field(title="Timetable Summary")] = None - timetable_description: Annotated[str | None, Field(title="Timetable Description")] = None + description: Annotated[str | None, Field(title="Description")] + timetable_summary: Annotated[str | None, Field(title="Timetable Summary")] + timetable_description: Annotated[str | None, Field(title="Timetable Description")] timetable_partitioned: Annotated[bool, Field(title="Timetable Partitioned")] timetable_periodic: Annotated[bool, Field(title="Timetable Periodic")] tags: Annotated[list[DagTagResponse], Field(title="Tags")] max_active_tasks: Annotated[int, Field(title="Max Active Tasks")] - max_active_runs: Annotated[int | None, Field(title="Max Active Runs")] = None + max_active_runs: Annotated[int | None, Field(title="Max Active Runs")] max_consecutive_failed_dag_runs: Annotated[int, Field(title="Max Consecutive Failed Dag Runs")] has_task_concurrency_limits: Annotated[bool, Field(title="Has Task Concurrency Limits")] has_import_errors: Annotated[bool, Field(title="Has Import Errors")] - next_dagrun_logical_date: Annotated[datetime | None, Field(title="Next Dagrun Logical Date")] = None + next_dagrun_logical_date: Annotated[datetime | None, Field(title="Next Dagrun Logical Date")] next_dagrun_data_interval_start: Annotated[ datetime | None, Field(title="Next Dagrun Data Interval Start") - ] = None - next_dagrun_data_interval_end: Annotated[ - datetime | None, Field(title="Next Dagrun Data Interval End") - ] = None - next_dagrun_run_after: Annotated[datetime | None, Field(title="Next Dagrun Run After")] = None - allowed_run_types: Annotated[list[DagRunType] | None, Field(title="Allowed Run Types")] = None + ] + next_dagrun_data_interval_end: Annotated[datetime | None, Field(title="Next Dagrun Data Interval End")] + next_dagrun_run_after: Annotated[datetime | None, Field(title="Next Dagrun Run After")] + allowed_run_types: Annotated[list[DagRunType] | None, Field(title="Allowed Run Types")] owners: Annotated[list[str], Field(title="Owners")] catchup: Annotated[bool, Field(title="Catchup")] - dag_run_timeout: Annotated[timedelta | None, Field(title="Dag Run Timeout")] = None - asset_expression: Annotated[dict[str, Any] | None, Field(title="Asset Expression")] = None - doc_md: Annotated[str | None, Field(title="Doc Md")] = None - start_date: Annotated[datetime | None, Field(title="Start Date")] = None - end_date: Annotated[datetime | None, Field(title="End Date")] = None - is_paused_upon_creation: Annotated[bool | None, Field(title="Is Paused Upon Creation")] = None - params: Annotated[dict[str, Any] | None, Field(title="Params")] = None + dag_run_timeout: Annotated[timedelta | None, Field(title="Dag Run Timeout")] + asset_expression: Annotated[dict[str, Any] | None, Field(title="Asset Expression")] + doc_md: Annotated[str | None, Field(title="Doc Md")] + start_date: Annotated[datetime | None, Field(title="Start Date")] + end_date: Annotated[datetime | None, Field(title="End Date")] + is_paused_upon_creation: Annotated[bool | None, Field(title="Is Paused Upon Creation")] + params: Annotated[dict[str, Any] | None, Field(title="Params")] render_template_as_native_obj: Annotated[bool, Field(title="Render Template As Native Obj")] - template_search_path: Annotated[list[str] | None, Field(title="Template Search Path")] = None - timezone: Annotated[str | None, Field(title="Timezone")] = None - last_parsed: Annotated[datetime | None, Field(title="Last Parsed")] = None - default_args: Annotated[dict[str, Any] | None, Field(title="Default Args")] = None + template_search_path: Annotated[list[str] | None, Field(title="Template Search Path")] + timezone: Annotated[str | None, Field(title="Timezone")] + last_parsed: Annotated[datetime | None, Field(title="Last Parsed")] + default_args: Annotated[dict[str, Any] | None, Field(title="Default Args")] rerun_with_latest_version: Annotated[bool | None, Field(title="Rerun With Latest Version")] = None owner_links: Annotated[dict[str, str] | None, Field(title="Owner Links")] = None is_favorite: Annotated[bool | None, Field(title="Is Favorite")] = False @@ -1765,13 +1761,14 @@ class DAGDetailsResponse(BaseModel): concurrency: Annotated[ int, Field( + deprecated=True, description="Return max_active_tasks as concurrency.\n\nDeprecated: Use max_active_tasks instead.", title="Concurrency", ), ] latest_dag_version: Annotated[ DagVersionResponse | None, Field(description="Return the latest DagVersion.") - ] = None + ] class DAGResponse(BaseModel): @@ -1783,33 +1780,31 @@ class DAGResponse(BaseModel): dag_display_name: Annotated[str, Field(title="Dag Display Name")] is_paused: Annotated[bool, Field(title="Is Paused")] is_stale: Annotated[bool, Field(title="Is Stale")] - last_parsed_time: Annotated[datetime | None, Field(title="Last Parsed Time")] = None - last_parse_duration: Annotated[float | None, Field(title="Last Parse Duration")] = None - last_expired: Annotated[datetime | None, Field(title="Last Expired")] = None - bundle_name: Annotated[str | None, Field(title="Bundle Name")] = None - bundle_version: Annotated[str | None, Field(title="Bundle Version")] = None - relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] = None + last_parsed_time: Annotated[datetime | None, Field(title="Last Parsed Time")] + last_parse_duration: Annotated[float | None, Field(title="Last Parse Duration")] + last_expired: Annotated[datetime | None, Field(title="Last Expired")] + bundle_name: Annotated[str | None, Field(title="Bundle Name")] + bundle_version: Annotated[str | None, Field(title="Bundle Version")] + relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] fileloc: Annotated[str, Field(title="Fileloc")] - description: Annotated[str | None, Field(title="Description")] = None - timetable_summary: Annotated[str | None, Field(title="Timetable Summary")] = None - timetable_description: Annotated[str | None, Field(title="Timetable Description")] = None + description: Annotated[str | None, Field(title="Description")] + timetable_summary: Annotated[str | None, Field(title="Timetable Summary")] + timetable_description: Annotated[str | None, Field(title="Timetable Description")] timetable_partitioned: Annotated[bool, Field(title="Timetable Partitioned")] timetable_periodic: Annotated[bool, Field(title="Timetable Periodic")] tags: Annotated[list[DagTagResponse], Field(title="Tags")] max_active_tasks: Annotated[int, Field(title="Max Active Tasks")] - max_active_runs: Annotated[int | None, Field(title="Max Active Runs")] = None + max_active_runs: Annotated[int | None, Field(title="Max Active Runs")] max_consecutive_failed_dag_runs: Annotated[int, Field(title="Max Consecutive Failed Dag Runs")] has_task_concurrency_limits: Annotated[bool, Field(title="Has Task Concurrency Limits")] has_import_errors: Annotated[bool, Field(title="Has Import Errors")] - next_dagrun_logical_date: Annotated[datetime | None, Field(title="Next Dagrun Logical Date")] = None + next_dagrun_logical_date: Annotated[datetime | None, Field(title="Next Dagrun Logical Date")] next_dagrun_data_interval_start: Annotated[ datetime | None, Field(title="Next Dagrun Data Interval Start") - ] = None - next_dagrun_data_interval_end: Annotated[ - datetime | None, Field(title="Next Dagrun Data Interval End") - ] = None - next_dagrun_run_after: Annotated[datetime | None, Field(title="Next Dagrun Run After")] = None - allowed_run_types: Annotated[list[DagRunType] | None, Field(title="Allowed Run Types")] = None + ] + next_dagrun_data_interval_end: Annotated[datetime | None, Field(title="Next Dagrun Data Interval End")] + next_dagrun_run_after: Annotated[datetime | None, Field(title="Next Dagrun Run After")] + allowed_run_types: Annotated[list[DagRunType] | None, Field(title="Allowed Run Types")] owners: Annotated[list[str], Field(title="Owners")] is_backfillable: Annotated[ bool, Field(description="Whether this Dag's schedule supports backfilling.", title="Is Backfillable") @@ -1836,26 +1831,26 @@ class DAGRunResponse(BaseModel): dag_run_id: Annotated[str, Field(title="Dag Run Id")] dag_id: Annotated[str, Field(title="Dag Id")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None - queued_at: Annotated[datetime | None, Field(title="Queued At")] = None - start_date: Annotated[datetime | None, Field(title="Start Date")] = None - end_date: Annotated[datetime | None, Field(title="End Date")] = None - duration: Annotated[float | None, Field(title="Duration")] = None - data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] = None - data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] + queued_at: Annotated[datetime | None, Field(title="Queued At")] + start_date: Annotated[datetime | None, Field(title="Start Date")] + end_date: Annotated[datetime | None, Field(title="End Date")] + duration: Annotated[float | None, Field(title="Duration")] + data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] + data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] run_after: Annotated[datetime, Field(title="Run After")] - last_scheduling_decision: Annotated[datetime | None, Field(title="Last Scheduling Decision")] = None + last_scheduling_decision: Annotated[datetime | None, Field(title="Last Scheduling Decision")] run_type: DagRunType state: DagRunState - triggered_by: DagRunTriggeredByType | None = None - triggering_user_name: Annotated[str | None, Field(title="Triggering User Name")] = None - conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None - note: Annotated[str | None, Field(title="Note")] = None + triggered_by: DagRunTriggeredByType | None + triggering_user_name: Annotated[str | None, Field(title="Triggering User Name")] + conf: Annotated[dict[str, Any] | None, Field(title="Conf")] + note: Annotated[str | None, Field(title="Note")] dag_versions: Annotated[list[DagVersionResponse], Field(title="Dag Versions")] - bundle_version: Annotated[str | None, Field(title="Bundle Version")] = None + bundle_version: Annotated[str | None, Field(title="Bundle Version")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] - partition_key: Annotated[str | None, Field(title="Partition Key")] = None - partition_date: Annotated[datetime | None, Field(title="Partition Date")] = None + partition_key: Annotated[str | None, Field(title="Partition Key")] + partition_date: Annotated[datetime | None, Field(title="Partition Date")] class DAGRunsBatchBody(BaseModel): @@ -2033,7 +2028,9 @@ class PluginResponse(BaseModel): ] react_apps: Annotated[list[ReactAppResponse], Field(title="React Apps")] appbuilder_views: Annotated[list[AppBuilderViewResponse], Field(title="Appbuilder Views")] - appbuilder_menu_items: Annotated[list[AppBuilderMenuItemResponse], Field(title="Appbuilder Menu Items")] + appbuilder_menu_items: Annotated[ + list[AppBuilderMenuItemResponse], Field(deprecated=True, title="Appbuilder Menu Items") + ] global_operator_extra_links: Annotated[list[str], Field(title="Global Operator Extra Links")] operator_extra_links: Annotated[list[str], Field(title="Operator Extra Links")] source: Annotated[str, Field(title="Source")] @@ -2085,28 +2082,28 @@ class TaskInstanceHistoryResponse(BaseModel): dag_id: Annotated[str, Field(title="Dag Id")] dag_run_id: Annotated[str, Field(title="Dag Run Id")] map_index: Annotated[int, Field(title="Map Index")] - start_date: Annotated[datetime | None, Field(title="Start Date")] = None - end_date: Annotated[datetime | None, Field(title="End Date")] = None - duration: Annotated[float | None, Field(title="Duration")] = None - state: TaskInstanceState | None = None + start_date: Annotated[datetime | None, Field(title="Start Date")] + end_date: Annotated[datetime | None, Field(title="End Date")] + duration: Annotated[float | None, Field(title="Duration")] + state: TaskInstanceState | None try_number: Annotated[int, Field(title="Try Number")] max_tries: Annotated[int, Field(title="Max Tries")] task_display_name: Annotated[str, Field(title="Task Display Name")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] - hostname: Annotated[str | None, Field(title="Hostname")] = None - unixname: Annotated[str | None, Field(title="Unixname")] = None + hostname: Annotated[str | None, Field(title="Hostname")] + unixname: Annotated[str | None, Field(title="Unixname")] pool: Annotated[str, Field(title="Pool")] pool_slots: Annotated[int, Field(title="Pool Slots")] - queue: Annotated[str | None, Field(title="Queue")] = None - priority_weight: Annotated[int | None, Field(title="Priority Weight")] = None - operator: Annotated[str | None, Field(title="Operator")] = None - operator_name: Annotated[str | None, Field(title="Operator Name")] = None - queued_when: Annotated[datetime | None, Field(title="Queued When")] = None - scheduled_when: Annotated[datetime | None, Field(title="Scheduled When")] = None - pid: Annotated[int | None, Field(title="Pid")] = None - executor: Annotated[str | None, Field(title="Executor")] = None + queue: Annotated[str | None, Field(title="Queue")] + priority_weight: Annotated[int | None, Field(title="Priority Weight")] + operator: Annotated[str | None, Field(title="Operator")] + operator_name: Annotated[str | None, Field(title="Operator Name")] + queued_when: Annotated[datetime | None, Field(title="Queued When")] + scheduled_when: Annotated[datetime | None, Field(title="Scheduled When")] + pid: Annotated[int | None, Field(title="Pid")] + executor: Annotated[str | None, Field(title="Executor")] executor_config: Annotated[str, Field(title="Executor Config")] - dag_version: DagVersionResponse | None = None + dag_version: DagVersionResponse | None class TaskInstanceResponse(BaseModel): @@ -2119,35 +2116,35 @@ class TaskInstanceResponse(BaseModel): dag_id: Annotated[str, Field(title="Dag Id")] dag_run_id: Annotated[str, Field(title="Dag Run Id")] map_index: Annotated[int, Field(title="Map Index")] - logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[datetime | None, Field(title="Logical Date")] run_after: Annotated[datetime, Field(title="Run After")] - start_date: Annotated[datetime | None, Field(title="Start Date")] = None - end_date: Annotated[datetime | None, Field(title="End Date")] = None - duration: Annotated[float | None, Field(title="Duration")] = None - state: TaskInstanceState | None = None + start_date: Annotated[datetime | None, Field(title="Start Date")] + end_date: Annotated[datetime | None, Field(title="End Date")] + duration: Annotated[float | None, Field(title="Duration")] + state: TaskInstanceState | None try_number: Annotated[int, Field(title="Try Number")] max_tries: Annotated[int, Field(title="Max Tries")] task_display_name: Annotated[str, Field(title="Task Display Name")] dag_display_name: Annotated[str, Field(title="Dag Display Name")] - hostname: Annotated[str | None, Field(title="Hostname")] = None - unixname: Annotated[str | None, Field(title="Unixname")] = None + hostname: Annotated[str | None, Field(title="Hostname")] + unixname: Annotated[str | None, Field(title="Unixname")] pool: Annotated[str, Field(title="Pool")] pool_slots: Annotated[int, Field(title="Pool Slots")] - queue: Annotated[str | None, Field(title="Queue")] = None - priority_weight: Annotated[int | None, Field(title="Priority Weight")] = None - operator: Annotated[str | None, Field(title="Operator")] = None - operator_name: Annotated[str | None, Field(title="Operator Name")] = None - queued_when: Annotated[datetime | None, Field(title="Queued When")] = None - scheduled_when: Annotated[datetime | None, Field(title="Scheduled When")] = None - pid: Annotated[int | None, Field(title="Pid")] = None - executor: Annotated[str | None, Field(title="Executor")] = None + queue: Annotated[str | None, Field(title="Queue")] + priority_weight: Annotated[int | None, Field(title="Priority Weight")] + operator: Annotated[str | None, Field(title="Operator")] + operator_name: Annotated[str | None, Field(title="Operator Name")] + queued_when: Annotated[datetime | None, Field(title="Queued When")] + scheduled_when: Annotated[datetime | None, Field(title="Scheduled When")] + pid: Annotated[int | None, Field(title="Pid")] + executor: Annotated[str | None, Field(title="Executor")] executor_config: Annotated[str, Field(title="Executor Config")] - note: Annotated[str | None, Field(title="Note")] = None - rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None + note: Annotated[str | None, Field(title="Note")] + rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] rendered_fields: Annotated[dict[str, Any] | None, Field(title="Rendered Fields")] = None - trigger: TriggerResponse | None = None - triggerer_job: JobResponse | None = None - dag_version: DagVersionResponse | None = None + trigger: TriggerResponse | None + triggerer_job: JobResponse | None + dag_version: DagVersionResponse | None class TaskResponse(BaseModel): @@ -2155,32 +2152,32 @@ class TaskResponse(BaseModel): Task serializer for responses. """ - task_id: Annotated[str | None, Field(title="Task Id")] = None - task_display_name: Annotated[str | None, Field(title="Task Display Name")] = None - owner: Annotated[str | None, Field(title="Owner")] = None - start_date: Annotated[datetime | None, Field(title="Start Date")] = None - end_date: Annotated[datetime | None, Field(title="End Date")] = None - trigger_rule: Annotated[str | None, Field(title="Trigger Rule")] = None + task_id: Annotated[str | None, Field(title="Task Id")] + task_display_name: Annotated[str | None, Field(title="Task Display Name")] + owner: Annotated[str | None, Field(title="Owner")] + start_date: Annotated[datetime | None, Field(title="Start Date")] + end_date: Annotated[datetime | None, Field(title="End Date")] + trigger_rule: Annotated[str | None, Field(title="Trigger Rule")] depends_on_past: Annotated[bool, Field(title="Depends On Past")] wait_for_downstream: Annotated[bool, Field(title="Wait For Downstream")] - retries: Annotated[float | None, Field(title="Retries")] = None - queue: Annotated[str | None, Field(title="Queue")] = None - pool: Annotated[str | None, Field(title="Pool")] = None - pool_slots: Annotated[float | None, Field(title="Pool Slots")] = None - execution_timeout: TimeDelta | None = None - retry_delay: TimeDelta | None = None + retries: Annotated[float | None, Field(title="Retries")] + queue: Annotated[str | None, Field(title="Queue")] + pool: Annotated[str | None, Field(title="Pool")] + pool_slots: Annotated[float | None, Field(title="Pool Slots")] + execution_timeout: TimeDelta | None + retry_delay: TimeDelta | None retry_exponential_backoff: Annotated[float, Field(title="Retry Exponential Backoff")] - priority_weight: Annotated[float | None, Field(title="Priority Weight")] = None - weight_rule: Annotated[str | None, Field(title="Weight Rule")] = None - ui_color: Annotated[str | None, Field(title="Ui Color")] = None - ui_fgcolor: Annotated[str | None, Field(title="Ui Fgcolor")] = None - template_fields: Annotated[list[str] | None, Field(title="Template Fields")] = None - downstream_task_ids: Annotated[list[str] | None, Field(title="Downstream Task Ids")] = None - doc_md: Annotated[str | None, Field(title="Doc Md")] = None - operator_name: Annotated[str | None, Field(title="Operator Name")] = None - params: Annotated[dict[str, Any] | None, Field(title="Params")] = None - class_ref: Annotated[dict[str, Any] | None, Field(title="Class Ref")] = None - is_mapped: Annotated[bool | None, Field(title="Is Mapped")] = None + priority_weight: Annotated[float | None, Field(title="Priority Weight")] + weight_rule: Annotated[str | None, Field(title="Weight Rule")] + ui_color: Annotated[str | None, Field(title="Ui Color")] + ui_fgcolor: Annotated[str | None, Field(title="Ui Fgcolor")] + template_fields: Annotated[list[str] | None, Field(title="Template Fields")] + downstream_task_ids: Annotated[list[str] | None, Field(title="Downstream Task Ids")] + doc_md: Annotated[str | None, Field(title="Doc Md")] + operator_name: Annotated[str | None, Field(title="Operator Name")] + params: Annotated[dict[str, Any] | None, Field(title="Params")] + class_ref: Annotated[dict[str, Any] | None, Field(title="Class Ref")] + is_mapped: Annotated[bool | None, Field(title="Is Mapped")] extra_links: Annotated[ list[str], Field(description="Extract and return extra_links.", title="Extra Links") ] diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index cf61ed5dafa72..b57883f7c4265 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -278,6 +278,11 @@ class TestAssetsOperations: ) dag_run_response = DAGRunResponse( + duration=None, + triggering_user_name=None, + bundle_version=None, + partition_key=None, + partition_date=None, dag_display_name=dag_id, dag_run_id=dag_id, dag_id=dag_id, @@ -296,6 +301,7 @@ class TestAssetsOperations: note=None, dag_versions=[ DagVersionResponse( + bundle_url=None, id=uuid.uuid4(), version_number=1, dag_id=dag_id, @@ -655,15 +661,17 @@ class TestConnectionsOperations: ) connection_response = ConnectionResponse( + description=None, + team_name=None, connection_id=connection_id, conn_type=conn_type, host=host, - schema_=schema_, + schema=schema_, login=login, password=password, port=port, extra=extra, - ) + ) # type: ignore[call-arg] connections_response = ConnectionCollectionResponse( connections=[connection_response], @@ -689,7 +697,9 @@ class TestConnectionsOperations: def test_get(self): def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path == f"/api/v2/connections/{self.connection_id}" - return httpx.Response(200, json=json.loads(self.connection_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connection_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.get(self.connection_id) @@ -698,7 +708,9 @@ def handle_request(request: httpx.Request) -> httpx.Response: def test_list(self): def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path == "/api/v2/connections" - return httpx.Response(200, json=json.loads(self.connections_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connections_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.list() @@ -707,7 +719,9 @@ def handle_request(request: httpx.Request) -> httpx.Response: def test_create(self): def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path == "/api/v2/connections" - return httpx.Response(200, json=json.loads(self.connection_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connection_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.create(connection=self.connection) @@ -729,7 +743,9 @@ def handle_request(request: httpx.Request) -> httpx.Response: "schema": self.schema_, } assert "schema_" not in request_body - return httpx.Response(200, json=json.loads(self.connection_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connection_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.create(connection=connection) @@ -775,7 +791,9 @@ def handle_request(request: httpx.Request) -> httpx.Response: def test_delete(self): def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path == f"/api/v2/connections/{self.connection_id}" - return httpx.Response(200, json=json.loads(self.connection_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connection_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.delete(self.connection_id) @@ -784,7 +802,9 @@ def handle_request(request: httpx.Request) -> httpx.Response: def test_update(self): def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path == f"/api/v2/connections/{self.connection_id}" - return httpx.Response(200, json=json.loads(self.connection_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connection_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.update(connection=self.connection) @@ -813,7 +833,9 @@ def handle_request(request: httpx.Request) -> httpx.Response: "team_name": None, } assert "schema_" not in request_body - return httpx.Response(200, json=json.loads(self.connection_response.model_dump_json())) + return httpx.Response( + 200, json=json.loads(self.connection_response.model_dump_json(by_alias=True)) + ) client = make_api_client(transport=httpx.MockTransport(handle_request)) response = client.connections.update(connection=connection) @@ -871,6 +893,9 @@ class TestDagOperations: dag_id = "dag_id" dag_display_name = "dag_display_name" dag_response = DAGResponse( + last_parse_duration=None, + bundle_version=None, + allowed_run_types=None, dag_id=dag_id, dag_display_name=dag_display_name, is_paused=False, @@ -901,6 +926,11 @@ class TestDagOperations: ) dag_details_response = DAGDetailsResponse( + last_parse_duration=None, + bundle_version=None, + allowed_run_types=None, + default_args=None, + latest_dag_version=None, dag_id=dag_id, dag_display_name="dag_display_name", is_paused=False, @@ -978,6 +1008,7 @@ class TestDagOperations: ) dag_version_response = DagVersionResponse( + bundle_url=None, id=uuid.uuid4(), version_number=1, dag_id=dag_id, @@ -1011,6 +1042,7 @@ class TestDagOperations: # DagRun related trigger_dag_run = TriggerDAGRunPostBody( + logical_date=None, conf=None, note=None, ) @@ -1018,6 +1050,11 @@ class TestDagOperations: dag_id = "dag_id" dag_run_id = "dag_run_id" dag_run_response = DAGRunResponse( + duration=None, + triggering_user_name=None, + bundle_version=None, + partition_key=None, + partition_date=None, dag_display_name=dag_run_id, dag_run_id=dag_run_id, dag_id=dag_id, @@ -1036,6 +1073,7 @@ class TestDagOperations: note=None, dag_versions=[ DagVersionResponse( + bundle_url=None, id=uuid.uuid4(), version_number=1, dag_id=dag_id, @@ -1175,6 +1213,11 @@ class TestDagRunOperations: dag_id = "dag_id" dag_run_id = "dag_run_id" dag_run_response = DAGRunResponse( + duration=None, + triggering_user_name=None, + bundle_version=None, + partition_key=None, + partition_date=None, dag_display_name=dag_run_id, dag_run_id=dag_run_id, dag_id=dag_id, @@ -1193,6 +1236,7 @@ class TestDagRunOperations: note=None, dag_versions=[ DagVersionResponse( + bundle_url=None, id=uuid.uuid4(), version_number=1, dag_id=dag_id, @@ -1387,6 +1431,7 @@ class TestPoolsOperations: ] ) pool_response = PoolResponse( + team_name=None, name=pool_name, slots=1, description="description", @@ -1456,6 +1501,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: class TestProvidersOperations: provider_response = ProviderResponse( + documentation_url=None, package_name="package_name", version="version", description="description", @@ -1487,6 +1533,7 @@ class TestVariablesOperations: } ) variable_response = VariableResponse( + team_name=None, key=key, value=value, description=description, diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_connections_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_connections_command.py index ba803ddd8e0cd..452d6182e7ddc 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_connections_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_connections_command.py @@ -51,7 +51,9 @@ class TestCliConnectionCommands: port=1234, extra="{}", description="Test connection description", - ) + schema=None, + team_name=None, + ) # type: ignore[call-arg] ], total_entries=1, ) diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py index 405eb030065f8..9dfc38082a54b 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py @@ -58,6 +58,9 @@ class TestDagCommands: file_token="file_token", bundle_name="bundle_name", is_stale=False, + last_parse_duration=None, + bundle_version=None, + allowed_run_types=None, ) dag_response_unpaused = DAGResponse( @@ -88,6 +91,9 @@ class TestDagCommands: file_token="file_token", bundle_name="bundle_name", is_stale=False, + last_parse_duration=None, + bundle_version=None, + allowed_run_types=None, ) dag_response_no_schedule = DAGResponse( @@ -118,6 +124,9 @@ class TestDagCommands: file_token="file_token", bundle_name="bundle_name", is_stale=False, + last_parse_duration=None, + bundle_version=None, + allowed_run_types=None, ) def test_pause_dag(self, api_client_maker, monkeypatch): diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py index f573585935fd4..3e7b2714f940f 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py @@ -41,10 +41,7 @@ class TestCliVariableCommands: variable_collection_response = VariableCollectionResponse( variables=[ VariableResponse( - key=key, - value=value, - description=description, - is_encrypted=False, + key=key, value=value, description=description, is_encrypted=False, team_name=None ), ], total_entries=1, diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index 413723ec31745..a27f7ee78463c 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2818,6 +2818,9 @@ def _create_task_instance( "run_after": run_after, # type: ignore "conf": conf, "consumed_asset_events": [], + # Nullable-but-required in the generated schema, so they must be + # passed explicitly; guarded for older Task SDKs that lack them. + **{f: None for f in ("end_date", "partition_key") if f in DagRun.model_fields}, **({"state": DagRunState.RUNNING} if "state" in DagRun.model_fields else {}), } ), diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py index e0984f88e24c2..be68d99af990d 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py @@ -467,6 +467,12 @@ def test_user_agent_dag_run_key_is_hashed_correctly( mock_supervisor_comms.send.return_value = ConnectionResult( conn_id="aws_default", conn_type="aws", + host=None, + schema=None, + login=None, + password=None, + port=None, + extra=None, ) with mock.patch.dict(os.environ, env_var, clear=True): dag_run_key = self.fetch_tags()["DagRunKey"] diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index edec29202c0db..b44e90b4958d4 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -1289,9 +1289,17 @@ def _create_listener_and_task_instance( "run_type": DagRunType.MANUAL, "run_after": timezone.datetime(2023, 1, 3, 13, 1, 1), "consumed_asset_events": [], - **( - {"state": SdkDagRunState.RUNNING} if "state" in SdkDagRun.model_fields else {} - ), + # Nullable-but-required on newer SDKs, absent on older ones. + **{ + field: value + for field, value in ( + ("state", SdkDagRunState.RUNNING), + ("data_interval_start", None), + ("data_interval_end", None), + ("partition_key", None), + ) + if field in SdkDagRun.model_fields + }, } ), task_reschedule_count=0, diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index caecd570700ec..2810a0f836cff 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -215,7 +215,7 @@ exclude_also = [ [dependency-groups] codegen = [ - "datamodel-code-generator[http]==0.33.0", + "datamodel-code-generator[http]>=0.71.0", "openapi-spec-validator>=0.7.1", "svcs>=25.1.0", "rich>=13.6.0", diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 80cee82cb1eaa..c625ba92a7249 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: http://0.0.0.0:8080/execution/openapi.json -# version: 0.33.0 +# version: 0.71.0 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file @@ -70,12 +70,12 @@ class ConnectionResponse(BaseModel): conn_id: Annotated[str, Field(title="Conn Id")] conn_type: Annotated[str, Field(title="Conn Type")] - host: Annotated[str | None, Field(title="Host")] = None - schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None - login: Annotated[str | None, Field(title="Login")] = None - password: Annotated[str | None, Field(title="Password")] = None - port: Annotated[int | None, Field(title="Port")] = None - extra: Annotated[str | None, Field(title="Extra")] = None + host: Annotated[str | None, Field(title="Host")] + schema_: Annotated[str | None, Field(alias="schema", title="Schema")] + login: Annotated[str | None, Field(title="Login")] + password: Annotated[str | None, Field(title="Password")] + port: Annotated[int | None, Field(title="Port")] + extra: Annotated[str | None, Field(title="Extra")] class ConnectionTestConnectionResponse(BaseModel): @@ -116,12 +116,12 @@ class DagResponse(BaseModel): dag_id: Annotated[str, Field(title="Dag Id")] is_paused: Annotated[bool, Field(title="Is Paused")] - bundle_name: Annotated[str | None, Field(title="Bundle Name")] = None - bundle_version: Annotated[str | None, Field(title="Bundle Version")] = None - relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] = None - owners: Annotated[str | None, Field(title="Owners")] = None + bundle_name: Annotated[str | None, Field(title="Bundle Name")] + bundle_version: Annotated[str | None, Field(title="Bundle Version")] + relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] + owners: Annotated[str | None, Field(title="Owners")] tags: Annotated[list[str], Field(title="Tags")] - next_dagrun: Annotated[AwareDatetime | None, Field(title="Next Dagrun")] = None + next_dagrun: Annotated[AwareDatetime | None, Field(title="Next Dagrun")] class DagRunAssetReference(BaseModel): @@ -134,13 +134,13 @@ class DagRunAssetReference(BaseModel): ) run_id: Annotated[str, Field(title="Run Id")] dag_id: Annotated[str, Field(title="Dag Id")] - logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None + logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] start_date: Annotated[AwareDatetime, Field(title="Start Date")] - end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None + end_date: Annotated[AwareDatetime | None, Field(title="End Date")] state: Annotated[str, Field(title="State")] - data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None - data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None - partition_key: Annotated[str | None, Field(title="Partition Key")] = None + data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] + data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] + partition_key: Annotated[str | None, Field(title="Partition Key")] class DagRunState(str, Enum): @@ -254,10 +254,10 @@ class TIAwaitingInputStatePayload(BaseModel): model_config = ConfigDict( extra="forbid", ) - state: Annotated[Literal["awaiting_input"] | None, Field(title="State")] = "awaiting_input" + state: Annotated[Literal["awaiting_input"], Field(title="State")] = "awaiting_input" timeout: Annotated[timedelta | None, Field(title="Timeout")] = None next_method: Annotated[str, Field(title="Next Method")] - next_kwargs: Annotated[dict[str, JsonValue] | None, Field(title="Next Kwargs")] = None + next_kwargs: Annotated[dict[str, JsonValue | None] | None, Field(title="Next Kwargs")] = None rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None @@ -269,13 +269,13 @@ class TIDeferredStatePayload(BaseModel): model_config = ConfigDict( extra="forbid", ) - state: Annotated[Literal["deferred"] | None, Field(title="State")] = "deferred" + state: Annotated[Literal["deferred"], Field(title="State")] = "deferred" classpath: Annotated[str, Field(title="Classpath")] - trigger_kwargs: Annotated[dict[str, JsonValue] | str | None, Field(title="Trigger Kwargs")] = None + trigger_kwargs: Annotated[dict[str, JsonValue | None] | str | None, Field(title="Trigger Kwargs")] = None trigger_timeout: Annotated[timedelta | None, Field(title="Trigger Timeout")] = None queue: Annotated[str | None, Field(title="Queue")] = None next_method: Annotated[str, Field(title="Next Method")] - next_kwargs: Annotated[dict[str, JsonValue] | None, Field(title="Next Kwargs")] = None + next_kwargs: Annotated[dict[str, JsonValue | None] | None, Field(title="Next Kwargs")] = None rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None @@ -287,7 +287,7 @@ class TIEnterRunningPayload(BaseModel): model_config = ConfigDict( extra="forbid", ) - state: Annotated[Literal["running"] | None, Field(title="State")] = "running" + state: Annotated[Literal["running"], Field(title="State")] = "running" hostname: Annotated[str, Field(title="Hostname")] unixname: Annotated[str, Field(title="Unixname")] pid: Annotated[int, Field(title="Pid")] @@ -314,7 +314,7 @@ class TIRescheduleStatePayload(BaseModel): model_config = ConfigDict( extra="forbid", ) - state: Annotated[Literal["up_for_reschedule"] | None, Field(title="State")] = "up_for_reschedule" + state: Annotated[Literal["up_for_reschedule"], Field(title="State")] = "up_for_reschedule" reschedule_date: Annotated[AwareDatetime, Field(title="Reschedule Date")] end_date: Annotated[AwareDatetime, Field(title="End Date")] @@ -327,7 +327,7 @@ class TIRetryStatePayload(BaseModel): model_config = ConfigDict( extra="forbid", ) - state: Annotated[Literal["up_for_retry"] | None, Field(title="State")] = "up_for_retry" + state: Annotated[Literal["up_for_retry"], Field(title="State")] = "up_for_retry" end_date: Annotated[AwareDatetime, Field(title="End Date")] rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None retry_delay_seconds: Annotated[float | None, Field(title="Retry Delay Seconds")] = None @@ -353,7 +353,7 @@ class TISuccessStatePayload(BaseModel): model_config = ConfigDict( extra="forbid", ) - state: Annotated[Literal["success"] | None, Field(title="State")] = "success" + state: Annotated[Literal["success"], Field(title="State")] = "success" end_date: Annotated[AwareDatetime, Field(title="End Date")] task_outlets: Annotated[list[AssetProfile] | None, Field(title="Task Outlets")] = None outlet_events: Annotated[list[dict[str, Any]] | None, Field(title="Outlet Events")] = None @@ -409,7 +409,7 @@ class TaskStateStorePutBody(BaseModel): model_config = ConfigDict( extra="forbid", ) - value: JsonValue + value: JsonValue | None expires_at: Annotated[AwareDatetime | None, Field(title="Expires At")] = None @@ -421,7 +421,7 @@ class TaskStateStoreResponse(BaseModel): model_config = ConfigDict( extra="forbid", ) - value: JsonValue + value: JsonValue | None class TaskStatesResponse(BaseModel): @@ -497,7 +497,7 @@ class VariablePostBody(BaseModel): model_config = ConfigDict( extra="forbid", ) - val: Annotated[str | None, Field(title="Val")] = None + val: Annotated[str | None, Field(title="Val")] description: Annotated[str | None, Field(title="Description")] = None @@ -510,7 +510,7 @@ class VariableResponse(BaseModel): extra="forbid", ) key: Annotated[str, Field(title="Key")] - value: Annotated[str | None, Field(title="Value")] = None + value: Annotated[str | None, Field(title="Value")] class XComResponse(BaseModel): @@ -519,12 +519,12 @@ class XComResponse(BaseModel): """ key: Annotated[str, Field(title="Key")] - value: JsonValue + value: JsonValue | None -class XComSequenceIndexResponse(RootModel[JsonValue]): +class XComSequenceIndexResponse(RootModel[JsonValue | None]): root: Annotated[ - JsonValue, + JsonValue | None, Field( description="XCom schema with minimal structure for index-based access.", title="XComSequenceIndexResponse", @@ -532,13 +532,13 @@ class XComSequenceIndexResponse(RootModel[JsonValue]): ] -class XComSequenceSliceResponse(RootModel[list[JsonValue]]): +class XComSequenceSliceResponse(RootModel[list[JsonValue | None]]): """ XCom schema with minimal structure for slice-based access. """ root: Annotated[ - list[JsonValue], + list[JsonValue | None], Field( description="XCom schema with minimal structure for slice-based access.", title="XComSequenceSliceResponse", @@ -617,7 +617,7 @@ class AssetReferenceAssetEventDagRun(BaseModel): ) name: Annotated[str, Field(title="Name")] uri: Annotated[str, Field(title="Uri")] - extra: Annotated[dict[str, JsonValue], Field(title="Extra")] + extra: Annotated[dict[str, JsonValue | None], Field(title="Extra")] class AssetResponse(BaseModel): @@ -628,7 +628,7 @@ class AssetResponse(BaseModel): name: Annotated[str, Field(title="Name")] uri: Annotated[str, Field(title="Uri")] group: Annotated[str, Field(title="Group")] - extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None + extra: Annotated[dict[str, JsonValue | None] | None, Field(title="Extra")] = None class AssetStateStorePutBody(BaseModel): @@ -639,7 +639,7 @@ class AssetStateStorePutBody(BaseModel): model_config = ConfigDict( extra="forbid", ) - value: JsonValue + value: JsonValue | None class AssetStateStoreResponse(BaseModel): @@ -650,7 +650,7 @@ class AssetStateStoreResponse(BaseModel): model_config = ConfigDict( extra="forbid", ) - value: JsonValue + value: JsonValue | None class ConnectionTestResultBody(BaseModel): @@ -687,8 +687,8 @@ class HITLDetailResponse(BaseModel): response_received: Annotated[bool, Field(title="Response Received")] responded_by_user: HITLUser | None = None - responded_at: Annotated[AwareDatetime | None, Field(title="Responded At")] = None - chosen_options: Annotated[list[str] | None, Field(title="Chosen Options")] = None + responded_at: Annotated[AwareDatetime | None, Field(title="Responded At")] + chosen_options: Annotated[list[str] | None, Field(title="Chosen Options")] params_input: Annotated[dict[str, Any] | None, Field(title="Params Input")] = None @@ -718,11 +718,11 @@ class AssetEventDagRunReference(BaseModel): extra="forbid", ) asset: AssetReferenceAssetEventDagRun - extra: Annotated[dict[str, JsonValue], Field(title="Extra")] - source_task_id: Annotated[str | None, Field(title="Source Task Id")] = None - source_dag_id: Annotated[str | None, Field(title="Source Dag Id")] = None - source_run_id: Annotated[str | None, Field(title="Source Run Id")] = None - source_map_index: Annotated[int | None, Field(title="Source Map Index")] = None + extra: Annotated[dict[str, JsonValue | None], Field(title="Extra")] + source_task_id: Annotated[str | None, Field(title="Source Task Id")] + source_dag_id: Annotated[str | None, Field(title="Source Dag Id")] + source_run_id: Annotated[str | None, Field(title="Source Run Id")] + source_map_index: Annotated[int | None, Field(title="Source Map Index")] source_aliases: Annotated[list[AssetAliasReferenceAssetEventDagRun], Field(title="Source Aliases")] timestamp: Annotated[AwareDatetime, Field(title="Timestamp")] partition_key: Annotated[str | None, Field(title="Partition Key")] = None @@ -735,7 +735,7 @@ class AssetEventResponse(BaseModel): id: Annotated[int, Field(title="Id")] timestamp: Annotated[AwareDatetime, Field(title="Timestamp")] - extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None + extra: Annotated[dict[str, JsonValue | None] | None, Field(title="Extra")] = None asset: AssetResponse created_dagruns: Annotated[list[DagRunAssetReference], Field(title="Created Dagruns")] source_task_id: Annotated[str | None, Field(title="Source Task Id")] = None @@ -763,19 +763,19 @@ class DagRun(BaseModel): ) dag_id: Annotated[str, Field(title="Dag Id")] run_id: Annotated[str, Field(title="Run Id")] - logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None - data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None - data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None + logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] + data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] + data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] run_after: Annotated[AwareDatetime, Field(title="Run After")] - start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None - end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None + start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] + end_date: Annotated[AwareDatetime | None, Field(title="End Date")] clear_number: Annotated[int | None, Field(title="Clear Number")] = 0 run_type: DagRunType state: DagRunState conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None triggering_user_name: Annotated[str | None, Field(title="Triggering User Name")] = None consumed_asset_events: Annotated[list[AssetEventDagRunReference], Field(title="Consumed Asset Events")] - partition_key: Annotated[str | None, Field(title="Partition Key")] = None + partition_key: Annotated[str | None, Field(title="Partition Key")] partition_date: Annotated[AwareDatetime | None, Field(title="Partition Date")] = None note: Annotated[str | None, Field(title="Note")] = None team_name: Annotated[str | None, Field(title="Team Name")] = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 08383b1422204..53efa708750dd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -34,7 +34,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -187,34 +194,6 @@ "title": "AssetProfile", "type": "object" }, - "AssetReferenceAssetEventDagRun": { - "additionalProperties": false, - "description": "Schema for AssetModel used in AssetEventDagRunReference.", - "properties": { - "name": { - "title": "Name", - "type": "string" - }, - "uri": { - "title": "Uri", - "type": "string" - }, - "extra": { - "additionalProperties": { - "$ref": "#/$defs/JsonValue" - }, - "title": "Extra", - "type": "object" - } - }, - "required": [ - "name", - "uri", - "extra" - ], - "title": "AssetReferenceAssetEventDagRun", - "type": "object" - }, "AssetResponse": { "description": "Asset schema for responses with fields that are needed for Runtime.", "properties": { @@ -234,7 +213,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -273,7 +259,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -304,7 +297,14 @@ "description": "Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "AssetStateStoreResult", @@ -347,17 +347,10 @@ "description": "Park a task instance awaiting human input (Human-in-the-loop), without a trigger.", "properties": { "state": { - "anyOf": [ - { - "const": "awaiting_input", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "awaiting_input", "default": "awaiting_input", - "title": "State" + "title": "State", + "type": "string" }, "timeout": { "anyOf": [ @@ -380,7 +373,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -500,6 +500,97 @@ "title": "ClearTaskStateStore", "type": "object" }, + "ConnectionResponse": { + "description": "Connection schema for responses with fields that are needed for Runtime.", + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Host" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Schema" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Port" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Extra" + } + }, + "required": [ + "conn_id", + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" + ], + "title": "ConnectionResponse", + "type": "object" + }, "ConnectionResult": { "properties": { "conn_id": { @@ -519,7 +610,6 @@ "type": "null" } ], - "default": null, "title": "Host" }, "schema": { @@ -531,7 +621,6 @@ "type": "null" } ], - "default": null, "title": "Schema" }, "login": { @@ -543,7 +632,6 @@ "type": "null" } ], - "default": null, "title": "Login" }, "password": { @@ -555,7 +643,6 @@ "type": "null" } ], - "default": null, "title": "Password" }, "port": { @@ -567,7 +654,6 @@ "type": "null" } ], - "default": null, "title": "Port" }, "extra": { @@ -579,7 +665,6 @@ "type": "null" } ], - "default": null, "title": "Extra" }, "type": { @@ -591,7 +676,13 @@ }, "required": [ "conn_id", - "conn_type" + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" ], "title": "ConnectionResult", "type": "object" @@ -931,7 +1022,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Name" }, "bundle_version": { @@ -943,7 +1033,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Version" }, "relative_fileloc": { @@ -955,7 +1044,6 @@ "type": "null" } ], - "default": null, "title": "Relative Fileloc" }, "owners": { @@ -967,7 +1055,6 @@ "type": "null" } ], - "default": null, "title": "Owners" }, "tags": { @@ -987,7 +1074,6 @@ "type": "null" } ], - "default": null, "title": "Next Dagrun" }, "type": { @@ -1000,7 +1086,12 @@ "required": [ "dag_id", "is_paused", - "tags" + "bundle_name", + "bundle_version", + "relative_fileloc", + "owners", + "tags", + "next_dagrun" ], "title": "DagResult", "type": "object" @@ -1027,7 +1118,6 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "start_date": { @@ -1045,7 +1135,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "state": { @@ -1062,7 +1151,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1075,7 +1163,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "partition_key": { @@ -1087,15 +1174,19 @@ "type": "null" } ], - "default": null, "title": "Partition Key" } }, "required": [ "run_id", "dag_id", + "logical_date", "start_date", - "state" + "end_date", + "state", + "data_interval_start", + "data_interval_end", + "partition_key" ], "title": "DagRunAssetReference", "type": "object" @@ -1150,7 +1241,6 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "data_interval_start": { @@ -1163,7 +1253,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1176,7 +1265,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "run_after": { @@ -1194,7 +1282,6 @@ "type": "null" } ], - "default": null, "title": "Start Date" }, "end_date": { @@ -1207,7 +1294,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "clear_number": { @@ -1269,7 +1355,6 @@ "type": "null" } ], - "default": null, "title": "Partition Key" }, "partition_date": { @@ -1319,10 +1404,16 @@ "required": [ "dag_id", "run_id", + "logical_date", + "data_interval_start", + "data_interval_end", "run_after", + "start_date", + "end_date", "run_type", "state", - "consumed_asset_events" + "consumed_asset_events", + "partition_key" ], "title": "DagRunResult", "type": "object" @@ -1374,17 +1465,10 @@ "description": "Update a task instance state to deferred.", "properties": { "state": { - "anyOf": [ - { - "const": "deferred", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "deferred", "default": "deferred", - "title": "State" + "title": "State", + "type": "string" }, "classpath": { "title": "Classpath", @@ -1394,7 +1478,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -1441,7 +1532,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -3271,17 +3369,10 @@ "description": "Update a task instance state to reschedule/up_for_reschedule.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_reschedule", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_reschedule", "default": "up_for_reschedule", - "title": "State" + "title": "State", + "type": "string" }, "reschedule_date": { "format": "date-time", @@ -3324,17 +3415,10 @@ "description": "Update a task instance state to up_for_retry.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_retry", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_retry", "default": "up_for_retry", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -3708,17 +3792,10 @@ "description": "Update a task's state to success. Includes task_outlets and outlet_events for registering asset events.", "properties": { "state": { - "anyOf": [ - { - "const": "success", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "success", "default": "success", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -4002,7 +4079,14 @@ "description": "Response to GetTaskStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "TaskStateStoreResult", @@ -4227,6 +4311,33 @@ "title": "VariableKeysResult", "type": "object" }, + "VariableResponse": { + "additionalProperties": false, + "description": "Variable schema for responses with fields that are needed for Runtime.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Value" + } + }, + "required": [ + "key", + "value" + ], + "title": "VariableResponse", + "type": "object" + }, "VariableResult": { "additionalProperties": false, "properties": { @@ -4243,7 +4354,6 @@ "type": "null" } ], - "default": null, "title": "Value" }, "type": { @@ -4254,7 +4364,8 @@ } }, "required": [ - "key" + "key", + "value" ], "title": "VariableResult", "type": "object" @@ -4286,7 +4397,14 @@ "type": "string" }, "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "XComResult", @@ -4342,97 +4460,6 @@ "title": "XComSequenceSliceResult", "type": "object" }, - "ConnectionResponse": { - "description": "Connection schema for responses with fields that are needed for Runtime.", - "properties": { - "conn_id": { - "title": "Conn Id", - "type": "string" - }, - "conn_type": { - "title": "Conn Type", - "type": "string" - }, - "host": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Host" - }, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Schema" - }, - "login": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Login" - }, - "password": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Password" - }, - "port": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "title": "Port" - }, - "extra": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Extra" - } - }, - "required": [ - "conn_id", - "conn_type", - "host", - "schema", - "login", - "password", - "port", - "extra" - ], - "title": "ConnectionResponse", - "type": "object" - }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4529,6 +4556,34 @@ "title": "AssetEventDagRunReference", "type": "object" }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, "DagRun": { "additionalProperties": false, "description": "Schema for DagRun model with minimal required fields needed for Runtime.", @@ -4880,33 +4935,6 @@ ], "title": "TaskInstance", "type": "object" - }, - "VariableResponse": { - "additionalProperties": false, - "description": "Variable schema for responses with fields that are needed for Runtime.", - "properties": { - "key": { - "title": "Key", - "type": "string" - }, - "value": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Value" - } - }, - "required": [ - "key", - "value" - ], - "title": "VariableResponse", - "type": "object" } } } diff --git a/task-sdk/tests/conftest.py b/task-sdk/tests/conftest.py index c1ef3b72c92f4..eed48a1554ca1 100644 --- a/task-sdk/tests/conftest.py +++ b/task-sdk/tests/conftest.py @@ -290,6 +290,8 @@ def _make_context( state=DagRunState.RUNNING, conf=conf, # type: ignore consumed_asset_events=list(consumed_asset_events), + end_date=None, + partition_key=None, ), task_reschedule_count=task_reschedule_count, max_tries=max_tries, diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 5c4bac5697268..a7f69e948f1a5 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1127,6 +1127,12 @@ def handle_request(request: httpx.Request) -> httpx.Response: json={ "conn_id": "test_conn", "conn_type": "mysql", + "host": None, + "schema": None, + "login": None, + "password": None, + "port": None, + "extra": None, }, ) return httpx.Response(status_code=400, json={"detail": "Bad Request"}) @@ -1480,6 +1486,10 @@ def handle_request(request: httpx.Request) -> httpx.Response: "run_type": "scheduled", "state": "success", "consumed_asset_events": [], + "data_interval_start": None, + "data_interval_end": None, + "end_date": None, + "partition_key": None, }, ) return httpx.Response(status_code=422) @@ -1515,6 +1525,10 @@ def handle_request(request: httpx.Request) -> httpx.Response: "run_type": "scheduled", "state": "success", "consumed_asset_events": [], + "data_interval_start": None, + "data_interval_end": None, + "end_date": None, + "partition_key": None, }, ) return httpx.Response(status_code=422) diff --git a/task-sdk/tests/task_sdk/bases/test_hook.py b/task-sdk/tests/task_sdk/bases/test_hook.py index 351f63e039dd6..b15a957cd5c24 100644 --- a/task-sdk/tests/task_sdk/bases/test_hook.py +++ b/task-sdk/tests/task_sdk/bases/test_hook.py @@ -110,6 +110,8 @@ async def test_aget_hook(self, mock_supervisor_comms): login="user", password="password", port=443, + schema=None, + extra=None, ) mock_supervisor_comms.asend.return_value = conn diff --git a/task-sdk/tests/task_sdk/definitions/test_connection.py b/task-sdk/tests/task_sdk/definitions/test_connection.py index 5746b1b14f75c..c973fc58c91a4 100644 --- a/task-sdk/tests/task_sdk/definitions/test_connection.py +++ b/task-sdk/tests/task_sdk/definitions/test_connection.py @@ -104,7 +104,16 @@ def test_get_uri(self): assert parsed_uri.path.lstrip("/") == "test_schema" def test_conn_get(self, mock_supervisor_comms): - conn_result = ConnectionResult(conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306) + conn_result = ConnectionResult( + conn_id="mysql_conn", + conn_type="mysql", + host="mysql", + port=3306, + schema=None, + login=None, + password=None, + extra=None, + ) mock_supervisor_comms.send.return_value = conn_result conn = Connection.get(conn_id="mysql_conn") diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py index c33299b313c81..25a334b73a16e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py @@ -183,7 +183,16 @@ class RequestCase: client_mock=ClientMock( method_path="connections.get", args=("test_conn",), - response=ConnectionResult(conn_id="test_conn", conn_type="mysql"), + response=ConnectionResult( + conn_id="test_conn", + conn_type="mysql", + host=None, + schema=None, + login=None, + password=None, + port=None, + extra=None, + ), # type: ignore[call-arg] ), ), RequestCase( @@ -192,7 +201,16 @@ class RequestCase: client_mock=ClientMock( method_path="connections.get", args=("test_conn",), - response=ConnectionResult(conn_id="test_conn", conn_type="mysql", password="secret"), + response=ConnectionResult( + conn_id="test_conn", + conn_type="mysql", + password="secret", + host=None, + schema=None, + login=None, + port=None, + extra=None, + ), # type: ignore[call-arg] ), mask_secret_args=("secret",), ), diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py b/task-sdk/tests/task_sdk/execution_time/test_comms.py index aa7264b87502d..05dd682d8901a 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_comms.py +++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py @@ -96,6 +96,7 @@ def test_recv_StartupDetails(self, socket_pair): "state": "success", "conf": None, "consumed_asset_events": [], + "partition_key": None, }, "max_tries": 0, "should_retry": False, diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index d625c91c91735..e34008ef9c7bd 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -246,7 +246,16 @@ def test_getattr_connection(self, mock_supervisor_comms): accessor = ConnectionAccessor() # Conn from the supervisor / API Server - conn_result = ConnectionResult(conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306) + conn_result = ConnectionResult( + conn_id="mysql_conn", + conn_type="mysql", + host="mysql", + port=3306, + schema=None, + login=None, + password=None, + extra=None, + ) mock_supervisor_comms.send.return_value = conn_result @@ -259,7 +268,16 @@ def test_getattr_connection(self, mock_supervisor_comms): def test_get_method_valid_connection(self, mock_supervisor_comms): """Test that the get method returns the requested connection using `conn.get`.""" accessor = ConnectionAccessor() - conn_result = ConnectionResult(conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306) + conn_result = ConnectionResult( + conn_id="mysql_conn", + conn_type="mysql", + host="mysql", + port=3306, + schema=None, + login=None, + password=None, + extra=None, + ) mock_supervisor_comms.send.return_value = conn_result @@ -289,6 +307,9 @@ def test_getattr_connection_for_extra_dejson(self, mock_supervisor_comms): host="mysql", port=3306, extra='{"extra_key": "extra_value"}', + schema=None, + login=None, + password=None, ) mock_supervisor_comms.send.return_value = conn_result @@ -306,7 +327,14 @@ def test_getattr_connection_for_extra_dejson_decode_error(self, mock_log, mock_s # Conn from the supervisor / API Server conn_result = ConnectionResult( - conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306, extra="This is not JSON!" + conn_id="mysql_conn", + conn_type="mysql", + host="mysql", + port=3306, + extra="This is not JSON!", + schema=None, + login=None, + password=None, ) mock_supervisor_comms.send.return_value = conn_result @@ -1070,6 +1098,11 @@ def test_source_task_instance_xcom_pull(self, sample_inlet_evnets_accessor, mock run_type="scheduled", state="success", consumed_asset_events=[], + logical_date=None, + data_interval_start=None, + data_interval_end=None, + end_date=None, + partition_key=None, ) mock_supervisor_comms.reset_mock() mock_supervisor_comms.send.side_effect = [dag_run_result] @@ -1115,6 +1148,8 @@ def test_dag_run_model_accepts_null_start_date(self): state="queued", conf=None, consumed_asset_events=[], + end_date=None, + partition_key=None, ) assert dag_run.start_date is None @@ -1193,6 +1228,10 @@ def test_get_connection_uses_backend_chain(self, mock_supervisor_comms): conn_type="http", host="example.com", port=443, + schema=None, + login=None, + password=None, + extra=None, ) conn_result = ConnectionResult.from_conn_response(conn_response) mock_supervisor_comms.send.return_value = conn_result @@ -1225,6 +1264,11 @@ def get_connection(self, conn_id): conn_id="test_conn", conn_type="postgres", host="db.example.com", + schema=None, + login=None, + password=None, + port=None, + extra=None, ) conn_result = ConnectionResult.from_conn_response(conn_response) mock_supervisor_comms.send.return_value = conn_result diff --git a/task-sdk/tests/task_sdk/execution_time/test_context_cache.py b/task-sdk/tests/task_sdk/execution_time/test_context_cache.py index aa3b20ca9ce59..58c9c51e15391 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context_cache.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context_cache.py @@ -97,6 +97,8 @@ def test_get_connection_from_api(self, mock_ensure_backends, mock_supervisor_com port=3306, login="user", password="pass", + schema=None, + extra=None, ) mock_ensure_backends.return_value = [ExecutionAPISecretsBackend()] @@ -281,6 +283,10 @@ async def test_async_get_connection_from_api(self, mock_supervisor_comms): conn_type="mysql", host="host", port=3306, + schema=None, + login=None, + password=None, + extra=None, ) # Configure asend to return the conn_result when awaited @@ -344,7 +350,16 @@ def teardown_method(): def test_get_connection_no_cache_when_disabled(self, mock_ensure_backends, mock_supervisor_comms): """Test that cache is not used when disabled.""" conn_id = "test_conn" - conn_result = ConnectionResult(conn_id=conn_id, conn_type="mysql", host="host") + conn_result = ConnectionResult( + conn_id=conn_id, + conn_type="mysql", + host="host", + schema=None, + login=None, + password=None, + port=None, + extra=None, + ) mock_ensure_backends.return_value = [ExecutionAPISecretsBackend()] diff --git a/task-sdk/tests/task_sdk/execution_time/test_secrets.py b/task-sdk/tests/task_sdk/execution_time/test_secrets.py index 5c9cb7a22f539..b97a85f369329 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_secrets.py +++ b/task-sdk/tests/task_sdk/execution_time/test_secrets.py @@ -36,6 +36,9 @@ def test_get_connection_via_supervisor_comms(self, mock_supervisor_comms): host="example.com", port=443, schema="https", + login=None, + password=None, + extra=None, ) conn_result = ConnectionResult.from_conn_response(conn_response) mock_supervisor_comms.send.return_value = conn_result diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 81b60c6a7d8a3..5928b4582e5ef 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1600,9 +1600,28 @@ class RequestTestCase: client_mock=ClientMock( method_path="connections.get", args=("test_conn",), - response=ConnectionResult(conn_id="test_conn", conn_type="mysql"), + response=ConnectionResult( + conn_id="test_conn", + conn_type="mysql", + host=None, + schema=None, + login=None, + password=None, + port=None, + extra=None, + ), # type: ignore[call-arg] ), - expected_body={"conn_id": "test_conn", "conn_type": "mysql", "type": "ConnectionResult"}, + expected_body={ + "conn_id": "test_conn", + "conn_type": "mysql", + "host": None, + "schema": None, + "login": None, + "password": None, + "port": None, + "extra": None, + "type": "ConnectionResult", + }, ), RequestTestCase( message=GetConnection(conn_id="test_conn"), @@ -1610,12 +1629,26 @@ class RequestTestCase: client_mock=ClientMock( method_path="connections.get", args=("test_conn",), - response=ConnectionResult(conn_id="test_conn", conn_type="mysql", password="password"), + response=ConnectionResult( + conn_id="test_conn", + conn_type="mysql", + password="password", + host=None, + schema=None, + login=None, + port=None, + extra=None, + ), # type: ignore[call-arg] ), expected_body={ "conn_id": "test_conn", "conn_type": "mysql", "password": "password", + "host": None, + "schema": None, + "login": None, + "port": None, + "extra": None, "type": "ConnectionResult", }, mask_secret_args=("password",), @@ -1626,12 +1659,26 @@ class RequestTestCase: client_mock=ClientMock( method_path="connections.get", args=("test_conn",), - response=ConnectionResult(conn_id="test_conn", conn_type="mysql", schema="mysql"), # type: ignore[call-arg] + response=ConnectionResult( + conn_id="test_conn", + conn_type="mysql", + schema="mysql", + host=None, + login=None, + password=None, + port=None, + extra=None, + ), # type: ignore[call-arg] ), expected_body={ "conn_id": "test_conn", "conn_type": "mysql", "schema": "mysql", + "host": None, + "login": None, + "password": None, + "port": None, + "extra": None, "type": "ConnectionResult", }, ), @@ -2426,6 +2473,10 @@ class RequestTestCase: consumed_asset_events=[], state=DagRunState.SUCCESS, triggering_user_name=None, + data_interval_start=None, + data_interval_end=None, + end_date=None, + partition_key=None, ), ), test_id="get_dag_run", @@ -2486,6 +2537,10 @@ class RequestTestCase: consumed_asset_events=[], state=DagRunState.SUCCESS, triggering_user_name=None, + data_interval_start=None, + data_interval_end=None, + end_date=None, + partition_key=None, ) ), ), @@ -3525,6 +3580,12 @@ def handle_request(request: httpx.Request) -> httpx.Response: # Minimal enough to pass validation, we don't care what fields are in here for the tests "conn_id": remote_conn, "conn_type": "aws", + "host": None, + "schema": None, + "login": None, + "password": None, + "port": None, + "extra": None, }, ) @@ -3862,7 +3923,7 @@ def test_fetch_remote_logging_conn_does_not_cache_none_result(mocker): conn_id=conn_id, conn_type="example", host=None, - schema_=None, + schema=None, login=None, password=None, port=None, @@ -3934,7 +3995,14 @@ def subprocess_main(): subprocess.check_call([sys.executable, "-c", dedent(script)]) client_with_ti_start.connections.get.return_value = ConnectionResult( - conn_id="test_conn", conn_type="mysql", login="a", password="password1" + conn_id="test_conn", + conn_type="mysql", + login="a", + password="password1", + host=None, + schema=None, + port=None, + extra=None, ) proc = ActivitySubprocess.start( dag_rel_path=os.devnull, diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 9714a49e87214..5a96450000d8e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2170,6 +2170,7 @@ def test_template_with_connection( password="passwordvalue", schema="schemavalues", extra='{"extra__asana__workspace": "extra1"}', + port=None, ) mock_supervisor_comms.send.return_value = conn @@ -2869,6 +2870,10 @@ def test_get_previous_dagrun_basic(self, create_runtime_ti, mock_supervisor_comm run_type="scheduled", state="success", consumed_asset_events=[], + data_interval_start=None, + data_interval_end=None, + end_date=None, + partition_key=None, ) mock_supervisor_comms.send.return_value = PreviousDagRunResult(dag_run=dag_run_data) @@ -2898,6 +2903,10 @@ def test_get_previous_dagrun_with_state(self, create_runtime_ti, mock_supervisor run_type="scheduled", state="success", consumed_asset_events=[], + data_interval_start=None, + data_interval_end=None, + end_date=None, + partition_key=None, ) mock_supervisor_comms.send.return_value = PreviousDagRunResult(dag_run=dag_run_data) diff --git a/uv.lock b/uv.lock index 317f7368d02a7..6f593f9fb857a 100644 --- a/uv.lock +++ b/uv.lock @@ -2213,7 +2213,7 @@ provides-extras = ["dev"] [package.metadata.requires-dev] codegen = [ { name = "apache-airflow-devel-common", editable = "devel-common" }, - { name = "datamodel-code-generator", extras = ["http"], specifier = "==0.33.0" }, + { name = "datamodel-code-generator", extras = ["http"], specifier = ">=0.71.0" }, ] dev = [ { name = "apache-airflow-ctl", extras = ["dev"], editable = "airflow-ctl" }, @@ -8934,7 +8934,7 @@ provides-extras = ["all", "datadog", "otel", "sentry", "statsd"] [package.metadata.requires-dev] codegen = [ - { name = "datamodel-code-generator", extras = ["http"], specifier = "==0.33.0" }, + { name = "datamodel-code-generator", extras = ["http"], specifier = ">=0.71.0" }, { name = "openapi-spec-validator", specifier = ">=0.7.1" }, { name = "rich", specifier = ">=13.6.0" }, { name = "svcs", specifier = ">=25.1.0" }, @@ -11323,23 +11323,22 @@ wheels = [ [[package]] name = "datamodel-code-generator" -version = "0.33.0" +version = "0.71.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "argcomplete" }, - { name = "black" }, + { name = "black", marker = "sys_platform != 'emscripten'" }, { name = "genson" }, { name = "inflect" }, - { name = "isort" }, + { name = "isort", marker = "sys_platform != 'emscripten'" }, { name = "jinja2" }, - { name = "packaging" }, { name = "pydantic" }, { name = "pyyaml" }, - { name = "tomli", marker = "python_full_version < '3.12'" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7d/a0/3f81c5c0c31d6f25f459e370e04553810f096e5dbd3cf7eda2a67709cd78/datamodel_code_generator-0.33.0.tar.gz", hash = "sha256:7635ef788201d69bd3e98ba88ce6afe479400dc2737fe9d5e21f87408f352c08", size = 458695, upload-time = "2025-08-14T13:50:36.965Z" } +sdist = { url = "https://files.pythonhosted.org/packages/54/f5/f4ce23d99503b147c9ec514dc995a96d3b4d2a3284252ad665f875a3145d/datamodel_code_generator-0.71.0.tar.gz", hash = "sha256:d27cd7a0d10f9b2db74a41db7f3e050c226da9cf0afb4916a7ab56275ebacbf2", size = 1684916, upload-time = "2026-07-24T15:32:04.334Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/d0/acd7a19dad4dc4118d2b6ba06df9a4a3725729e91fa3f2c63a765d4e7d44/datamodel_code_generator-0.33.0-py3-none-any.whl", hash = "sha256:e229264aa612b2d5bb4901bcd6c520a799ae0d5c19262577a0f876eb48afaaa3", size = 121148, upload-time = "2025-08-14T13:50:35.306Z" }, + { url = "https://files.pythonhosted.org/packages/c6/4d/556cb290170f41b97ce50fd872e10a266f141d7a38352bb3071e4ae61f41/datamodel_code_generator-0.71.0-py3-none-any.whl", hash = "sha256:680b68338d59e98a0559eeb54d8e5ca33c35b3ec0bef922ec2cc783f2cb28e9a", size = 452379, upload-time = "2026-07-24T15:32:02.467Z" }, ] [package.optional-dependencies]