[v3-3-test] Raise the datamodel-code-generator floor for the codegen toolchain (#70670)

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

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

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",
         }
 

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"
 ]
 

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
 

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,13 +1042,19 @@ class TestDagOperations:
 
     # DagRun related
     trigger_dag_run = TriggerDAGRunPostBody(
+        logical_date=None,
         conf=None,
         note=None,
     )
 
     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,

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

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

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,

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 {}),
                 }
             ),

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"]

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,

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",