Fix test fixture cleanup for bundles and teams#68376
Conversation
|
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
|
The _fetch_remote_logging_conn function in the supervisor was caching None results when a connection lookup failed. In the two-token execution model, an initial lookup with a restricted token might fail, and caching this None result prevented subsequent lookups with an execution token from succeeding. This change ensures that only successful connection lookups are cached, allowing for retries when a better token becomes available. Fixes: apache#68366
| @@ -1181,10 +1181,10 @@ def _fetch_remote_logging_conn(conn_id: str, client: Client) -> Connection | Non | |||
| from airflow.sdk.definitions.connection import Connection | |||
|
|
|||
There was a problem hiding this comment.
I think changes to this file is not related to the issue attached to this pr. And I think changes to this file can be dropped
|
Thank you so much for the PR. I think the changes to supervisor.py can be dropped, seems unrelated to the issue tagged in the description. And also, I think this is initial part(fix) of the issue, I would suggest updating the description to be related to tagged issue. |
ferruzzi
left a comment
There was a problem hiding this comment.
That was fast, thanks! We definitely need to get it manually tested and I left some comments to look at.
| delete(DagBundleModel).where(DagBundleModel.name.in_(bundle_names)) | ||
| ) | ||
| self.session.commit() | ||
| self.created_bundle_names.difference_update(bundle_names) |
There was a problem hiding this 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.
| with attempt: | ||
| dag_ids = list(self.dagbag.dag_ids) | ||
| if not dag_ids: | ||
| bundle_names = set(self.created_bundle_names) |
There was a problem hiding this 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.
| finally: | ||
| if created: | ||
| with create_session() as session: | ||
| session.execute(delete(DagBundleModel).where(DagBundleModel.name == "testing")) |
There was a problem hiding this 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.
| if created: | ||
| session.rollback() | ||
| session.execute(delete(Team).where(Team.name == "testing")) | ||
|
|
There was a problem hiding this 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.
Summary
dag_makertest fixture and remove them during fixture cleanup.testing_dag_bundlefixture when it creates the sharedtestingbundle row.testing_teamfixture when it creates the sharedtestingteam row.Details
dag_maker.cleanup()already removes DAG-related rows such asDagModel,DagRun,TaskInstance,DagVersion,XCom,TaskMap, andAssetEvent, but it did not removeDagBundleModelrows thatdag_makercreated. This could leave bundle rows behind and force later tests to defensively callclear_db_dag_bundles().This change records bundle names created by
dag_makerand deletes those bundles after DAG model/version cleanup, preserving FK-safe cleanup order. It also convertstesting_dag_bundleandtesting_teaminto fixtures with teardown paths so rows created for a test do not leak into later tests.Testing
python -m py_compile devel-common\src\tests_common\pytest_plugin.pygit diff --checkTargeted pytest was not run locally because
uvandpytestare not installed in this environment.Fixes: #68374