Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Determine needs_expansion on client side when DB isolated #39257

Closed
wants to merge 1 commit into from
Closed

Determine needs_expansion on client side when DB isolated #39257

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 6 additions & 2 deletions airflow/cli/commands/task_command.py
View file
Open in desktop
Expand Up @@ -166,6 +166,7 @@ def _get_ti_db_access(
exec_date_or_run_id: str | None = None,
pool: str | None = None,
create_if_necessary: CreateIfNecessary = False,
should_expand=False,
session: Session = NEW_SESSION,
) -> tuple[TaskInstance | TaskInstancePydantic, bool]:
"""Get the task instance through DagRun.run_id, if that fails, get the TI the old way."""
Expand All @@ -177,7 +178,7 @@ def _get_ti_db_access(

if not exec_date_or_run_id and not create_if_necessary:
raise ValueError("Must provide `exec_date_or_run_id` if not `create_if_necessary`.")
if needs_expansion(task):
if should_expand:
if map_index < 0:
raise RuntimeError("No map_index passed to mapped task")
elif map_index >= 0:
Expand Down Expand Up @@ -218,18 +219,21 @@ def _get_ti(
if dag is None:
raise ValueError("Cannot get task instance for a task not assigned to a DAG")

should_expand = needs_expansion(task)
ti, dr_created = _get_ti_db_access(
dag=dag,
task=task,
map_index=map_index,
exec_date_or_run_id=exec_date_or_run_id,
pool=pool,
create_if_necessary=create_if_necessary,
should_expand=should_expand,
)

# setting ti.task is necessary for AIP-44 since the task object does not serialize perfectly
# if we update the serialization logic for Operator to also serialize the dag object on it,
# then this would not be necessary;
ti.task = task
ti.refresh_from_task(task, pool_override=pool)
return ti, dr_created


Expand Down