This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 987f9c08a89 Add None to Mapped annotations of three nullable ORM
columns (#72579)
987f9c08a89 is described below
commit 987f9c08a89b45204d7ea8444c576f1074c33e45
Author: PoAn Yang <[email protected]>
AuthorDate: Fri Sep 11 11:26:34 2026 +0900
Add None to Mapped annotations of three nullable ORM columns (#72579)
Signed-off-by: PoAn Yang <[email protected]>
---
.pre-commit-config.yaml | 11 +
airflow-core/src/airflow/models/backfill.py | 2 +-
airflow-core/src/airflow/models/callback.py | 2 +-
airflow-core/src/airflow/models/dagrun.py | 2 +-
.../check_mapped_column_nullable_annotations.py | 227 ++++++++++++++++++++
...est_check_mapped_column_nullable_annotations.py | 235 +++++++++++++++++++++
6 files changed, 476 insertions(+), 3 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 4777b57ea24..e463adad7e7 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1358,6 +1358,17 @@ repos:
language: python
pass_filenames: false
files:
^airflow-core/src/airflow/utils/[^/]+$|^scripts/ci/prek/known_airflow_core_utils_modules\.txt$|^scripts/ci/prek/check_no_new_airflow_core_utils_modules\.py$
+ - id: check-mapped-column-nullable-annotations
+ name: Check that nullable mapped_column attributes are annotated
Mapped[X | None]
+ entry: ./scripts/ci/prek/check_mapped_column_nullable_annotations.py
+ language: python
+ pass_filenames: true
+ files: >
+ (?x)
+ ^airflow-core/src/airflow/(models|jobs)/.*\.py$|
+ ^providers/edge3/src/airflow/providers/edge3/models/.*\.py$|
+ ^providers/fab/src/airflow/providers/fab/auth_manager/models/.*\.py$|
+ ^scripts/ci/prek/check_mapped_column_nullable_annotations\.py$
- id: bandit
name: bandit
description: "Bandit is a tool for finding common security issues in
Python code"
diff --git a/airflow-core/src/airflow/models/backfill.py
b/airflow-core/src/airflow/models/backfill.py
index 13b9e6d004f..203be25188e 100644
--- a/airflow-core/src/airflow/models/backfill.py
+++ b/airflow-core/src/airflow/models/backfill.py
@@ -221,7 +221,7 @@ class BackfillDagRun(Base):
backfill_id: Mapped[int] = mapped_column(Integer, nullable=False)
dag_run_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
exception_reason: Mapped[str | None] = mapped_column(StringID(),
nullable=True)
- logical_date: Mapped[datetime] = mapped_column(UtcDateTime, nullable=True)
+ logical_date: Mapped[datetime | None] = mapped_column(UtcDateTime,
nullable=True)
partition_key: Mapped[str | None] = mapped_column(StringID(),
nullable=True)
sort_ordinal: Mapped[int] = mapped_column(Integer, nullable=False)
diff --git a/airflow-core/src/airflow/models/callback.py
b/airflow-core/src/airflow/models/callback.py
index 352885a420e..6295b7dfa7b 100644
--- a/airflow-core/src/airflow/models/callback.py
+++ b/airflow-core/src/airflow/models/callback.py
@@ -144,7 +144,7 @@ class Callback(Base, BaseWorkload):
created_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
# Used for callbacks of type CallbackType.TRIGGERER
- trigger_id: Mapped[int] = mapped_column(Integer, ForeignKey("trigger.id"),
nullable=True)
+ trigger_id: Mapped[int | None] = mapped_column(Integer,
ForeignKey("trigger.id"), nullable=True)
trigger = relationship("Trigger", back_populates="callback", uselist=False)
def __init__(self, priority_weight: int = 1, prefix: str = "", **kwargs):
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index 46fbebb0800..725bf09c77d 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -286,7 +286,7 @@ class DagRun(Base, LoggingMixin):
# This is nullable because it's too costly to migrate dagruns created prior
# to this column's addition (Airflow 3.2.0). If you want a reasonable
# meaningful non-null value, use ``dr.created_at or dr.run_after``.
- created_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=True,
default=timezone.utcnow)
+ created_at: Mapped[datetime | None] = mapped_column(UtcDateTime,
nullable=True, default=timezone.utcnow)
updated_at: Mapped[datetime | None] = mapped_column(
UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow,
nullable=True
)
diff --git a/scripts/ci/prek/check_mapped_column_nullable_annotations.py
b/scripts/ci/prek/check_mapped_column_nullable_annotations.py
new file mode 100755
index 00000000000..f3b41f7917f
--- /dev/null
+++ b/scripts/ci/prek/check_mapped_column_nullable_annotations.py
@@ -0,0 +1,227 @@
+#!/usr/bin/env python
+# 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.
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+# "rich>=13.0.0",
+# ]
+# ///
+"""Check that ``mapped_column(..., nullable=True)`` attributes are annotated
``Mapped[X | None]``.
+
+SQLAlchemy 2 derives a column's nullability from its ``Mapped[...]``
annotation only when
+``nullable=`` is *not* passed explicitly. Once ``nullable=True`` is spelled
out the DDL is
+right, but nothing cross-checks the annotation any more: ``Mapped[datetime]``
on such a column
+tells mypy the attribute can never be ``None``, so callers skip the ``None``
guard and fail at
+runtime with ``AttributeError`` the first time they touch a NULL row.
+
+The check is deliberately one-directional. An Optional annotation on a NOT
NULL column is
+harmless and is not reported.
+
+Modes
+-----
+Default (files passed by prek):
+ Check only the supplied files.
+
+``--all-files``:
+ Walk the directories that define ORM models (``models`` and ``jobs`` in
airflow-core,
+ and the ``models`` packages of the edge3 and FAB providers) and check
every ``.py`` file.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+from common_prek_utils import AIRFLOW_ROOT_PATH
+from rich.console import Console
+from rich.markup import escape
+
+console = Console(color_system="standard", width=200)
+
+REPO_ROOT = AIRFLOW_ROOT_PATH
+# Keep in sync with the hook's ``files`` pattern in .pre-commit-config.yaml.
+SCAN_ROOTS: tuple[str, ...] = (
+ "airflow-core/src/airflow/models",
+ "airflow-core/src/airflow/jobs",
+ "providers/edge3/src/airflow/providers/edge3/models",
+ "providers/fab/src/airflow/providers/fab/auth_manager/models",
+)
+
+_NONE_ADMITTING_NAMES = frozenset({"Any", "Optional"})
+
+
+@dataclass(frozen=True)
+class NullableAnnotationMismatch:
+ path: Path
+ lineno: int
+ attribute: str
+ annotation: str
+
+
+def _extract_trailing_name(node: ast.expr) -> str | None:
+ """Return the last identifier of a ``Name``/``Attribute`` chain
(``orm.Mapped`` -> ``Mapped``)."""
+ if isinstance(node, ast.Name):
+ return node.id
+ if isinstance(node, ast.Attribute):
+ return node.attr
+ return None
+
+
+def _resolve_string_annotation(node: ast.expr) -> ast.expr:
+ if isinstance(node, ast.Constant) and isinstance(node.value, str):
+ try:
+ return ast.parse(node.value, mode="eval").body
+ except SyntaxError:
+ return node
+ return node
+
+
+def admits_none(annotation: ast.expr) -> bool:
+ """Return True if the annotation can hold ``None``: ``X | None``,
``Optional[X]``, ``Union[X, None]``, ``Any``."""
+ annotation = _resolve_string_annotation(annotation)
+ if isinstance(annotation, ast.Constant):
+ return annotation.value is None
+ if isinstance(annotation, ast.BinOp) and isinstance(annotation.op,
ast.BitOr):
+ return admits_none(annotation.left) or admits_none(annotation.right)
+ if _extract_trailing_name(annotation) in _NONE_ADMITTING_NAMES:
+ return True
+ if isinstance(annotation, ast.Subscript):
+ outer = _extract_trailing_name(annotation.value)
+ if outer == "Optional":
+ return True
+ if outer == "Union":
+ members = annotation.slice.elts if isinstance(annotation.slice,
ast.Tuple) else [annotation.slice]
+ return any(admits_none(member) for member in members)
+ return False
+
+
+def _explicit_nullable(call: ast.Call) -> bool | None:
+ """Return the literal ``nullable=`` value, or ``None`` when absent or not
a plain boolean."""
+ for keyword in call.keywords:
+ if keyword.arg == "nullable":
+ value = keyword.value
+ if isinstance(value, ast.Constant) and isinstance(value.value,
bool):
+ return value.value
+ return None
+ return None
+
+
+def iter_mismatches(path: Path) -> Iterator[NullableAnnotationMismatch]:
+ """Yield every ``Mapped[...] = mapped_column(..., nullable=True)`` whose
annotation cannot be ``None``."""
+ try:
+ source = path.read_text(encoding="utf-8")
+ except (OSError, UnicodeDecodeError):
+ return
+ if "mapped_column" not in source:
+ return
+ try:
+ tree = ast.parse(source, filename=str(path))
+ except SyntaxError:
+ return
+ for node in ast.walk(tree):
+ if not isinstance(node, ast.AnnAssign) or node.value is None:
+ continue
+ annotation = _resolve_string_annotation(node.annotation)
+ if not (
+ isinstance(annotation, ast.Subscript) and
_extract_trailing_name(annotation.value) == "Mapped"
+ ):
+ continue
+ call = node.value
+ if not (isinstance(call, ast.Call) and
_extract_trailing_name(call.func) == "mapped_column"):
+ continue
+ if _explicit_nullable(call) is not True or
admits_none(annotation.slice):
+ continue
+ yield NullableAnnotationMismatch(
+ path=path,
+ lineno=node.lineno,
+ attribute=ast.unparse(node.target),
+ annotation=ast.unparse(annotation),
+ )
+
+
+def iter_python_files(roots: Iterable[Path]) -> Iterator[Path]:
+ for root in roots:
+ for path in sorted(root.rglob("*.py")):
+ relative_parts = path.relative_to(root).parts
+ if any(part.startswith(".") or part == "node_modules" for part in
relative_parts):
+ continue
+ yield path
+
+
+def _format_display_path(path: Path) -> str:
+ try:
+ return str(path.resolve().relative_to(REPO_ROOT))
+ except ValueError:
+ return str(path)
+
+
+def check_files(files: Iterable[Path]) -> int:
+ mismatches = [mismatch for path in files for mismatch in
iter_mismatches(path)]
+ if not mismatches:
+ return 0
+ console.print(
+ f"[red]Found {len(mismatches)} nullable mapped_column attribute(s)
whose "
+ "Mapped annotation does not admit None:[/]\n",
+ highlight=False,
+ )
+ for mismatch in mismatches:
+ console.print(
+ f" {_format_display_path(mismatch.path)}:{mismatch.lineno}: "
+ f"[bold]{escape(mismatch.attribute)}[/]:
{escape(mismatch.annotation)}",
+ highlight=False,
+ soft_wrap=True,
+ )
+ console.print(
+ "\nA column declared with [cyan]nullable=True[/] can hold NULL, so its
annotation must say so: "
+ f"annotate it as [cyan]{escape('Mapped[X | None]')}[/], or drop
[cyan]nullable=True[/] if the column "
+ "is meant to be NOT NULL. Otherwise mypy believes the attribute is
never None and lets callers "
+ "skip the guard that the NULL rows require.",
+ highlight=False,
+ )
+ return 1
+
+
+def main(argv: list[str] | None = None) -> int:
+ parser = argparse.ArgumentParser(
+ description="Check that nullable mapped_column attributes are
annotated Mapped[X | None].",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog=__doc__,
+ )
+ parser.add_argument("files", nargs="*", metavar="FILE", help="Files to
check (provided by prek)")
+ parser.add_argument(
+ "--all-files",
+ action="store_true",
+ help="Check every Python file in the ORM model directories",
+ )
+ args = parser.parse_args(argv)
+
+ if args.all_files:
+ return check_files(iter_python_files(REPO_ROOT / root for root in
SCAN_ROOTS))
+
+ if not args.files:
+ console.print("[yellow]No files provided. Pass filenames or use
--all-files.[/yellow]")
+ return 0
+
+ return check_files(Path(file) for file in args.files)
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git
a/scripts/tests/ci/prek/test_check_mapped_column_nullable_annotations.py
b/scripts/tests/ci/prek/test_check_mapped_column_nullable_annotations.py
new file mode 100644
index 00000000000..95bed4307af
--- /dev/null
+++ b/scripts/tests/ci/prek/test_check_mapped_column_nullable_annotations.py
@@ -0,0 +1,235 @@
+# 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 re
+import textwrap
+
+import pytest
+from ci.prek import check_mapped_column_nullable_annotations as hook
+from ci.prek.check_mapped_column_nullable_annotations import check_files,
iter_mismatches, main
+
+MODEL_PREAMBLE = """\
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Optional, Union
+
+from sqlalchemy import DateTime, Integer, orm
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
+
+IS_NULLABLE = True
+
+
+class Base(DeclarativeBase):
+ pass
+
+
+class Model(Base):
+ __tablename__ = "model"
+ id: Mapped[int] = mapped_column(Integer, primary_key=True)
+"""
+COLUMN_LINE = MODEL_PREAMBLE.count("\n") + 1
+_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
+
+
+def model_with(attribute_line: str) -> str:
+ return MODEL_PREAMBLE + textwrap.indent(textwrap.dedent(attribute_line), "
")
+
+
+class TestIterMismatches:
+ @pytest.mark.parametrize(
+ ("attribute_line", "expected_annotation"),
+ [
+ pytest.param(
+ "created_at: Mapped[datetime] = mapped_column(DateTime,
nullable=True)",
+ "Mapped[datetime]",
+ id="plain",
+ ),
+ pytest.param(
+ "trigger_id: Mapped[int] = mapped_column(Integer,
nullable=True, default=None)",
+ "Mapped[int]",
+ id="other-kwargs",
+ ),
+ pytest.param(
+ 'created_at: Mapped["datetime"] = mapped_column(DateTime,
nullable=True)',
+ "Mapped['datetime']",
+ id="string-inner",
+ ),
+ pytest.param(
+ 'created_at: "Mapped[datetime]" = mapped_column(DateTime,
nullable=True)',
+ "Mapped[datetime]",
+ id="string-whole",
+ ),
+ pytest.param(
+ "created_at: orm.Mapped[datetime] =
orm.mapped_column(DateTime, nullable=True)",
+ "orm.Mapped[datetime]",
+ id="qualified",
+ ),
+ ],
+ )
+ def test_reports_nullable_column_whose_annotation_cannot_be_none(
+ self, write_python_file, attribute_line, expected_annotation
+ ):
+ path = write_python_file(model_with(attribute_line))
+
+ mismatches = list(iter_mismatches(path))
+
+ assert len(mismatches) == 1
+ mismatch = mismatches[0]
+ assert mismatch.path == path
+ assert mismatch.lineno == COLUMN_LINE
+ assert mismatch.attribute == attribute_line.split(":")[0]
+ assert mismatch.annotation == expected_annotation
+
+ @pytest.mark.parametrize(
+ "annotation",
+ [
+ "Mapped[datetime | None]",
+ "Mapped[None | datetime]",
+ "Mapped[Optional[datetime]]",
+ "Mapped[Union[datetime, None]]",
+ "Mapped[Union[None, datetime]]",
+ "Mapped[Any]",
+ "Mapped[dict[str, Any] | None]",
+ 'Mapped["datetime | None"]',
+ '"Mapped[datetime | None]"',
+ ],
+ )
+ def test_accepts_annotations_that_admit_none(self, write_python_file,
annotation):
+ path = write_python_file(
+ model_with(f"created_at: {annotation} = mapped_column(DateTime,
nullable=True)")
+ )
+
+ assert list(iter_mismatches(path)) == []
+
+ @pytest.mark.parametrize(
+ "attribute_line",
+ [
+ pytest.param("created_at: Mapped[datetime] =
mapped_column(DateTime)", id="nullable-derived"),
+ pytest.param(
+ "created_at: Mapped[datetime] = mapped_column(DateTime,
nullable=False)", id="not-nullable"
+ ),
+ pytest.param(
+ "created_at: Mapped[datetime] = mapped_column(DateTime,
nullable=IS_NULLABLE)",
+ id="non-literal-nullable",
+ ),
+ pytest.param(
+ "created_at: Mapped[datetime | None] = mapped_column(DateTime,
nullable=False)",
+ id="optional-on-not-null-is-not-reported",
+ ),
+ pytest.param("children: Mapped[list[Model]] =
relationship(Model)", id="relationship"),
+ pytest.param("created_at = mapped_column(DateTime,
nullable=True)", id="unannotated"),
+ pytest.param("created_at: Mapped[datetime]", id="no-value"),
+ pytest.param("count: int = 0", id="not-mapped"),
+ ],
+ )
+ def test_ignores_attributes_outside_the_rule(self, write_python_file,
attribute_line):
+ path = write_python_file(model_with(attribute_line))
+
+ assert list(iter_mismatches(path)) == []
+
+ def test_reports_multiline_call_at_the_annotation_line(self,
write_python_file):
+ path = write_python_file(
+ model_with(
+ """\
+ updated_at: Mapped[datetime] = mapped_column(
+ DateTime,
+ default=datetime.now,
+ nullable=True,
+ )
+ """
+ )
+ )
+
+ assert [(m.attribute, m.lineno) for m in iter_mismatches(path)] ==
[("updated_at", COLUMN_LINE)]
+
+ def test_reports_every_mismatch_in_a_file(self, write_python_file):
+ path = write_python_file(
+ model_with(
+ """\
+ created_at: Mapped[datetime] = mapped_column(DateTime,
nullable=True)
+ updated_at: Mapped[datetime | None] = mapped_column(DateTime,
nullable=True)
+ trigger_id: Mapped[int] = mapped_column(Integer, nullable=True)
+ """
+ )
+ )
+
+ assert [m.attribute for m in iter_mismatches(path)] == ["created_at",
"trigger_id"]
+
+ def test_skips_files_without_mapped_column(self, write_python_file):
+ path = write_python_file("created_at: Mapped[datetime] =
Column(DateTime, nullable=True)\n")
+
+ assert list(iter_mismatches(path)) == []
+
+ def test_skips_unparsable_and_missing_files(self, write_python_file,
tmp_path):
+ broken = write_python_file("created_at: Mapped[datetime] =
mapped_column(DateTime, nullable=True\n")
+
+ assert list(iter_mismatches(broken)) == []
+ assert list(iter_mismatches(tmp_path / "missing.py")) == []
+
+
+class TestCheckFiles:
+ def test_reports_each_mismatch_with_location_and_fails(self,
write_python_file, capsys):
+ path = write_python_file(
+ model_with("created_at: Mapped[datetime] = mapped_column(DateTime,
nullable=True)")
+ )
+
+ assert check_files([path]) == 1
+
+ output = _ANSI_RE.sub("", capsys.readouterr().out)
+ assert f"{path}:{COLUMN_LINE}" in output
+ assert "created_at" in output
+ assert "Mapped[X | None]" in output
+
+ def test_passes_silently_when_clean(self, write_python_file, capsys):
+ path = write_python_file(
+ model_with("created_at: Mapped[datetime | None] =
mapped_column(DateTime, nullable=True)")
+ )
+
+ assert check_files([path]) == 0
+ assert capsys.readouterr().out == ""
+
+
+class TestMain:
+ def test_checks_given_files(self, write_python_file):
+ path = write_python_file(
+ model_with("created_at: Mapped[datetime] = mapped_column(DateTime,
nullable=True)")
+ )
+
+ assert main([str(path)]) == 1
+
+ def test_no_files_is_a_noop(self):
+ assert main([]) == 0
+
+ def test_all_files_walks_scan_roots_and_skips_hidden_dirs(self, tmp_path,
monkeypatch):
+ monkeypatch.setattr(hook, "REPO_ROOT", tmp_path)
+ monkeypatch.setattr(hook, "SCAN_ROOTS", ("dist",))
+ bad = model_with("created_at: Mapped[datetime] =
mapped_column(DateTime, nullable=True)")
+ for relative in (
+ "dist/src/pkg/models.py",
+ "dist/.venv/lib/models.py",
+ "dist/node_modules/x/models.py",
+ ):
+ target = tmp_path / relative
+ target.parent.mkdir(parents=True)
+ target.write_text(bad)
+
+ assert main(["--all-files"]) == 1
+ assert [str(path.relative_to(tmp_path)) for path in
hook.iter_python_files([tmp_path / "dist"])] == [
+ "dist/src/pkg/models.py"
+ ]
