kalluripradeep commented on code in PR #64503:
URL: https://github.com/apache/airflow/pull/64503#discussion_r3034699539
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -1443,6 +1443,25 @@ def mynameis(arg):
assert indices == [0, 1, 2, 3]
+def test_verify_integrity_handles_stale_data_error(dag_maker, session):
+ """Test that StaleDataError during _create_task_instances is caught and
session is rolled back."""
+ from unittest.mock import patch
+
+ from sqlalchemy.orm.exc import StaleDataError
+
Review Comment:
Fixed, moved `StaleDataError` up to module level and switched to
`mock.patch.object`.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1879,9 +1880,10 @@ def _create_task_instances(
extra_tags={"task_type": task_type},
)
session.flush()
- except IntegrityError:
+ except (IntegrityError, StaleDataError) as exc:
self.log.info(
Review Comment:
Fair point — though this matches how `IntegrityError` is already handled in
the same block (stats emitted per-attempt, not per-flush-success). Happy to
address it in a follow-up if preferred.
##########
airflow-core/tests/unit/utils/test_retries.py:
##########
@@ -78,6 +79,29 @@ def test_function(session):
assert mock_rollback.call_count == 3
mock_rollback.assert_has_calls([mock.call(), mock.call(), mock.call()])
+ @pytest.mark.db_test
+ def test_retry_db_transaction_with_stale_data_error(self, caplog):
+ """Test that StaleDataError is retried just like DBAPIError"""
+ mock_obj = mock.MagicMock()
+ mock_session = mock.MagicMock()
+ mock_rollback = mock.MagicMock()
+ mock_session.rollback = mock_rollback
Review Comment:
Added `spec=['execute', 'rollback']` to the session mock.
--
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]
