From a82887adf7d046580a78e830872aedbe1e061717 Mon Sep 17 00:00:00 2001 From: Jed Cunningham Date: Tue, 14 Jan 2025 20:28:48 -0700 Subject: [PATCH] Fix typing/docstring on dag helper methods --- airflow/models/dag.py | 6 +++--- tests/models/test_dag.py | 21 +++++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/airflow/models/dag.py b/airflow/models/dag.py index ffd3d2dba56cc..0f7d4d4925c32 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -779,13 +779,13 @@ def get_is_paused(self, session=NEW_SESSION) -> None: return session.scalar(select(DagModel.is_paused).where(DagModel.dag_id == self.dag_id)) @provide_session - def get_bundle_name(self, session=NEW_SESSION) -> None: + def get_bundle_name(self, session=NEW_SESSION) -> str | None: """Return the bundle name this DAG is in.""" return session.scalar(select(DagModel.bundle_name).where(DagModel.dag_id == self.dag_id)) @provide_session - def get_latest_bundle_version(self, session=NEW_SESSION) -> None: - """Return the bundle name this DAG is in.""" + def get_latest_bundle_version(self, session=NEW_SESSION) -> str | None: + """Return the latest version of the bundle this DAG is in.""" return session.scalar(select(DagModel.latest_bundle_version).where(DagModel.dag_id == self.dag_id)) @methodtools.lru_cache(maxsize=None) diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py index 95fa85848dcc9..9f599e7acd47b 100644 --- a/tests/models/test_dag.py +++ b/tests/models/test_dag.py @@ -2089,6 +2089,27 @@ def test_dag_owner_links(self): orm_dag_owners = session.query(DagOwnerAttributes).all() assert not orm_dag_owners + def test_get_bundle_name(self, testing_dag_bundle): + dag = DAG("dag") + + # until we've sycned, it'll be None + assert dag.get_bundle_name() is None + + DAG.bulk_write_to_db("testing", None, [dag]) + assert dag.get_bundle_name() == "testing" + + def test_get_latest_bundle_version(self, testing_dag_bundle): + dag = DAG("dag") + + # until we've sycned, it'll be None + assert dag.get_latest_bundle_version() is None + + # Now, it can be none or a str + DAG.bulk_write_to_db("testing", None, [dag]) + assert dag.get_latest_bundle_version() is None + DAG.bulk_write_to_db("testing", "abc", [dag]) + assert dag.get_latest_bundle_version() == "abc" + class TestDagModel: def _clean(self):