ferruzzi commented on code in PR #68376:
URL: https://github.com/apache/airflow/pull/68376#discussion_r3397832692
##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -1399,7 +1401,8 @@ def cleanup(self):
for attempt in run_with_db_retries(logger=self.log):
with attempt:
dag_ids = list(self.dagbag.dag_ids)
- if not dag_ids:
+ bundle_names = set(self.created_bundle_names)
Review Comment:
Any reason to not use `self.created_bundle_names` directly? I think in this
case it is overly defensive; I don't see where anything is modifying the set
once it reaches this point.
##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -2924,40 +2935,60 @@ def mock_xcom_backend():
def testing_dag_bundle():
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
- if AIRFLOW_V_3_0_PLUS:
- from sqlalchemy import func, select
+ if not AIRFLOW_V_3_0_PLUS:
+ yield
+ return
- from airflow.models.dagbundle import DagBundleModel
- from airflow.utils.session import create_session
+ from sqlalchemy import delete, func, select
- with create_session() as session:
- if (
- session.scalar(
-
select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name ==
"testing")
- )
- == 0
- ):
- testing = DagBundleModel(name="testing")
- session.add(testing)
+ from airflow.models.dagbundle import DagBundleModel
+ from airflow.utils.session import create_session
+
+ created = False
+ with create_session() as session:
+ if (
+ session.scalar(
+
select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name ==
"testing")
+ )
+ == 0
+ ):
+ testing = DagBundleModel(name="testing")
+ session.add(testing)
+ created = True
+
+ try:
+ yield
+ finally:
+ if created:
+ with create_session() as session:
+
session.execute(delete(DagBundleModel).where(DagBundleModel.name == "testing"))
Review Comment:
I am pretty sure this is one of the "quick fixes" I tried before creating
this issue and this will blow up. We'll see when the CI finishes. I didn't
dive too deep into it once I saw it wasn't working, but I suspect that this
fails because (some of??) the DagModels in tests are being created after this
fixture fires and _those_ aren't being cleaned, which causes an F-Key Reference
error.
If the CI is failing, I'd suggest you start there. If the CI passes then
I'm mistaken and you can resolve this with my apologies.
##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -1420,11 +1425,17 @@ def cleanup(self):
)
self.session.execute(delete(DagRun).where(DagRun.dag_id.in_(dag_ids)))
self.session.execute(delete(TaskInstance).where(TaskInstance.dag_id.in_(dag_ids)))
-
self.session.execute(delete(XCom).where(XCom.dag_id.in_(dag_ids)))
-
self.session.execute(delete(DagModel).where(DagModel.dag_id.in_(dag_ids)))
-
self.session.execute(delete(TaskMap).where(TaskMap.dag_id.in_(dag_ids)))
-
self.session.execute(delete(AssetEvent).where(AssetEvent.source_dag_id.in_(dag_ids)))
+ if dag_ids:
+
self.session.execute(delete(XCom).where(XCom.dag_id.in_(dag_ids)))
+
self.session.execute(delete(DagModel).where(DagModel.dag_id.in_(dag_ids)))
+
self.session.execute(delete(TaskMap).where(TaskMap.dag_id.in_(dag_ids)))
+
self.session.execute(delete(AssetEvent).where(AssetEvent.source_dag_id.in_(dag_ids)))
+ if AIRFLOW_V_3_0_PLUS and bundle_names:
+ self.session.execute(
+
delete(DagBundleModel).where(DagBundleModel.name.in_(bundle_names))
+ )
self.session.commit()
+ self.created_bundle_names.difference_update(bundle_names)
Review Comment:
Similar to above, you can likely drop this. We should be deleting all
bundle_names, right? The SQL above is batch-deleting ll of them, then here you
remove all of them from the set. MAYBE if you wanted to do this then validate
that it is now an empty set, that could be a reason for it, but as of now it's
just emptying the set before teardown.
##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -2924,40 +2935,60 @@ def mock_xcom_backend():
def testing_dag_bundle():
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
- if AIRFLOW_V_3_0_PLUS:
- from sqlalchemy import func, select
+ if not AIRFLOW_V_3_0_PLUS:
+ yield
+ return
- from airflow.models.dagbundle import DagBundleModel
- from airflow.utils.session import create_session
+ from sqlalchemy import delete, func, select
- with create_session() as session:
- if (
- session.scalar(
-
select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name ==
"testing")
- )
- == 0
- ):
- testing = DagBundleModel(name="testing")
- session.add(testing)
+ from airflow.models.dagbundle import DagBundleModel
+ from airflow.utils.session import create_session
+
+ created = False
+ with create_session() as session:
+ if (
+ session.scalar(
+
select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name ==
"testing")
+ )
+ == 0
+ ):
+ testing = DagBundleModel(name="testing")
+ session.add(testing)
+ created = True
+
+ try:
+ yield
+ finally:
+ if created:
+ with create_session() as session:
+
session.execute(delete(DagBundleModel).where(DagBundleModel.name == "testing"))
@pytest.fixture
def testing_team():
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
- if AIRFLOW_V_3_0_PLUS:
- from sqlalchemy import select
+ if not AIRFLOW_V_3_0_PLUS:
+ yield None
+ return
- from airflow.models.team import Team
- from airflow.utils.session import create_session
+ from sqlalchemy import delete, select
- with create_session() as session:
- team = session.scalar(select(Team).where(Team.name == "testing"))
- if not team:
- team = Team(name="testing")
- session.add(team)
- session.flush()
- yield team
+ from airflow.models.team import Team
+ from airflow.utils.session import create_session
+
+ created = False
+ with create_session() as session:
+ team = session.scalar(select(Team).where(Team.name == "testing"))
+ if not team:
+ team = Team(name="testing")
+ session.add(team)
+ session.flush()
+ created = True
+ yield team
+ if created:
+ session.rollback()
+ session.execute(delete(Team).where(Team.name == "testing"))
Review Comment:
SQL is one of my weak points, but this smells funny. Rolling back before
deleting feels wrong here? 100% willing to be wrong here and learn something
from you if this is an area you are comfortable with.
If I understand this right, we're creating the Team and yielding it, then
after the test finishes with it we're rolling it back, then deleting it? I'm
guessing the rollback is in case the test modified it during the yield? Which
maybe makes sense, but if we're deleting it anyway, does it matter if we roll
it back first? If the session is in a good/unmodified state then the rollback
is unnecessary. If the session is in a bad/modified state, then the test has
already failed and we should nuke it anyway, no?
Again, this is a blind spot for me, I cold absolutely be missing something.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
