You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "harjeevanmaan (via GitHub)" <gi...@apache.org> on 2024/08/11 19:46:56 UTC
[PR] Better map_index and downstream/upstream tasks handling [airflow]
harjeevanmaan opened a new pull request, #41384:
URL: https://github.com/apache/airflow/pull/41384
- Fixes "Clearing a Mapped Task (with Downstream included) clears all Downstream tasks instead of mapped downstream tasks #41278"
- Fixes "control state of individual taskflow in mapped task-group #40543"
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
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](https://github.com/apache/airflow/tree/main/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on PR #41384:
URL: https://github.com/apache/airflow/pull/41384#issuecomment-2286761387
All unit tests passed. Working on the static checks now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1728770907
##########
airflow/models/taskinstance.py:
##########
@@ -3794,22 +3793,61 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau
return or_(*filter_condition)
@classmethod
- def ti_selector_condition(cls, vals: Collection[str | tuple[str, int]]) -> ColumnOperators:
+ def ti_selector_condition(TaskInstance, vals: Collection[str | tuple[str, int]]) -> ColumnOperators:
Review Comment:
Why change this? Using `cls` here is perfectly fine.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1747196300
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
+ files = ["a", "b", "c"]
+
+ start_date = datetime(2024, 1, 1)
+ files = ["file1", "file2", "file3"]
+
+ with dag_maker(dag_id="task_group_mapping_example", start_date=start_date, schedule=None, catchup=False):
+
+ @task_group(group_id="etl")
+ def etl_pipeline(file):
+ e = EmptyOperator(task_id="e")
+ t = EmptyOperator(task_id="t")
+ last = EmptyOperator(task_id="last")
+
+ e >> t >> last
+
+ etl_pipeline.expand(file=files)
+
+ dag_instance = dag_maker.dag
+ dag_maker.create_dagrun(
+ run_id="manual_run_2024_01_01",
+ state=State.SUCCESS,
+ execution_date=start_date,
+ start_date=start_date,
+ data_interval=(start_date, start_date),
+ run_type=DagRunType.SCHEDULED,
+ )
+
+ # with map_index
+ task_id = "etl.e"
+ task_id_or_regex = [task_id]
+ map_indexes = [0, 1]
+ task_ids = [(task_id, map_index) for map_index in map_indexes]
+ partial_dag = dag_instance.partial_subset(
+ task_ids_or_regex=task_id_or_regex,
+ include_downstream=True,
+ include_upstream=False,
+ )
+
+ # handling downstream tasks
+ if len(partial_dag.task_dict) > 1:
+ task_ids.extend(tid for tid in partial_dag.task_dict if tid != task_id)
+
+ with create_session() as session:
Review Comment:
Already using the `dag_maker `fixture. Added the `session `too
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on PR #41384:
URL: https://github.com/apache/airflow/pull/41384#issuecomment-2288994355
I've committed changes to fix the static-check failures. I would appreciate any feedback or review.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1724942144
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
The session needs to be a `kwarg` of `ti_selector_condition` since the entire logic of filtering revolves around being able to filter out using SQL queries. If the `session` kwarg is not passed the function is decorated by `@provide_session,` which will handle session creation.
One of the tests failed once I removed the session kwarg because of inconsistency between a `taskinstance` and `session,` the `@provide_session` created a new session which lead to this inconsistency.
But if the session is persisted as is shown in the PR, this inconsistency would not occur.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on PR #41384:
URL: https://github.com/apache/airflow/pull/41384#issuecomment-2285786429
@uranusjr Working on it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1724942144
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
@uranusjr Thanks for your review!. The session needs to be a `kwarg` of `ti_selector_condition` since the entire logic of filtering revolves around being able to filter out using SQL queries. If the `session` kwarg is not passed the function is decorated by `@provide_session,` which will handle session creation.
One of the tests failed once I removed the session kwarg because of inconsistency between a `taskinstance` and `session,` the `@provide_session` created a new session which lead to this inconsistency.
But if the session is persisted as is shown in the PR, this inconsistency would not occur.
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
@uranusjr Thanks for your review! The session needs to be a `kwarg` of `ti_selector_condition` since the entire logic of filtering revolves around being able to filter out using SQL queries. If the `session` kwarg is not passed the function is decorated by `@provide_session,` which will handle session creation.
One of the tests failed once I removed the session kwarg because of inconsistency between a `taskinstance` and `session,` the `@provide_session` created a new session which lead to this inconsistency.
But if the session is persisted as is shown in the PR, this inconsistency would not occur.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1720797446
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
I've made the changes. If there are any other modifications needed, please let me know so I can include them all in a single commit
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1729203608
##########
tests/models/test_taskinstance.py:
##########
@@ -5290,3 +5290,70 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition():
+ from airflow.decorators import dag, task_group
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.operators.empty import EmptyOperator
+ from airflow.utils.session import create_session
+ from airflow.utils.state import State
+ from airflow.utils.timezone import datetime
+ from airflow.utils.types import DagRunType
+ from tests.test_utils.db import clear_db_runs
Review Comment:
Any other changes I need to make?
Hopefully the logic itself is correct now?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1741682178
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
Review Comment:
Better to also use a fixture here. `clean_dags_and_dagruns` would be useful.
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
+ files = ["a", "b", "c"]
+
+ start_date = datetime(2024, 1, 1)
+ files = ["file1", "file2", "file3"]
+
+ with dag_maker(dag_id="task_group_mapping_example", start_date=start_date, schedule=None, catchup=False):
+
+ @task_group(group_id="etl")
+ def etl_pipeline(file):
+ e = EmptyOperator(task_id="e")
+ t = EmptyOperator(task_id="t")
+ last = EmptyOperator(task_id="last")
+
+ e >> t >> last
+
+ etl_pipeline.expand(file=files)
+
+ dag_instance = dag_maker.dag
+ dag_maker.create_dagrun(
+ run_id="manual_run_2024_01_01",
+ state=State.SUCCESS,
+ execution_date=start_date,
+ start_date=start_date,
+ data_interval=(start_date, start_date),
+ run_type=DagRunType.SCHEDULED,
+ )
+
+ # with map_index
+ task_id = "etl.e"
+ task_id_or_regex = [task_id]
+ map_indexes = [0, 1]
+ task_ids = [(task_id, map_index) for map_index in map_indexes]
+ partial_dag = dag_instance.partial_subset(
+ task_ids_or_regex=task_id_or_regex,
+ include_downstream=True,
+ include_upstream=False,
+ )
+
+ # handling downstream tasks
+ if len(partial_dag.task_dict) > 1:
+ task_ids.extend(tid for tid in partial_dag.task_dict if tid != task_id)
+
+ with create_session() as session:
Review Comment:
There’s also a `session` fixture. (There’s actually also one on `dag_maker` too.)
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
+ files = ["a", "b", "c"]
+
+ start_date = datetime(2024, 1, 1)
+ files = ["file1", "file2", "file3"]
+
+ with dag_maker(dag_id="task_group_mapping_example", start_date=start_date, schedule=None, catchup=False):
+
+ @task_group(group_id="etl")
+ def etl_pipeline(file):
+ e = EmptyOperator(task_id="e")
+ t = EmptyOperator(task_id="t")
+ last = EmptyOperator(task_id="last")
+
+ e >> t >> last
+
+ etl_pipeline.expand(file=files)
+
+ dag_instance = dag_maker.dag
+ dag_maker.create_dagrun(
+ run_id="manual_run_2024_01_01",
+ state=State.SUCCESS,
+ execution_date=start_date,
+ start_date=start_date,
+ data_interval=(start_date, start_date),
+ run_type=DagRunType.SCHEDULED,
+ )
+
+ # with map_index
+ task_id = "etl.e"
+ task_id_or_regex = [task_id]
+ map_indexes = [0, 1]
+ task_ids = [(task_id, map_index) for map_index in map_indexes]
+ partial_dag = dag_instance.partial_subset(
+ task_ids_or_regex=task_id_or_regex,
+ include_downstream=True,
+ include_upstream=False,
+ )
+
+ # handling downstream tasks
+ if len(partial_dag.task_dict) > 1:
+ task_ids.extend(tid for tid in partial_dag.task_dict if tid != task_id)
Review Comment:
There’s only exactly one case in this test… eitehr we need the `extend` call or we don’t, there’s no need for the `if` condition.
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
+ files = ["a", "b", "c"]
+
+ start_date = datetime(2024, 1, 1)
+ files = ["file1", "file2", "file3"]
+
+ with dag_maker(dag_id="task_group_mapping_example", start_date=start_date, schedule=None, catchup=False):
+
+ @task_group(group_id="etl")
+ def etl_pipeline(file):
+ e = EmptyOperator(task_id="e")
+ t = EmptyOperator(task_id="t")
+ last = EmptyOperator(task_id="last")
+
+ e >> t >> last
+
+ etl_pipeline.expand(file=files)
+
+ dag_instance = dag_maker.dag
+ dag_maker.create_dagrun(
+ run_id="manual_run_2024_01_01",
+ state=State.SUCCESS,
+ execution_date=start_date,
+ start_date=start_date,
+ data_interval=(start_date, start_date),
+ run_type=DagRunType.SCHEDULED,
+ )
Review Comment:
Manual or scheduled? (It doesn’t really matter in this test but it’s just weird.)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1729202429
##########
tests/models/test_taskinstance.py:
##########
@@ -5290,3 +5290,70 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition():
+ from airflow.decorators import dag, task_group
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.operators.empty import EmptyOperator
+ from airflow.utils.session import create_session
+ from airflow.utils.state import State
+ from airflow.utils.timezone import datetime
+ from airflow.utils.types import DagRunType
+ from tests.test_utils.db import clear_db_runs
Review Comment:
I used `dag_maker` as you suggested
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1726181192
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
My point is that the method _should not_ do a fetch in the first place. The method calculates filtering parameters for a query. Using another query to do that is not a generally preferred approach since it makes it difficult to manage where queries happen. It would be best if the multiple queries can be combined into one with JOIN or other SQL constructs, or, if that is not viable, the queries should at least happen in the same function, not in different levels in the call stack.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1720796866
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
Thanks for the review!
Removal of session won't affect anything, I did this because I had to comment out the @provide_session in order to debug. Here, session: Session = NEW_SESSION, it already has a default supplied by the decorator.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1724942144
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
@uranusjr the session needs to be a `kwarg` of `ti_selector_condition` since the entire logic of filtering revolves around being able to filter out using SQL queries. If the `session` kwarg is not passed the function is decorated by `@provide_session,` which will handle session creation.
One of the tests failed once I removed the session kwarg because of inconsistency between a `taskinstance` and `session,` the `@provide_session` created a new session which lead to this inconsistency.
But if the session is persisted as is shown in the PR, this inconsistency would not occur.
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
I've made the changes. If there are any other modifications needed, please let me know so I can include them all in a single commit
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1728768745
##########
tests/models/test_taskinstance.py:
##########
@@ -5290,3 +5290,70 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition():
+ from airflow.decorators import dag, task_group
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.operators.empty import EmptyOperator
+ from airflow.utils.session import create_session
+ from airflow.utils.state import State
+ from airflow.utils.timezone import datetime
+ from airflow.utils.types import DagRunType
+ from tests.test_utils.db import clear_db_runs
Review Comment:
Why do we need to add these imports here?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1747197048
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
+ files = ["a", "b", "c"]
+
+ start_date = datetime(2024, 1, 1)
+ files = ["file1", "file2", "file3"]
+
+ with dag_maker(dag_id="task_group_mapping_example", start_date=start_date, schedule=None, catchup=False):
+
+ @task_group(group_id="etl")
+ def etl_pipeline(file):
+ e = EmptyOperator(task_id="e")
+ t = EmptyOperator(task_id="t")
+ last = EmptyOperator(task_id="last")
+
+ e >> t >> last
+
+ etl_pipeline.expand(file=files)
+
+ dag_instance = dag_maker.dag
+ dag_maker.create_dagrun(
+ run_id="manual_run_2024_01_01",
+ state=State.SUCCESS,
+ execution_date=start_date,
+ start_date=start_date,
+ data_interval=(start_date, start_date),
+ run_type=DagRunType.SCHEDULED,
+ )
+
+ # with map_index
+ task_id = "etl.e"
+ task_id_or_regex = [task_id]
+ map_indexes = [0, 1]
+ task_ids = [(task_id, map_index) for map_index in map_indexes]
+ partial_dag = dag_instance.partial_subset(
+ task_ids_or_regex=task_id_or_regex,
+ include_downstream=True,
+ include_upstream=False,
+ )
+
+ # handling downstream tasks
+ if len(partial_dag.task_dict) > 1:
+ task_ids.extend(tid for tid in partial_dag.task_dict if tid != task_id)
Review Comment:
Thanks for catching it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1727369223
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
My apologies @uranusjr, you are absolutely correct. I should have done that from the very beginning. I've completely removed any dependency on `session` and used simple `subquery` to filter out the `TaskMap` objects
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1720773915
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
It’s quite weird we need to pass in `session` here. Is it possible to do this with some SQL JOIN and/or subquery magic?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #41384: Better map_index and downstream/upstream tasks handling
URL: https://github.com/apache/airflow/pull/41384
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1728769858
##########
tests/models/test_taskinstance.py:
##########
@@ -5290,3 +5290,70 @@ def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition():
+ from airflow.decorators import dag, task_group
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.operators.empty import EmptyOperator
+ from airflow.utils.session import create_session
+ from airflow.utils.state import State
+ from airflow.utils.timezone import datetime
+ from airflow.utils.types import DagRunType
+ from tests.test_utils.db import clear_db_runs
Review Comment:
Also it may be better to use the `dag_maker` fixture instead.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1728767974
##########
tests/models/test_dag.py:
##########
@@ -3429,8 +3429,8 @@ def consumer(value):
state=TaskInstanceState.SUCCESS,
session=session,
)
- assert dr1 in session, "Check session is passed down all the way"
-
+ dag_runs = session.query(DagRun).filter(DagRun.run_id == dr1.run_id).all()
+ assert len(dag_runs) > 0, "Check session is passed down all the way"
Review Comment:
We can use `func.count()` here to improve performance. This doesn’t matter that much since this is just a test (and if `all()` returns anything to cause performance issues, we should fix it anyway) but that’s good practice.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "harjeevanmaan (via GitHub)" <gi...@apache.org>.
harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1720796866
##########
airflow/api/common/mark_tasks.py:
##########
@@ -164,7 +164,7 @@ def get_all_dag_task_query(
qry_dag = select(TaskInstance).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.ti_selector_condition(task_ids),
+ TaskInstance.ti_selector_condition(task_ids, session=session),
Review Comment:
Thanks for the review!
Removal of session won't affect anything, I did this because I had to comment out the @provide_session in order to debug. Here, session: Session = NEW_SESSION, it already has a default supplied by the decorator.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #41384:
URL: https://github.com/apache/airflow/pull/41384#issuecomment-2285650554
Tests are failing
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Better map_index and downstream/upstream tasks handling [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #41384:
URL: https://github.com/apache/airflow/pull/41384#issuecomment-2433857547
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org