Optimise and migrate to SA2-compatible syntax for TaskReschedule #33720

Taragolis · 2023-08-25T09:07:26Z

related: #28723

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

airflow/models/taskreschedule.py

Taragolis · 2023-08-27T19:05:51Z

Interesting errors which happen on, it might refers either and error in code during change behaviour or the fact that something wrong with test on Postgres. Locally I've also have this error, wit flakey behaviour: 5-10 sequential fail, 2-3 success, and again failure

FAILED tests/sensors/test_base.py::TestBaseSensor::test_ok_with_reschedule - AssertionError: assert datetime.datetime(2023, 8, 27, 16, 39, 45, 511220, tzinfo=Timezone('UTC')) == datetime.datetime(2023, 8, 27, 16, 39, 55, 511220, tzinfo=Timezone('UTC'))
 +  where datetime.datetime(2023, 8, 27, 16, 39, 45, 511220, tzinfo=Timezone('UTC')) = <airflow.models.taskreschedule.TaskReschedule object at 0x7fee0a9bdf70>.start_date
FAILED tests/sensors/test_base.py::TestBaseSensor::test_ok_with_custom_reschedule_exception - AssertionError: assert datetime.datetime(2023, 8, 27, 16, 39, 46, 338455, tzinfo=Timezone('UTC')) == datetime.datetime(2023, 8, 27, 16, 40, 46, 338455, tzinfo=Timezone('UTC'))
 +  where datetime.datetime(2023, 8, 27, 16, 39, 46, 338455, tzinfo=Timezone('UTC')) = <airflow.models.taskreschedule.TaskReschedule object at 0x7fee0a900e80>.start_date
FAILED tests/sensors/test_base.py::TestBaseSensor::test_reschedule_and_retry_timeout - Failed: DID NOT RAISE <class 'airflow.exceptions.AirflowSensorTimeout'>
FAILED tests/sensors/test_base.py::TestBaseSensor::test_reschedule_and_retry_timeout_and_silent_fail - Failed: DID NOT RAISE <class 'airflow.exceptions.AirflowSensorTimeout'>

uranusjr · 2023-08-28T07:09:17Z

airflow/models/taskreschedule.py

+
+    @classmethod
+    @provide_session
+    def find_last_for_task_instance(


I wonder if we really need this method at all. The logic is rather simple and does not really provide any encapsulation benefits. (Same for find_for_task_instance above but it’s already there so there’s probably an argument for keeping it.)

I change logic a bit methods

New (all marked as private):
stmt_for_task_instance -> Query builder
find_task_reschedule_for_ti -> Return single task_reschedule record, first or last
find_task_reschedules_for_ti -> replacement for find_for_task_instance, not used in production code, only in tests

Deprecated Methods, for backward compatibility
query_for_task_instance
find_for_task_instance

In general only find_task_reschedule_for_ti and stmt_for_task_instance uses, all other are redundant, and in theory could be dropped. However find_task_reschedules_for_ti help me to find my mistake during tests

But I mean we don’t really need find_task_reschedule_for_ti and find_task_reschedules_for_ti since those are just one liners that pass most of the arguments (except session) to stmt_for_task_instance anyway. We can just call stmt_for_task_instance and execute the statement as appropriated directly where needed.

I've implement by this way due to difference between legacy Query API and new Statement Based (i don't know correct name): you don't have to provide a DB session for build query statement

We could directly use stmt_for_task_instance into

airflow/models/taskinstance.py

airflow/ti_deps/deps/ready_to_reschedule.py

Unfortunetly we can't do that in airflow/sensors/base.py, because we do not open session here, however it is not a showstopper, if applicable I could create private method in BaseSensorOperator which would open a session, execute query against stmt_for_task_instance and return start_date

You don’t even need a private method, just use with create_session(), that’s what @provide_session does anyway.

Nice! I just need to figure out how to deal with tests now, I'm not a fan of mocking DB execution if we could avoid it.

I've just thought, are find_task_reschedule_for_ti / find_task_reschedules_for_ti would be required for AIP-44?

Actually might be keep only find_task_reschedule_for_ti and use it again, WDYT?

