Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,11 @@ def date_param():
"dags update example_bash_operator --no-is-paused",
# Dag Run commands
"dagrun list --dag-id example_bash_operator --state success --limit=1",
# Task instance commands
'tasks list example_bash_operator "manual__{date_param}"',
'tasks get example_bash_operator "manual__{date_param}" runme_0',
"tasks clear example_bash_operator --dry-run",
'tasks update example_bash_operator "manual__{date_param}" runme_0 --new-state=failed',
# XCom commands - need a Dag run with completed tasks
'xcom add example_bash_operator "manual__{date_param}" runme_0 {xcom_key} \'{{"test": "value"}}\'',
'xcom get example_bash_operator "manual__{date_param}" runme_0 {xcom_key}',
Expand Down
3 changes: 2 additions & 1 deletion airflow-ctl/docs/images/command_hashes.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
main:27a22c00dcf32e7a1a4f06672dc8e3c8
main:164bc97843d5be583c0b48f7a34dc8c8
assets:6419e20452692f577c4c6f570b74be0c
auth:d79e9c7d00c432bdbcbc2a86e2e32053
backfill:74c8737b0a62a86ed3605fa9e6165874
Expand All @@ -12,4 +12,5 @@ providers:34502fe09dc0b8b0a13e7e46efdffda6
variables:f8fc76d3d398b2780f4e97f7cd816646
version:31f4efdf8de0dbaaa4fac71ff7efecc3
plugins:4864fd8f356704bd2b3cd1aec3567e35
tasks:7ab24cac521242b6b6012e2bcd317831
auth login:9fe2bb1dd5c602beea2eefb33a2b20a8
134 changes: 69 additions & 65 deletions airflow-ctl/docs/images/output_main.svg
Comment thread
Suraj-kumar00 marked this conversation as resolved.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
117 changes: 117 additions & 0 deletions airflow-ctl/docs/images/output_tasks.svg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
7 changes: 7 additions & 0 deletions airflow-ctl/src/airflowctl/api/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
PoolsOperations,
ProvidersOperations,
ServerResponseError,
TasksOperations,
VariablesOperations,
VersionOperations,
XComOperations,
Expand Down Expand Up @@ -474,6 +475,12 @@ def plugins(self):
"""Operations related to plugins."""
return PluginsOperations(self)

@lru_cache() # type: ignore[prop-decorator]
@property
def tasks(self):
"""Operations related to tasks."""
return TasksOperations(self)


# API Client Decorator for CLI Actions
@contextlib.contextmanager
Expand Down
46 changes: 46 additions & 0 deletions airflow-ctl/src/airflowctl/api/operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
BulkBodyPoolBody,
BulkBodyVariableBody,
BulkResponse,
ClearTaskInstancesBody,
Config,
ConnectionBody,
ConnectionCollectionResponse,
Expand All @@ -59,6 +60,7 @@
ImportErrorCollectionResponse,
ImportErrorResponse,
JobCollectionResponse,
PatchTaskInstanceBody,
PluginCollectionResponse,
PluginImportErrorCollectionResponse,
PoolBody,
Expand All @@ -68,6 +70,8 @@
ProviderCollectionResponse,
QueuedEventCollectionResponse,
QueuedEventResponse,
TaskInstanceCollectionResponse,
TaskInstanceResponse,
TriggerDAGRunPostBody,
VariableBody,
VariableCollectionResponse,
Expand Down Expand Up @@ -950,3 +954,45 @@ def list_import_errors(self) -> PluginImportErrorCollectionResponse | ServerResp
return PluginImportErrorCollectionResponse.model_validate_json(self.response.content)
except ServerResponseError as e:
raise e


class TasksOperations(BaseOperations):
Comment thread
Suraj-kumar00 marked this conversation as resolved.
"""Tasks operations."""

def get(self, dag_id: str, dag_run_id: str, task_id: str) -> TaskInstanceResponse | ServerResponseError:
"""Get a task instance."""
self.response = self.client.get(f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}")
return TaskInstanceResponse.model_validate_json(self.response.content)

def list(self, dag_id: str, dag_run_id: str) -> TaskInstanceCollectionResponse | ServerResponseError:
"""List task instances."""
return super().execute_list(
path=f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
data_model=TaskInstanceCollectionResponse,
)

def clear(
self, dag_id: str, body: ClearTaskInstancesBody
) -> TaskInstanceCollectionResponse | ServerResponseError:
"""Clear task instances."""
self.response = self.client.post(
f"dags/{dag_id}/clearTaskInstances",
json=body.model_dump(mode="json", exclude_unset=True),
)
return TaskInstanceCollectionResponse.model_validate_json(self.response.content)

def update(
self,
dag_id: str,
dag_run_id: str,
task_id: str,
body: PatchTaskInstanceBody,
map_index: int | None = None,
) -> TaskInstanceCollectionResponse | ServerResponseError:
"""Update task instance state. When map_index is given, only that mapped instance is affected."""
if map_index is not None:
path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/{map_index}"
else:
path = f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}"
self.response = self.client.patch(path, json=body.model_dump(mode="json", exclude_unset=True))
return TaskInstanceCollectionResponse.model_validate_json(self.response.content)
51 changes: 46 additions & 5 deletions airflow-ctl/src/airflowctl/ctl/cli_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
import inspect
import os
import sys
import types as builtin_types
import typing
from argparse import Namespace
from collections.abc import Callable, Iterable
from enum import Enum
Expand All @@ -52,6 +54,22 @@
BUILD_DOCS = "BUILDING_AIRFLOW_DOCS" in os.environ


