From 8ff8cc55793a9d2561c237d0183da5ee82001d8b Mon Sep 17 00:00:00 2001 From: johncmerfeld Date: Wed, 13 Nov 2024 11:25:15 -0600 Subject: [PATCH 1/4] make boolean check explicit --- airflow/utils/db.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 5748adf0e46e6..87cb09c4fe6e1 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -1557,7 +1557,9 @@ def __setstate__(self, state: Any) -> None: self._session = get_current_task_instance_session() def __bool__(self) -> bool: - return check_query_exists(self._select_asc, session=self._session) + if check := check_query_exists(self._select_asc, session=self._session) is not None: + return check + return False def __eq__(self, other: Any) -> bool: if not isinstance(other, collections.abc.Sequence): From c8170abcd1999d1f10669110bf297f64d335fd94 Mon Sep 17 00:00:00 2001 From: johncmerfeld Date: Tue, 19 Nov 2024 08:51:17 -0600 Subject: [PATCH 2/4] alter where bool cast is --- airflow/utils/db.py | 7 +++---- tests/utils/test_db.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 87cb09c4fe6e1..5efa382a54dc7 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -1458,7 +1458,8 @@ def check_query_exists(query_stmt: Select, *, session: Session) -> bool: :meta private: """ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery()) - return session.scalar(count_stmt) + # we must cast to bool because scalar() can return None + return bool(session.scalar(count_stmt)) def exists_query(*where: ClauseElement, session: Session) -> bool: @@ -1557,9 +1558,7 @@ def __setstate__(self, state: Any) -> None: self._session = get_current_task_instance_session() def __bool__(self) -> bool: - if check := check_query_exists(self._select_asc, session=self._session) is not None: - return check - return False + return check_query_exists(self._select_asc, session=self._session) def __eq__(self, other: Any) -> bool: if not isinstance(other, collections.abc.Sequence): diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 1d3412d361412..b05932ab689b2 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -37,6 +37,7 @@ from airflow.models import Base as airflow_base from airflow.settings import engine from airflow.utils.db import ( + LazySelectSequence, _get_alembic_config, check_migrations, compare_server_default, @@ -56,6 +57,12 @@ pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] +class EmptyLazySelectSequence(LazySelectSequence): + _data = [] + + def __init__(self): + super().__init__(None, None, session="MockSession") + class TestDb: def test_database_schema_and_sqlalchemy_model_are_in_sync(self): import airflow.models @@ -251,3 +258,10 @@ def test_alembic_configuration(self): import airflow assert config.config_file_name == os.path.join(os.path.dirname(airflow.__file__), "alembic.ini") + + @mock.patch("sqlalchemy.orm.Session.scalar") + def test_bool_lazy_select_sequence(self, mock_scalar): + mock_scalar.return_value = None + + lss = EmptyLazySelectSequence() + assert not bool(lss) \ No newline at end of file From 1deed35df019aaf208941cf71ed3b96b63ffba53 Mon Sep 17 00:00:00 2001 From: johncmerfeld Date: Thu, 21 Nov 2024 17:15:04 -0600 Subject: [PATCH 3/4] fix test --- tests/utils/test_db.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index b05932ab689b2..7b99f1480ee17 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -32,7 +32,7 @@ from alembic.migration import MigrationContext from alembic.runtime.environment import EnvironmentContext from alembic.script import ScriptDirectory -from sqlalchemy import MetaData +from sqlalchemy import Column, Integer, MetaData, Table, select from airflow.models import Base as airflow_base from airflow.settings import engine @@ -57,12 +57,6 @@ pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] -class EmptyLazySelectSequence(LazySelectSequence): - _data = [] - - def __init__(self): - super().__init__(None, None, session="MockSession") - class TestDb: def test_database_schema_and_sqlalchemy_model_are_in_sync(self): import airflow.models @@ -259,9 +253,15 @@ def test_alembic_configuration(self): assert config.config_file_name == os.path.join(os.path.dirname(airflow.__file__), "alembic.ini") - @mock.patch("sqlalchemy.orm.Session.scalar") - def test_bool_lazy_select_sequence(self, mock_scalar): - mock_scalar.return_value = None + def test_bool_lazy_select_sequence(self): + class MockSession: + def __init__(self): + pass + + def scalar(self, stmt): + return None + + t = Table("t", MetaData(), Column("id", Integer, primary_key=True)) + lss = LazySelectSequence.from_select(select(t.c.id), order_by=[], session=MockSession()) - lss = EmptyLazySelectSequence() - assert not bool(lss) \ No newline at end of file + assert not bool(lss) From 24b36d35d938303139119ad1c685cb80502d5b02 Mon Sep 17 00:00:00 2001 From: johncmerfeld Date: Fri, 22 Nov 2024 09:19:11 -0600 Subject: [PATCH 4/4] make test more explicit --- tests/utils/test_db.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 7b99f1480ee17..471bdeecc1e37 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -264,4 +264,4 @@ def scalar(self, stmt): t = Table("t", MetaData(), Column("id", Integer, primary_key=True)) lss = LazySelectSequence.from_select(select(t.c.id), order_by=[], session=MockSession()) - assert not bool(lss) + assert bool(lss) is False