Taragolis · 2023-08-29T17:16:38Z

wit flakey behaviour: 5-10 sequential fail, 2-3 success, and again failure

🤦 🤦 🤦 🤦 🤦 🤦 🤦 🤦 🤦 🤦 🤦 🤦 🤦
I've forgot about assignments: stmt = stmt.order_by() instead of stmt.order_by()

tests/ti_deps/deps/test_ready_to_reschedule_dep.py

Taragolis · 2023-08-27T16:16:41Z

airflow/models/taskreschedule.py

+    @staticmethod
+    @provide_session
+    def query_for_task_instance(
+        task_instance: TaskInstance,
+        descending: bool = False,
+        session: Session = NEW_SESSION,
+        try_number: int | None = None,
+    ) -> Query:
+        """
+        Return query for task reschedules for a given the task instance (deprecated).
+
+        :param session: the database session object
+        :param task_instance: the task instance to find task reschedules for
+        :param descending: If True then records are returned in descending order
+        :param try_number: Look for TaskReschedule of the given try_number. Default is None which
+            looks for the same try_number of the given task_instance.
+        """
+        warnings.warn(
+            "`query_for_task_instance` use SQLAlchemy's Legacy Query API.",
+            category=RemovedInAirflow3Warning,
+            stacklevel=2,
+        )
+
+        if try_number is None:
+            try_number = task_instance.try_number
+
+        TR = TaskReschedule
+        qry = session.query(TR).filter(
+            TR.dag_id == task_instance.dag_id,
+            TR.task_id == task_instance.task_id,
+            TR.run_id == task_instance.run_id,
+            TR.map_index == task_instance.map_index,
+            TR.try_number == try_number,
+        )
+        if descending:
+            return qry.order_by(desc(TR.id))
+        else:
+            return qry.order_by(asc(TR.id))


I've decided to keep query_for_task_instance as is, because if some one use it before, than expects that this method return sqlalchemy.orm.Query and there is no direct replacement for this object, which are implements same methods.

If we think that this part never been a part of public interface, we could remove it, this PR also remove usage of this method

airflow/models/taskreschedule.py

uranusjr · 2023-08-31T06:08:29Z

tests/models/test_taskinstance.py

+def task_reschedules_for_ti():
+    def wrapper(ti):
+        with create_session() as session:
+            return session.scalars(TaskReschedule.stmt_for_task_instance(ti=ti, descending=False)).all()
+
+    return wrapper


Suggested change

def task_reschedules_for_ti():

def wrapper(ti):

with create_session() as session:

return session.scalars(TaskReschedule.stmt_for_task_instance(ti=ti, descending=False)).all()

return wrapper

def task_reschedules_for_ti(session):

def wrapper(ti):

return session.scalars(TaskReschedule.stmt_for_task_instance(ti=ti, descending=False)).all()

return wrapper

(I think this should work. session here is a function-level fixture.)

I would be better try to avoid to keep session longer than it required

phanikumv · 2023-09-05T17:53:20Z

We need to run full tests on the changes related to SqlAlchemy 2.0 , hence applied the label

phanikumv · 2023-09-05T17:56:12Z

@Taragolis can you please rebase so that we can run all the tests on this PR

potiuk · 2023-09-05T18:17:49Z

@Taragolis can you please rebase so that we can run all the tests on this PR

I rebased it.

Taragolis · 2023-09-20T14:48:54Z

@Taragolis Taragolis 3 weeks ago
I've just thought, are find_task_reschedule_for_ti / find_task_reschedules_for_ti would be required for AIP-44?

@Taragolis Taragolis 3 weeks ago
Actually might be keep only find_task_reschedule_for_ti and use it again, WDYT?

cc: @potiuk @phanikumv @uranusjr @mhenc @vincbeck

phanikumv · 2023-10-16T14:30:44Z

@Taragolis could you please resolve the conflict here

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>

boring-cyborg bot added the area:core-operators Operators, Sensors and hooks within Core Airflow label Aug 25, 2023

Taragolis commented Aug 25, 2023

View reviewed changes

airflow/models/taskreschedule.py Outdated Show resolved Hide resolved