def _is_list_annotation(annotation: Any) -> bool:
"""Check whether a Pydantic field annotation is a list type (including Optional[list[...]])."""
origin = typing.get_origin(annotation)
if origin is list:
return True
# Handle both typing.Union (Optional[list[...]]) and PEP-604 X | Y (types.UnionType)
if origin is typing.Union or isinstance(annotation, builtin_types.UnionType):
return any(_is_list_annotation(arg) for arg in typing.get_args(annotation) if arg is not type(None))
return False


def _get_bool_flag_default(field_info: Any) -> bool:
"""Return the CLI default for a generated bool flag: the datamodel field's default (the API default), or False."""
return field_info.default if isinstance(field_info.default, bool) else False
Comment on lines +68 to +70


def lazy_load_command(import_path: str) -> Callable:
"""Create a lazy loader for command."""
_, _, name = import_path.rpartition(".")
Expand Down Expand Up @@ -399,7 +417,17 @@ def __init__(self, file_path: str | Path | None = None):
# Exclude parameters that are not needed for CLI from datamodels
self.excluded_parameters = ["schema_"]
# This list is used to determine if the command/operation needs to output data
self.output_command_list = ["list", "get", "create", "delete", "update", "trigger", "add", "edit"]
self.output_command_list = [
"list",
"get",
"create",
"delete",
"update",
"trigger",
"add",
"edit",
"clear",
]
self.exclude_operation_names = ["LoginOperations", "VersionOperations", "BaseOperations"]
self.exclude_method_names = [
"error",
Expand All @@ -411,6 +439,8 @@ def __init__(self, file_path: str | Path | None = None):
]
self.excluded_output_keys = [
"total_entries",
"next_cursor",
"previous_cursor",
]

def _inspect_operations(self) -> None:
Expand Down Expand Up @@ -587,7 +617,9 @@ def _create_arg_for_non_primitive_type(
arg_type=self._python_type_from_string(field_type.annotation),
arg_action=argparse.BooleanOptionalAction if field_type.annotation is bool else None, # type: ignore
arg_help=f"{field} for {parameter_key} operation",
arg_default=False if field_type.annotation is bool else None,
arg_default=_get_bool_flag_default(field_type)
if field_type.annotation is bool
else None,
)
)
else:
Expand All @@ -602,7 +634,7 @@ def _create_arg_for_non_primitive_type(
arg_type=self._python_type_from_string(annotation),
arg_action=argparse.BooleanOptionalAction if annotation is bool else None, # type: ignore
arg_help=f"{field} for {parameter_key} operation",
arg_default=False if annotation is bool else None,
arg_default=_get_bool_flag_default(field_type) if annotation is bool else None,
)
)
return commands
Expand Down Expand Up @@ -717,10 +749,19 @@ def _get_func(args: Namespace, api_operation: dict, api_client: Client = NEW_API
datamodel_param_name = parameter_key
if expanded_parameter in self.excluded_parameters:
continue
if expanded_parameter in args_dict.keys():
if (
expanded_parameter in args_dict.keys()
and args_dict[expanded_parameter] is not None
):
val = args_dict[expanded_parameter]
Comment on lines +752 to +756

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I understand correctly, the None filtering is not strictly required for fields where null and “not provided” are equivalent. More importantly, this new path does not protect boolean fields, because the generated CLI currently defaults bool arguments to False. For ClearTaskInstancesBody, some API defaults are True or None (dry_run, only_failed, reset_dag_runs, run_on_latest_version), so airflowctl tasks clear <dag_id> may send explicit False values and override the API defaults.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, you're right on both counts.

Fixed by changing the default from False to None for bool fields in _create_arg_for_non_primitive_type, so unset flags are now omitted from the request body and the API defaults apply. This also brings the path in sync with the primitive arg path, which already defaulted to None.

Added a regression test (test_command_factory_body_bool_field_defaults_to_none) using ClearTaskInstancesBody to lock it in.

if isinstance(val, str) and expanded_parameter in datamodel.model_fields:
if _is_list_annotation(
datamodel.model_fields[expanded_parameter].annotation
):
val = [v.strip() for v in val.split(",") if v.strip()]
method_params[parameter_key][
self._sanitize_method_param_key(expanded_parameter)
] = args_dict[expanded_parameter]
] = val

if datamodel:
if datamodel_param_name:
Expand Down
6 changes: 6 additions & 0 deletions airflow-ctl/src/airflowctl/ctl/help_texts.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -102,3 +102,9 @@ xcom:
plugins:
list: "List all installed Airflow plugins"
list-import-errors: "List all plugin import errors"

tasks:
get: "Retrieve a task instance by Dag ID, run ID, and task ID"
list: "List all task instances for a given Dag run"
clear: "Clear task instance state, optionally filtering by task IDs or run scope"
update: "Update task instance state or note; use --map-index to target a single mapped instance"
Loading
Loading