This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bb6be350c0f Auto-generate the Pydantic datamodels for TaskSDK in
pre-commit (#47026)
bb6be350c0f is described below
commit bb6be350c0f40ea683a04eaccef3366fbf602dfd
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Tue Feb 25 15:01:33 2025 +0000
Auto-generate the Pydantic datamodels for TaskSDK in pre-commit (#47026)
We have so far been sort-of running this manually when we remember, which is
obviously error prone and forgotten once or twice.
This automates it so it can't be forgotten.
In order to get the license text in there, and to format it with ruff (not
just Black) I had to add a very simple custom code formatter.
---
.pre-commit-config.yaml | 9 ++
airflow/api_fastapi/execution_api/app.py | 8 ++
.../execution_api/datamodels/taskinstance.py | 2 +-
contributing-docs/08_static_code_checks.rst | 2 +
dev/breeze/doc/images/output_static-checks.svg | 34 ++++----
dev/breeze/doc/images/output_static-checks.txt | 2 +-
dev/breeze/src/airflow_breeze/pre_commit_ids.py | 1 +
dev/datamodel_code_formatter.py | 66 +++++++++++++++
task_sdk/dev/generate_models.py | 99 ++++++++++++++++++++++
task_sdk/pyproject.toml | 12 ++-
.../src/airflow/sdk/api/datamodels/_generated.py | 8 +-
.../airflow/sdk/definitions/_internal/templater.py | 11 ++-
.../src/airflow/sdk/execution_time/task_runner.py | 2 +-
13 files changed, 227 insertions(+), 29 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ab73a9dfe8f..f10e2c69a7c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1448,6 +1448,15 @@ repos:
entry: ./scripts/ci/pre_commit/update_fastapi_api_spec.py
pass_filenames: false
files: ^airflow/api_fastapi/.*\.py$
+ exclude: ^airflow/api_fastapi/execution_api/.*
+ additional_dependencies: ['rich>=12.4.4']
+ - id: generate-tasksdk-datamodels
+ name: Generate Datamodels for TaskSDK client
+ language: python
+ entry: uv run --active --group codegen --project
apache-airflow-task-sdk --directory task_sdk -s dev/generate_models.py
+ pass_filenames: false
+ files: ^airflow/api_fastapi/execution_api/.*\.py$
+ require_serial: true
additional_dependencies: ['rich>=12.4.4']
- id: update-er-diagram
name: Update ER diagram
diff --git a/airflow/api_fastapi/execution_api/app.py
b/airflow/api_fastapi/execution_api/app.py
index cee567f2d49..a9bbf4dd788 100644
--- a/airflow/api_fastapi/execution_api/app.py
+++ b/airflow/api_fastapi/execution_api/app.py
@@ -72,6 +72,14 @@ def create_task_execution_api_app() -> FastAPI:
if schema_name not in openapi_schema["components"]["schemas"]:
openapi_schema["components"]["schemas"][schema_name] = schema
+ # The `JsonValue` component is missing any info. causes issues when
generating models
+ openapi_schema["components"]["schemas"]["JsonValue"] = {
+ "title": "Any valid JSON value",
+ "anyOf": [
+ {"type": t} for t in ("string", "number", "integer", "object",
"array", "boolean", "null")
+ ],
+ }
+
app.openapi_schema = openapi_schema
return app.openapi_schema
diff --git a/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
b/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index acf56833666..a176a2f9d5b 100644
--- a/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++ b/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -227,7 +227,7 @@ class DagRun(StrictBaseModel):
run_after: UtcDateTime
start_date: UtcDateTime
end_date: UtcDateTime | None
- clear_number: int
+ clear_number: int = 0
run_type: DagRunType
conf: Annotated[dict[str, Any], Field(default_factory=dict)]
diff --git a/contributing-docs/08_static_code_checks.rst
b/contributing-docs/08_static_code_checks.rst
index e45e7d718fe..74e1a692e48 100644
--- a/contributing-docs/08_static_code_checks.rst
+++ b/contributing-docs/08_static_code_checks.rst
@@ -288,6 +288,8 @@ require Breeze Docker image to be built locally.
+-----------------------------------------------------------+--------------------------------------------------------+---------+
| generate-pypi-readme | Generate PyPI
README | |
+-----------------------------------------------------------+--------------------------------------------------------+---------+
+| generate-tasksdk-datamodels | Generate
Datamodels for TaskSDK client | * |
++-----------------------------------------------------------+--------------------------------------------------------+---------+
| generate-volumes-for-sources | Generate volumes
for docker compose | |
+-----------------------------------------------------------+--------------------------------------------------------+---------+
| identity | Print checked
files | |
diff --git a/dev/breeze/doc/images/output_static-checks.svg
b/dev/breeze/doc/images/output_static-checks.svg
index 20a07c78d0e..42e49195398 100644
--- a/dev/breeze/doc/images/output_static-checks.svg
+++ b/dev/breeze/doc/images/output_static-checks.svg
@@ -369,23 +369,23 @@
</text><text class="breeze-static-checks-r5" x="0" y="1044.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-42)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1044.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-42)">compile-www-assets-dev | create-missing-init-py-files-tests | debug-statements | </text><text
class="breeze-static-checks-r5" x="1451.8" y="1044.8" textLength="12.2"
clip-path="url(#breeze-static-checks [...]
</text><text class="breeze-static-checks-r5" x="0" y="1069.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-43)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1069.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-43)">detect-private-key | doctoc | end-of-file-fixer | fix-encoding-pragma | flynt |  </text><text
class="breeze-static-checks-r5" x="1451.8" y="1069.2" textLength="12.2"
clip-path=" [...]
</text><text class="breeze-static-checks-r5" x="0" y="1093.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-44)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1093.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-44)">generate-airflow-diagrams | generate-openapi-spec | generate-pypi-readme |       </text><text
class="breeze-static-checks-r5" x="1451.8" y="1093.6" textLength="12.2" clip-p
[...]
-</text><text class="breeze-static-checks-r5" x="0" y="1118" textLength="12.2"
clip-path="url(#breeze-static-checks-line-45)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1118" textLength="988.2"
clip-path="url(#breeze-static-checks-line-45)">generate-volumes-for-sources | identity | insert-license | kubeconform |         </text><text
class="breeze-static-checks-r5" x="1451.8" y="1118" textLength [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1142.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-46)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1142.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-46)">lint-chart-schema | lint-css | lint-dockerfile | lint-helm-chart |               </text><text
class="breeze-static-checks- [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1166.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-47)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1166.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-47)">lint-json-schema | lint-markdown | lint-openapi | mixed-line-ending |            </text><text
class="breeze-static-checks-r5" x="1451.8" [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1191.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-48)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1191.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-48)">mypy-airflow | mypy-dev | mypy-docs | mypy-providers | mypy-task-sdk |           </text><text
class="breeze-static-checks-r5" x="145 [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1215.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-49)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1215.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-49)">pretty-format-json | pylint | python-no-log-warn | replace-bad-characters |      </text><text
class="breeze-static-checks-r5" x="1451.8" y="1215.6" textLength="12.2" c [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1240" textLength="12.2"
clip-path="url(#breeze-static-checks-line-50)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1240" textLength="988.2"
clip-path="url(#breeze-static-checks-line-50)">rst-backticks | ruff | ruff-format | shellcheck | trailing-whitespace |          </text><text
class="breeze-static-checks-r5" x="1451.8" y="1 [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1264.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-51)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1264.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-51)">ts-compile-format-lint-ui | ts-compile-format-lint-www | update-black-version |  </text><text
class="breeze-static-checks-r5" x="1451.8" y="1264.4" textLength="12.2"
clip-path="url(#breeze-static-c [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1288.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-52)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1288.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-52)">update-breeze-cmd-output | update-breeze-readme-config-hash |                    </text><text
class="breeze-static-ch [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1313.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-53)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1313.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-53)">update-chart-dependencies | update-er-diagram | update-extras |                  </text><text
class="breeze-static-ch [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1337.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-54)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1337.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-54)">update-in-the-wild-to-be-sorted | update-inlined-dockerfile-scripts |            </text><text
class="breeze-static-checks-r5" x="1451.8" y="1337.6" textLengt [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1362" textLength="12.2"
clip-path="url(#breeze-static-checks-line-55)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1362" textLength="988.2"
clip-path="url(#breeze-static-checks-line-55)">update-installed-providers-to-be-sorted | update-installers-and-pre-commit |     </text><text
class="breeze-static-checks-r5" x="1451.8" y="1362" textLength="12.2"
clip-path="url(#breeze-static-ch [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1386.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-56)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1386.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-56)">update-local-yml-file | update-migration-references |                           &#
[...]
-</text><text class="breeze-static-checks-r5" x="0" y="1410.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-57)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1410.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-57)">update-openapi-spec-tags-to-be-sorted | update-providers-build-files |           </text><text
class="breeze-static-checks-r5" x="1451.8" y="1410.8" textLength="12 [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1435.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-58)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1435.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-58)">update-providers-dependencies | update-reproducible-source-date-epoch |          </text><text
class="breeze-static-checks-r5" x="1451.8" y="1435.2" textLength="12.2" c [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1459.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-59)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1459.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-59)">update-spelling-wordlist-to-be-sorted | update-supported-versions |              </text><text
class="breeze-static-checks-r5" x="1451.8" y="1459.6" [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1484" textLength="12.2"
clip-path="url(#breeze-static-checks-line-60)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1484" textLength="988.2"
clip-path="url(#breeze-static-checks-line-60)">update-vendored-in-k8s-json-schema | update-version | validate-operators-init |  </text><text
class="breeze-static-checks-r5" x="1451.8" y="1484" textLength="12.2"
clip-path="url(#breeze-static-checks- [...]
-</text><text class="breeze-static-checks-r5" x="0" y="1508.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-61)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1508.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-61)">yamllint | zizmor)                                  
[...]
+</text><text class="breeze-static-checks-r5" x="0" y="1118" textLength="12.2"
clip-path="url(#breeze-static-checks-line-45)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1118" textLength="988.2"
clip-path="url(#breeze-static-checks-line-45)">generate-tasksdk-datamodels | generate-volumes-for-sources | identity |          </text><text
class="breeze-static-checks-r5" x="1451.8" y="1118" textLength="12. [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1142.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-46)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1142.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-46)">insert-license | kubeconform | lint-chart-schema | lint-css | lint-dockerfile |  </text><text
class="breeze-static-checks-r5" x="1451.8" y="1142.4" textLength="12.2"
clip-path=" [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1166.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-47)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1166.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-47)">lint-helm-chart | lint-json-schema | lint-markdown | lint-openapi |              </text><text
class="breeze-static-checks-r5" x [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1191.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-48)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1191.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-48)">mixed-line-ending | mypy-airflow | mypy-dev | mypy-docs | mypy-providers |       </text><text
class="breeze-static-checks-r5" x="1451.8" y="1191.2" text [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1215.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-49)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1215.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-49)">mypy-task-sdk | pretty-format-json | pylint | python-no-log-warn |               </text><text
class="breeze-static-checks- [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1240" textLength="12.2"
clip-path="url(#breeze-static-checks-line-50)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1240" textLength="988.2"
clip-path="url(#breeze-static-checks-line-50)">replace-bad-characters | rst-backticks | ruff | ruff-format | shellcheck |       </text><text
class="breeze-static-checks-r5" x="1451.8" y="1240" textLength [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1264.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-51)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1264.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-51)">trailing-whitespace | ts-compile-format-lint-ui | ts-compile-format-lint-www |   </text><text
class="breeze-static-checks-r5" x="1451.8" y="1264.4" textLength="12.2"
clip-path="url(#breeze-sta [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1288.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-52)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1288.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-52)">update-black-version | update-breeze-cmd-output |                            
[...]
+</text><text class="breeze-static-checks-r5" x="0" y="1313.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-53)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1313.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-53)">update-breeze-readme-config-hash | update-chart-dependencies | update-er-diagram </text><text
class="breeze-static-checks-r5" x="1451.8" y="1313.2" textLength="12.2"
clip-path="url(#breeze-static-checks-line [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1337.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-54)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1337.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-54)">| update-extras | update-in-the-wild-to-be-sorted |                           
[...]
+</text><text class="breeze-static-checks-r5" x="0" y="1362" textLength="12.2"
clip-path="url(#breeze-static-checks-line-55)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1362" textLength="988.2"
clip-path="url(#breeze-static-checks-line-55)">update-inlined-dockerfile-scripts | update-installed-providers-to-be-sorted |    </text><text
class="breeze-static-checks-r5" x="1451.8" y="1362" textLength="12.2"
clip-path="url(#breeze-static-checks- [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1386.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-56)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1386.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-56)">update-installers-and-pre-commit | update-local-yml-file |                       </text><text
class="b [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1410.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-57)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1410.8" textLength="988.2"
clip-path="url(#breeze-static-checks-line-57)">update-migration-references | update-openapi-spec-tags-to-be-sorted |            </text><text
class="breeze-static-checks-r5" x="1451.8" y="1410.8" textLengt [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1435.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-58)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1435.2" textLength="988.2"
clip-path="url(#breeze-static-checks-line-58)">update-providers-build-files | update-providers-dependencies |                   </text><text
class="breeze-static-checks- [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1459.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-59)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1459.6" textLength="988.2"
clip-path="url(#breeze-static-checks-line-59)">update-reproducible-source-date-epoch | update-spelling-wordlist-to-be-sorted |  </text><text
class="breeze-static-checks-r5" x="1451.8" y="1459.6" textLength="12.2"
clip-path="url(#breeze-static-checks-line [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1484" textLength="12.2"
clip-path="url(#breeze-static-checks-line-60)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1484" textLength="988.2"
clip-path="url(#breeze-static-checks-line-60)">update-supported-versions | update-vendored-in-k8s-json-schema | update-version |</text><text
class="breeze-static-checks-r5" x="1451.8" y="1484" textLength="12.2"
clip-path="url(#breeze-static-checks-line-60)"> [...]
+</text><text class="breeze-static-checks-r5" x="0" y="1508.4"
textLength="12.2" clip-path="url(#breeze-static-checks-line-61)">│</text><text
class="breeze-static-checks-r7" x="451.4" y="1508.4" textLength="988.2"
clip-path="url(#breeze-static-checks-line-61)">validate-operators-init | yamllint | zizmor)                            
[...]
</text><text class="breeze-static-checks-r5" x="0" y="1532.8"
textLength="12.2" clip-path="url(#breeze-static-checks-line-62)">│</text><text
class="breeze-static-checks-r4" x="24.4" y="1532.8" textLength="268.4"
clip-path="url(#breeze-static-checks-line-62)">--show-diff-on-failure</text><text
class="breeze-static-checks-r6" x="402.6" y="1532.8" textLength="24.4"
clip-path="url(#breeze-static-checks-line-62)">-s</text><text
class="breeze-static-checks-r1" x="451.4" y="1532.8" textLength=" [...]
</text><text class="breeze-static-checks-r5" x="0" y="1557.2"
textLength="12.2" clip-path="url(#breeze-static-checks-line-63)">│</text><text
class="breeze-static-checks-r4" x="24.4" y="1557.2" textLength="292.8"
clip-path="url(#breeze-static-checks-line-63)">--initialize-environment</text><text
class="breeze-static-checks-r1" x="451.4" y="1557.2" textLength="549"
clip-path="url(#breeze-static-checks-line-63)">Initialize environment before running checks.</text><text
c [...]
</text><text class="breeze-static-checks-r5" x="0" y="1581.6"
textLength="12.2" clip-path="url(#breeze-static-checks-line-64)">│</text><text
class="breeze-static-checks-r4" x="24.4" y="1581.6" textLength="353.8"
clip-path="url(#breeze-static-checks-line-64)">--max-initialization-attempts</text><text
class="breeze-static-checks-r1" x="451.4" y="1581.6" textLength="854"
clip-path="url(#breeze-static-checks-line-64)">Maximum number of attempts to initialize env
[...]
diff --git a/dev/breeze/doc/images/output_static-checks.txt
b/dev/breeze/doc/images/output_static-checks.txt
index 5ecb06b3bb6..2517d29745d 100644
--- a/dev/breeze/doc/images/output_static-checks.txt
+++ b/dev/breeze/doc/images/output_static-checks.txt
@@ -1 +1 @@
-26a564635ee0c2aaecaea40dd7e93779
+e7ca86fda9d023fbd26f9a8a1ce83467
diff --git a/dev/breeze/src/airflow_breeze/pre_commit_ids.py
b/dev/breeze/src/airflow_breeze/pre_commit_ids.py
index b5d308f8291..ca65bb7193e 100644
--- a/dev/breeze/src/airflow_breeze/pre_commit_ids.py
+++ b/dev/breeze/src/airflow_breeze/pre_commit_ids.py
@@ -109,6 +109,7 @@ PRE_COMMIT_LIST = [
"generate-airflow-diagrams",
"generate-openapi-spec",
"generate-pypi-readme",
+ "generate-tasksdk-datamodels",
"generate-volumes-for-sources",
"identity",
"insert-license",
diff --git a/dev/datamodel_code_formatter.py b/dev/datamodel_code_formatter.py
new file mode 100644
index 00000000000..0f59fde22e0
--- /dev/null
+++ b/dev/datamodel_code_formatter.py
@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+import libcst as cst
+from datamodel_code_generator.format import CustomCodeFormatter
+
+
+def license_text() -> str:
+ license = (
+ Path(__file__).parents[1].joinpath("scripts", "ci",
"license-templates", "LICENSE.txt").read_text()
+ )
+
+ return "\n".join(f"# {line}" if line else "#" for line in
license.splitlines()) + "\n"
+
+
+class CodeFormatter(CustomCodeFormatter):
+ def apply(self, code: str) -> str:
+ code = license_text() + code
+
+ # Swap "class JsonValue[RootValue]:" for the import from pydantic
+ class JsonValueNodeRemover(cst.CSTTransformer):
+ def leave_ImportFrom(
+ self, original_node: cst.ImportFrom, updated_node:
cst.ImportFrom
+ ) -> cst.BaseSmallStatement |
cst.FlattenSentinel[cst.BaseSmallStatement] | cst.RemovalSentinel:
+ if original_node.module and original_node.module.value ==
"pydantic":
+ new_names = updated_node.names +
(cst.ImportAlias(name=cst.Name("JsonValue")),) # type: ignore[operator]
+ return updated_node.with_changes(names=new_names)
+ return super().leave_ImportFrom(original_node, updated_node)
+
+ def leave_ClassDef(
+ self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+ ) -> cst.BaseStatement | cst.FlattenSentinel[cst.BaseStatement] |
cst.RemovalSentinel:
+ if original_node.name.value == "JsonValue":
+ return cst.RemoveFromParent()
+ return super().leave_ClassDef(original_node, updated_node)
+
+ source_tree = cst.parse_module(code)
+ modified_tree = source_tree.visit(JsonValueNodeRemover())
+ code = modified_tree.code
+
+ result = subprocess.check_output(
+ ["ruff", "check", "--fix-only", "--unsafe-fixes", "--quiet",
"--preview", "-"],
+ input=code,
+ text=True,
+ )
+
+ return result
diff --git a/task_sdk/dev/generate_models.py b/task_sdk/dev/generate_models.py
new file mode 100644
index 00000000000..60a17333de9
--- /dev/null
+++ b/task_sdk/dev/generate_models.py
@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from datamodel_code_generator import (
+ DataModelType,
+ DatetimeClassType,
+ InputFileType,
+ LiteralType,
+ PythonVersion,
+ generate as generate_models,
+)
+
+os.environ["_AIRFLOW__AS_LIBRARY"] = "1"
+sys.path.insert(
+ 0,
+ str(Path(__file__).parents[2].joinpath("scripts", "ci",
"pre_commit").resolve()),
+) # make sure common utils are importable
+
+from common_precommit_utils import (
+ AIRFLOW_SOURCES_ROOT_PATH,
+)
+
+sys.path.insert(0, str(AIRFLOW_SOURCES_ROOT_PATH)) # make sure setup is
imported from Airflow
+
+from airflow.api_fastapi.execution_api.app import create_task_execution_api_app
+
+task_sdk_root = Path(__file__).parents[1]
+
+if TYPE_CHECKING:
+ from fastapi import FastAPI
+
+
+def load_config():
+ try:
+ from tomllib import load as load_tomllib
+ except ImportError:
+ from tomli import load as load_tomllib
+
+ pyproject = task_sdk_root / "pyproject.toml"
+ # Simulate what `datamodel-code-generator` does on the CLI
+ with pyproject.open("rb") as fh:
+ cfg = load_tomllib(fh)["tool"]["datamodel-codegen"]
+
+ cfg = {k.replace("-", "_"): v for k, v in cfg.items()}
+
+ # Translate config file option names to generate() kwarg names
+ if (use_default := cfg.pop("use_default", None)) is not None:
+ cfg["apply_default_values_for_required_fields"] = use_default
+
+ if cfg.get("use_annotated"):
+ cfg["field_constraints"] = True
+
+ cfg["output"] = Path(cfg["output"])
+ cfg["output_model_type"] = DataModelType(cfg["output_model_type"])
+ cfg["output_datetime_class"] =
DatetimeClassType(cfg["output_datetime_class"])
+ cfg["input_file_type"] = InputFileType(cfg["input_file_type"])
+ cfg["target_python_version"] = PythonVersion(cfg["target_python_version"])
+ cfg["enum_field_as_literal"] = LiteralType(cfg["enum_field_as_literal"])
+ return cfg
+
+
+def generate_file(app: FastAPI):
+ # The persisted openapi spec will list all endpoints (public and ui), this
+ # is used for code generation.
+ for route in app.routes:
+ if getattr(route, "name") == "webapp":
+ continue
+ route.__setattr__("include_in_schema", True)
+
+ os.chdir(task_sdk_root)
+
+ openapi_schema = json.dumps(app.openapi())
+ args = load_config()
+ args["input_filename"] = args.pop("url")
+ generate_models(openapi_schema, **args)
+
+
+generate_file(create_task_execution_api_app())
diff --git a/task_sdk/pyproject.toml b/task_sdk/pyproject.toml
index 6cd0215e82c..aa8c3e42124 100644
--- a/task_sdk/pyproject.toml
+++ b/task_sdk/pyproject.toml
@@ -109,11 +109,16 @@ exclude_also = [
[dependency-groups]
codegen = [
- "datamodel-code-generator[http]>=0.26.5",
+ "datamodel-code-generator[http]>=0.28.1",
+ "apache-airflow",
+ "rich",
+ "ruff",
]
-[tool.black]
-# This is needed for datamodel-codegen to treat this as the "project" file
+[tool.uv.sources]
+# These names must match the names as defined in the pyproject.toml of the
workspace items,
+# *not* the workspace folder paths
+apache-airflow = {workspace = true}
# To use:
#
@@ -134,6 +139,7 @@ use-schema-description=true # Desc becomes class doc
comment
use-standard-collections=true # list[] not List[]
use-subclass-enum=true # enum, not union of Literals
use-union-operator=true # 3.9+annotations, not `Union[]`
+custom-formatters = ['dev.datamodel_code_formatter',]
url = 'http://0.0.0.0:9091/execution/openapi.json'
output = 'src/airflow/sdk/api/datamodels/_generated.py'
diff --git a/task_sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task_sdk/src/airflow/sdk/api/datamodels/_generated.py
index 7814095007c..71f110b64f4 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:9091/execution/openapi.json
-# version: 0.26.5
+# version: 0.28.1
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
@@ -290,13 +290,13 @@ class DagRun(BaseModel):
)
dag_id: Annotated[str, Field(title="Dag Id")]
run_id: Annotated[str, Field(title="Run Id")]
- logical_date: Annotated[datetime | None, Field(title="Logical Date")]
+ logical_date: Annotated[datetime | None, Field(title="Logical Date")] =
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
run_after: Annotated[datetime, Field(title="Run After")]
start_date: Annotated[datetime, Field(title="Start Date")]
end_date: Annotated[datetime | None, Field(title="End Date")] = None
- clear_number: Annotated[int, Field(title="Clear Number")] = 0
+ clear_number: Annotated[int | None, Field(title="Clear Number")] = 0
run_type: DagRunType
conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None
@@ -311,7 +311,7 @@ class TIRunContext(BaseModel):
"""
dag_run: DagRun
- task_reschedule_count: Annotated[int, Field(title="Task Reschedule
Count")] = 0
+ task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule
Count")] = 0
max_tries: Annotated[int, Field(title="Max Tries")]
variables: Annotated[list[VariableResponse] | None,
Field(title="Variables")] = None
connections: Annotated[list[ConnectionResponse] | None,
Field(title="Connections")] = None
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/templater.py
b/task_sdk/src/airflow/sdk/definitions/_internal/templater.py
index b50c4dbb3ca..a0914113496 100644
--- a/task_sdk/src/airflow/sdk/definitions/_internal/templater.py
+++ b/task_sdk/src/airflow/sdk/definitions/_internal/templater.py
@@ -27,11 +27,11 @@ import jinja2
import jinja2.nativetypes
import jinja2.sandbox
-from airflow.io.path import ObjectStoragePath
from airflow.sdk.definitions._internal.mixins import ResolveMixin
from airflow.utils.helpers import render_template_as_native,
render_template_to_string
if TYPE_CHECKING:
+ from airflow.io.path import ObjectStoragePath
from airflow.models.operator import Operator
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.dag import DAG
@@ -154,6 +154,13 @@ class Templater:
*RecursionError* on circular dependencies)
:return: Templated content
"""
+ try:
+ # Delay this to runtime, it invokes provider manager otherwise
+ from airflow.io.path import ObjectStoragePath
+ except ImportError:
+ # A placeholder class so isinstance checks work
+ class ObjectStoragePath: ... # type: ignore[no-redef]
+
# "content" is a bad name, but we're stuck to it being public API.
value = content
del content
@@ -203,7 +210,7 @@ class Templater:
serialized_path = value.serialize()
path_version = value.__version__
serialized_path["path"] =
self._render(jinja_env.from_string(serialized_path["path"]), context)
- return ObjectStoragePath.deserialize(data=serialized_path,
version=path_version)
+ return value.deserialize(data=serialized_path, version=path_version)
def _render_nested_template_fields(
self,
diff --git a/task_sdk/src/airflow/sdk/execution_time/task_runner.py
b/task_sdk/src/airflow/sdk/execution_time/task_runner.py
index 6868019c38b..610c989e01f 100644
--- a/task_sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task_sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -159,7 +159,7 @@ class RuntimeTaskInstance(TaskInstance):
# TODO: Assess if we need to pass these through
timezone.coerce_datetime
"dag_run": dag_run, # type: ignore[typeddict-item] #
Removable after #46522
"task_instance_key_str":
f"{self.task.dag_id}__{self.task.task_id}__{dag_run.run_id}",
- "task_reschedule_count":
self._ti_context_from_server.task_reschedule_count,
+ "task_reschedule_count":
self._ti_context_from_server.task_reschedule_count or 0,
"prev_start_date_success": lazy_object_proxy.Proxy(
lambda: get_previous_dagrun_success(self.id).start_date
),