Taragolis force-pushed the optimise-taskreschedule branch from 5d46d0f to f62c3e1 Compare August 27, 2023 16:04

uranusjr reviewed Aug 28, 2023

View reviewed changes

Taragolis force-pushed the optimise-taskreschedule branch from f62c3e1 to 0ed825e Compare August 29, 2023 17:37

Taragolis requested review from potiuk and phanikumv August 29, 2023 20:51

Taragolis commented Aug 29, 2023

View reviewed changes

Taragolis marked this pull request as ready for review August 29, 2023 20:58

Taragolis requested review from kaxil, XD-DENG and ashb as code owners August 29, 2023 20:58

Taragolis force-pushed the optimise-taskreschedule branch from beaaf08 to e72b0bd Compare August 30, 2023 18:46

uranusjr reviewed Aug 31, 2023

View reviewed changes

airflow/models/taskreschedule.py Outdated Show resolved Hide resolved

uranusjr reviewed Aug 31, 2023

View reviewed changes

uranusjr mentioned this pull request Aug 31, 2023

SqlAlchemy 2.0 changes for taskreschedule.py #33891

Closed

Taragolis force-pushed the optimise-taskreschedule branch 2 times, most recently from e44ee27 to 489f8bb Compare September 4, 2023 22:43

uranusjr approved these changes Sep 5, 2023

View reviewed changes

phanikumv added the full tests needed We need to run full set of tests for this PR to merge label Sep 5, 2023

potiuk force-pushed the optimise-taskreschedule branch from 489f8bb to 132b538 Compare September 5, 2023 18:17

Taragolis force-pushed the optimise-taskreschedule branch from 132b538 to d588e3f Compare September 5, 2023 22:45

phanikumv approved these changes Sep 12, 2023

View reviewed changes

ephraimbuddy approved these changes Oct 16, 2023

View reviewed changes

Optimise and migrate to SA2-compatible syntax for TaskReschedule

e521db4

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>

Taragolis force-pushed the optimise-taskreschedule branch from d588e3f to e521db4 Compare October 16, 2023 17:28

ephraimbuddy merged commit 85fd0e1 into apache:main Oct 16, 2023
66 checks passed

ephraimbuddy added the type:improvement Changelog: Improvements label Oct 16, 2023

Taragolis deleted the optimise-taskreschedule branch October 16, 2023 21:31

eladkal added this to the Airflow 2.8.0 milestone Oct 17, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Optimise and migrate to SA2-compatible syntax for TaskReschedule #33720

Optimise and migrate to SA2-compatible syntax for TaskReschedule #33720

Taragolis commented Aug 25, 2023 •

edited by phanikumv

Taragolis commented Aug 27, 2023

uranusjr Aug 28, 2023

Taragolis Aug 29, 2023

uranusjr Aug 30, 2023 •

edited

Taragolis Aug 30, 2023

uranusjr Aug 30, 2023

Taragolis Aug 30, 2023

Taragolis Aug 31, 2023

Taragolis Aug 31, 2023

Taragolis commented Aug 29, 2023

Taragolis Aug 27, 2023

uranusjr Aug 31, 2023 •

edited

Taragolis Aug 31, 2023

phanikumv commented Sep 5, 2023

phanikumv commented Sep 5, 2023

potiuk commented Sep 5, 2023

Taragolis commented Sep 20, 2023

phanikumv commented Oct 16, 2023

Optimise and migrate to SA2-compatible syntax for TaskReschedule #33720

Optimise and migrate to SA2-compatible syntax for TaskReschedule #33720

Conversation

Taragolis commented Aug 25, 2023 • edited by phanikumv

Taragolis commented Aug 27, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

uranusjr Aug 30, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Taragolis commented Aug 29, 2023

Choose a reason for hiding this comment

uranusjr Aug 31, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

phanikumv commented Sep 5, 2023

phanikumv commented Sep 5, 2023

potiuk commented Sep 5, 2023

Taragolis commented Sep 20, 2023

phanikumv commented Oct 16, 2023

Taragolis commented Aug 25, 2023 •

edited by phanikumv

uranusjr Aug 30, 2023 •

edited

uranusjr Aug 31, 2023 •

edited