airflow-api-client 0.1.0__py3-none-any.whl

airflow_api_client/commands/list/dags.py

@@ -0,0 +1,53 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""List DAGs command"""
+
+from airflow_api_client.filters import filter_obj
+from airflow_api_client.io import output
+from airflow_api_client.models import DAGBag
+from airflow_api_client.typedefs import App
+
+
+class ListDagsCommand:
+    """Helper class for listing DAGs"""
+
+    @classmethod
+    def add_subparser(cls, subparsers):
+        """Add subparser for dags"""
+        _ = subparsers.add_parser("dags", help="List all DAGs")
+
+    @staticmethod
+    def list_dags(app: App, add_extra_attrs_fn) -> None:
+        """List all DAGs in Airflow"""
+
+        # Use the lighter API call to get the list of DAGs unless extra details are needed
+        details = any(
+            [app.args.extra_attrs, app.args.filter, app.args.output_format in ["yaml", "json"]]
+        )
+
+        all_dags = [
+            dag
+            for dag in DAGBag(app).get_all_dags(details=details)
+            if filter_obj(dag, app.args.filter)
+        ]
+
+        if app.args.output_format in ("json", "yaml"):
+            output({"dags": [d.to_dict() for d in all_dags]}, app.args)
+            return
+
+        headers = ["DAG ID", "Display Name", "Is Paused", "Owners", "Tags", "Next Dag Run"]
+        data = []
+        for dag in all_dags:
+            attrs = [
+                dag.dag_id,
+                dag.dag_display_name,
+                "Yes" if dag.is_paused else "No",
+                ", ".join(dag.owners),
+                ", ".join([t.name for t in dag.tags]),
+                dag.next_dagrun_logical_date or "N/A",
+            ]
+
+            add_extra_attrs_fn(app, dag, headers, attrs)
+
+            data.append(attrs)
+
+        output(data, app.args, headers=headers)

airflow_api_client/commands/list/main.py

@@ -0,0 +1,39 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""Main list command with subcommand dispatch"""
+
+from airflow_api_client.commands.list.base import ListCommandBase
+from airflow_api_client.commands.list.dag_runs import ListDagRunsCommand
+from airflow_api_client.commands.list.dags import ListDagsCommand
+from airflow_api_client.commands.list.task_instances import ListTaskInstancesCommand
+from airflow_api_client.log import logger
+
+DESCRIPTION = """
+Command to list various resources in Airflow, such as DAGs, DAG runs, and task instances.
+Supports filtering and custom attribute selection for flexible querying.
+"""
+
+
+class ListCommand(ListCommandBase):
+    """Command to list various resources in Airflow"""
+
+    name = "list"
+    help = DESCRIPTION
+
+    @classmethod
+    def add_arguments(cls, parser) -> None:
+        cls.add_common_arguments(parser)
+
+        subparsers = parser.add_subparsers(dest="list_resource", required=True)
+        ListDagsCommand.add_subparser(subparsers)
+        ListDagRunsCommand.add_subparser(subparsers)
+        ListTaskInstancesCommand.add_subparser(subparsers)
+
+    def execute(self, **kwargs) -> None:
+        logger.debug("Executing ListCommand")
+
+        if self.app.args.list_resource == "dags":
+            ListDagsCommand.list_dags(self.app, self._add_extra_attrs)
+        elif self.app.args.list_resource == "dag-runs":
+            ListDagRunsCommand.list_dag_runs(self.app, self._add_extra_attrs)
+        elif self.app.args.list_resource == "task-instances":
+            ListTaskInstancesCommand.list_task_instances(self.app, self._add_extra_attrs)

airflow_api_client/commands/list/task_instances.py

@@ -0,0 +1,56 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""List task instances command"""
+
+from airflow_api_client.filters import filter_obj
+from airflow_api_client.io import output
+from airflow_api_client.models import ResourceSpec
+from airflow_api_client.typedefs import App
+
+
+class ListTaskInstancesCommand:
+    """Helper class for listing task instances"""
+
+    @classmethod
+    def add_subparser(cls, subparsers):
+        """Add subparser for task-instances"""
+        tasks_parser = subparsers.add_parser("task-instances", help="List task instances")
+        tasks_parser.add_argument(
+            "spec",
+            metavar="DAG_ID/RUN_ID[/TASK_ID]",
+            help=(
+                "Retrieve information about the task instance(s) related to the  specified DAG run."
+                " The format should be DAG_ID:DAG_RUN_ID  Optionally you can append :TASK_ID."
+                " to retrieve a specific task instance."
+            ),
+        )
+
+    @staticmethod
+    def list_task_instances(app: App, add_extra_attrs_fn) -> None:
+        """List task instances for a given DAG run"""
+
+        spec = ResourceSpec.from_string(app, app.args.spec)
+
+        task_instances = (
+            [spec.task_instance] if spec.task_instance is not None else spec.dag_run.task_instances
+        )
+        task_instances = [t for t in task_instances if filter_obj(t, app.args.filter)]
+
+        if app.args.output_format in ("json", "yaml"):
+            output({"task_instances": [t.to_dict() for t in task_instances]}, app.args)
+        else:
+            headers = ["Task ID", "State", "Duration", "Try Number", "Start Date", "End Date"]
+            data = []
+            for ti in task_instances:
+                attrs = [
+                    ti.task_id,
+                    ti.state.value,
+                    ti.duration,
+                    ti.try_number,
+                    ti.start_date,
+                    ti.end_date,
+                ]
+
+                add_extra_attrs_fn(app, ti, headers, attrs)
+
+                data.append(attrs)
+            output(data, app.args, headers=headers)

airflow_api_client/commands/pause.py

@@ -0,0 +1,35 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""Code related to the `pause` sub-command."""
+
+from airflow_api_client.commands import Command
+from airflow_api_client.io import output
+from airflow_api_client.log import logger
+from airflow_api_client.models import DAG
+
+DESCRIPTION = """
+Command to pause a DAG. This command allows you to pause a DAG by its ID. When a DAG is paused, it
+will not be scheduled to run until it is resumed. This can be useful for temporarily stopping a DAG
+from running without deleting it or changing its schedule.
+"""
+
+
+class PauseDagCommand(Command):
+    """Command to pause a DAG."""
+
+    name = "pause"
+    help = DESCRIPTION
+
+    @classmethod
+    def add_arguments(cls, parser) -> None:
+        parser.add_argument("dag_id", help="DAG ID to pause")
+
+    def execute(self, **kwargs) -> None:
+        """Pause a DAG"""
+        logger.debug("Executing PauseDagCommand")
+
+        dag = DAG(self.app, self.app.args.dag_id)
+        response = dag.pause()
+        logger.info(f"DAG '{self.app.args.dag_id}' has been paused")
+
+        if self.app.args.output_format in ("json", "yaml"):
+            output(response.to_dict(), self.app.args)

airflow_api_client/commands/profile.py

@@ -0,0 +1,216 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""Code related to the `profile` sub-command."""
+
+import configparser
+import os
+import sys
+from functools import cached_property
+from getpass import getpass
+from pathlib import Path
+
+from airflow_api_client.commands import Command
+from airflow_api_client.io import output
+from airflow_api_client.log import logger
+
+try:
+    import keyring
+except ImportError:
+    keyring = None
+
+DESCRIPTION = """
+Manage configuration profiles. Profiles store connection details for Airflow instances
+in the config file and optionally store passwords securely in the system keyring.
+"""
+
+
+class ProfileCommand(Command):
+    """Command to manage configuration profiles."""
+
+    name = "profile"
+    help = DESCRIPTION
+    no_client = True
+
+    @classmethod
+    def add_arguments(cls, parser) -> None:
+        subparsers = parser.add_subparsers(dest="profile_action", required=True)
+
+        subparsers.add_parser("list", help="List all profiles")
+
+        delete_parser = subparsers.add_parser("delete", help="Delete an existing profile")
+        delete_parser.add_argument(
+            "profile_name",
+            help=("Name of the profile to delete."),
+        )
+        delete_parser.add_argument(
+            "--keep-keyring-password",
+            action="store_true",
+            help="Don't delete password from system keyring",
+        )
+
+        create_parser = subparsers.add_parser("create", help="Create a new minimal profile")
+        create_parser.add_argument(
+            "--profile-name",
+            help=("Name of the profile to create. (default: 'default')"),
+            default="default",
+        )
+        create_parser.add_argument("--endpoint", "-e", required=False, help="Airflow API endpoint")
+        create_parser.add_argument("--username", "-u", help="Airflow username")
+        create_parser.add_argument(
+            "--password",
+            "-p",
+            help=(
+                "Airflow password. If not provided, the password will be read from the"
+                " AIRFLOW_API_CLIENT_PASSWORD environment variable, or prompted for interactively."
+                " When keyring is available, the password is stored securely in the system keyring."
+                " Use --no-keyring-password to store in plaintext in the config file instead."
+            ),
+        )
+        create_parser.add_argument(
+            "--no-keyring-password",
+            action="store_true",
+            help=(
+                "Don't store password in system keyring, even if available (not recommended)."
+                " Providing --no-keyring-password will store the password in plaintext in the"
+                " config file,"
+            ),
+        )
+
+    @cached_property
+    def _config_path(self) -> Path:
+        """Determine the path to the config file based on environment variables and defaults"""
+        default_config_dir = os.getenv("XDG_CONFIG_HOME", str(Path.home() / ".config"))
+        return (
+            self.app.args.config_file
+            or Path(default_config_dir) / "airflow-api-client" / "config.ini"
+        )
+
+    def execute(self, **kwargs) -> None:
+        """Execute profile subcommand"""
+        if self.app.args.profile_action == "list":
+            self._list()
+        elif self.app.args.profile_action == "create":
+            self._create()
+        elif self.app.args.profile_action == "delete":
+            self._delete()
+
+    def get_password(self, args, endpoint, username) -> str:
+        """Get password from command line arguments or environment variable, or fall-back to
+        prompting the user"""
+        if args.password:
+            return args.password
+
+        env_password = os.getenv("AIRFLOW_API_CLIENT_PASSWORD")
+        if env_password:
+            return env_password
+
+        if sys.stdin.isatty():
+            return getpass(f"Enter Airflow password for {username} at {endpoint}: ")
+
+        return sys.stdin.readline().strip()
+
+    def _list(self) -> None:
+        """List all profiles in the config file"""
+        if not self._config_path.is_file():
+            raise RuntimeError(f"No config file found at {self._config_path}")
+
+        config = configparser.RawConfigParser()
+        config.read(self._config_path)
+
+        sections = config.sections()
+        if not sections:
+            print("No profiles found.")
+            return
+
+        headers = {"profile": "Profile", "endpoint": "Endpoint", "username": "Username"}
+        data = []
+
+        for section in sections:
+            if keyring and not config.has_option(section, "password"):
+                password = "(stored in keyring)"
+            elif config.has_option(section, "password"):
+                password = "(stored in config)"
+            else:
+                password = "(not set)"
+
+            data.append(
+                {
+                    "profile": section,
+                    "endpoint": config.get(section, "endpoint", fallback="(not set)"),
+                    "username": config.get(section, "username", fallback="(not set)"),
+                    "verbose": config.get(section, "verbose", fallback="default"),
+                    "password": password,
+                    "auth_provider": config.get(section, "auth_provider", fallback="default"),
+                }
+            )
+
+        output(data, self.app.args, headers=headers)
+
+    def _delete(self) -> None:
+        """Delete an existing profile"""
+        if not self._config_path.is_file():
+            raise RuntimeError(f"No config file found at {self._config_path}")
+
+        config = configparser.RawConfigParser()
+        config.read(self._config_path)
+        profile = self.app.args.profile_name
+
+        if not config.has_section(profile):
+            raise RuntimeError(f"Profile '{profile}' not found in config file")
+
+        if keyring and not self.app.args.keep_keyring_password:
+            try:
+                keyring.delete_password(
+                    f"airflow-api-client/{config.get(profile, 'endpoint')}",
+                    config.get(profile, "username"),
+                )
+                logger.info("Password removed from system keyring")
+            except keyring.errors.PasswordDeleteError:
+                logger.debug("No keyring entry found for this profile")
+
+        config.remove_section(profile)
+        with self._config_path.open("w", encoding="utf-8") as f:
+            config.write(f)
+
+        logger.info(f"Profile '{profile}' deleted from {self._config_path}")
+
+    def _create(self) -> None:
+        """Create a new profile"""
+        self._config_path.parent.mkdir(parents=True, exist_ok=True)
+
+        config = configparser.RawConfigParser()
+        if self._config_path.is_file():
+            config.read(self._config_path)
+
+        profile = self.app.args.profile_name
+
+        if config.has_section(profile):
+            logger.warning(f"Profile '{profile}' already exists, overwriting")
+        else:
+            config.add_section(profile)
+
+        if (not self.app.args.endpoint or not self.app.args.username) and not sys.stdin.isatty():
+            raise RuntimeError(
+                "Endpoint and username must be provided as arguments when running in"
+                " non-interactive mode. Use --endpoint and --username options or run the command in"
+                " an interactive terminal."
+            )
+
+        endpoint = self.app.args.endpoint or input("Enter Airflow API endpoint: ")
+        username = self.app.args.username or input("Enter Airflow username: ")
+
+        endpoint = endpoint.strip().rstrip("/")
+        config.set(profile, "endpoint", endpoint)
+        config.set(profile, "username", username)
+
+        password = self.get_password(self.app.args, endpoint, username)
+
+        if keyring and not self.app.args.no_keyring_password:
+            keyring.set_password(f"airflow-api-client/{endpoint}", username, password)
+            logger.info("Password stored in system keyring")
+        else:
+            config.set(profile, "password", password)
+
+        with self._config_path.open("w", encoding="utf-8") as f:
+            config.write(f)
+
+        logger.info(f"Profile '{profile}' saved to {self._config_path}")

airflow_api_client/commands/raw.py

@@ -0,0 +1,99 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""Code related to the `raw` sub-command."""
+
+import json
+
+from airflow_api_client.commands import Command
+from airflow_api_client.io import output
+from airflow_api_client.log import logger
+
+DESCRIPTION = """
+Make a raw API request to Airflow. This command allows you to interact with the Airflow API
+directly, bypassing any predefined commands or logic. You can specify the endpoint, method, and
+payload for your request, giving you full control over the API interaction.
+"""
+
+
+class RawCommand(Command):
+    """Command to make a raw API request."""
+
+    name = "raw"
+    help = DESCRIPTION
+
+    @classmethod
+    def add_arguments(cls, parser) -> None:
+        """Add arguments for the `raw` command."""
+        parser.add_argument(
+            "method",
+            type=str.upper,
+            choices=["GET", "POST", "PUT", "DELETE", "PATCH"],
+            help="The HTTP method to use for the request",
+        )
+        parser.add_argument(
+            "stub_url",
+            type=str,
+            help="The API endpoint to call (e.g., /dags/{dag_id}/pause)",
+        )
+        parser.add_argument(
+            "--payload",
+            type=str,
+            default=None,
+            help="The JSON payload to send with the request (for POST and PUT methods)",
+        )
+        parser.add_argument(
+            "--header",
+            action="append",
+            help=(
+                "Custom headers to include in the request (can be specified multiple times, e.g.,"
+                " --header 'Foo=Bar' --header 'Baz=Qux')"
+            ),
+        )
+
+    def execute(self, **kwargs) -> None:
+        """Run a raw API request to Airflow based on the provided arguments."""
+
+        logger.debug("Executing RawCommand")
+
+        endpoint = f"{self.app.args.endpoint}/{self.app.args.stub_url.lstrip('/')}"
+
+        headers = {
+            "Authorization": f"Bearer {self.app.client.auth.client_token.token}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+
+        headers.update(dict([h.split("=", maxsplit=1) for h in self.app.args.header or []]))
+        try:
+            payload = json.loads(self.app.args.payload) if self.app.args.payload else None
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON payload: {e}")
+            return
+
+        logger.debug(
+            f"Making {self.app.args.method} request to {endpoint}",
+            headers=headers,
+            payload=payload,
+            timeout=self.app.args.timeout,
+        )
+
+        response = self.app.client.rest_client.request(
+            self.app.args.method,
+            endpoint,
+            body=payload,
+            headers=headers,
+            _request_timeout=self.app.args.timeout,
+        )
+
+        if self.app.args.output_format not in ["json", "yaml"]:
+            # There's no pretty-printing arbitrary output for this command, so we default to JSON if
+            # the user specified an unsupported format
+            self.app.args.output_format = "json"
+
+        logger.info(f"Received response from Airflow API. Status code: {response.status}")
+        response_b = response.response.read()
+
+        try:
+            output(json.loads(response_b), self.app.args)
+        except json.JSONDecodeError:
+            logger.debug("Response is not valid JSON, printing raw response")
+            output(response_b, self.app.args)

airflow_api_client/commands/run.py

@@ -0,0 +1,110 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""Code related to the `run` sub-command."""
+
+import json
+import time
+from pathlib import Path
+
+from airflow_api_client.commands import Command
+from airflow_api_client.io import output
+from airflow_api_client.log import logger
+from airflow_api_client.models import DAG, TERMINAL_DAGRUN_STATES, DAGRun
+
+DESCRIPTION = """
+Command to trigger a DAG run with an optional payload. The response is printed to the console
+in the specified output format (JSON, YAML, or human-readable). If a payload is provided, it is
+used as the `conf` for the DAG run. If no payload is provided, the user is prompted to enter
+values for each parameter defined in the DAG's parameters. After triggering the DAG, the command
+polls for completion and retrieves any XCom entries associated with the DAG run.
+"""
+
+
+class RunDagCommand(Command):
+    """Command to trigger a DAG run with an optional payload."""
+
+    name = "run"
+    help = DESCRIPTION
+
+    @classmethod
+    def add_arguments(cls, parser) -> None:
+        parser.add_argument("dag_id", help="DAG ID to run")
+        parser.add_argument("--conf", "-c", type=str, help="JSON conf or @file")
+        parser.add_argument(
+            "--run-id",
+            type=str,
+            help=(
+                "Optional run_id for the DAG run (default: autogenerated). If"
+                " provided, it must be unique across all DAG runs. "
+            ),
+        )
+        parser.add_argument(
+            "--no-wait",
+            action="store_true",
+            help="Don't wait for DAG run to complete. Trigger and forget",
+        )
+
+    def execute(self, **kwargs) -> None:
+        """Triggers a DAG with an optional payload. The response is printed to the console in the
+        specified output format (JSON, YAML, or human-readable). If a payload is provided, it is
+        used as the `conf` for the DAG run. If no payload is provided, the user is prompted to
+        enter values for each parameter defined in the DAG's parameters. After triggering the DAG,
+        the function polls for completion and retrieves any XCom entries associated with the
+        DAG run.
+        """
+        logger.debug("Executing RunDagCommand")
+
+        conf = self.app.args.conf
+
+        dag = DAG(self.app, self.app.args.dag_id)
+
+        if conf is None:
+            conf = {"conf": dag.create_payload()}
+        else:
+            if conf.startswith("@"):
+                conf_path = Path(conf[1:]).absolute()
+                if not conf_path.is_file():
+                    logger.error("Conf file not found", path=str(conf_path))
+                    raise FileNotFoundError(f"Conf file not found: {conf_path}")
+                with conf_path.open("r", encoding="utf-8") as f:
+                    conf = f.read()
+            conf = {"conf": json.loads(conf)}
+
+        try:
+            conf = json.loads(conf) if isinstance(conf, str) else conf
+        except json.JSONDecodeError as e:
+            logger.error("Invalid JSON provided for --conf", error=str(e))
+            raise
+
+        dag_run: DAGRun = dag.run(conf=conf, run_id=self.app.args.run_id)
+
+        if self.app.args.no_wait:
+            return
+
+        # Poll for completion. Prefer this over calling airflow_client's wait_dag_run_until_finished
+        # until that's no longer marked as experimental.
+        while (current_state := dag_run.state) not in TERMINAL_DAGRUN_STATES:
+            logger.debug("Polling for DAG run completion", dag=dag, state=current_state)
+            time.sleep(2)
+
+        logger.info("DAG run completed", dag=dag, final_state=dag_run.state)
+
+        if self.app.args.output_format in ("json", "yaml"):
+            result = dag_run.to_dict()
+            result["task_instances"] = {}
+            for task_instance in dag_run.task_instances:
+                result["task_instances"][task_instance.task_id] = {
+                    "state": task_instance.state,
+                    "logs": list(task_instance.get_log()),
+                    "xcom": task_instance.get_xcom_entries().to_dict(),
+                }
+
+            output(result, self.app.args)
+            return
+
+        output([["Logs"]], self.app.args, headers=[])
+        for task_instance in dag_run.task_instances:
+            task_instance.get_log().output()
+
+        output("\n", self.app.args)
+        output([["XCom results"]], self.app.args, headers=[])
+        dag_run.xcoms.output()

airflow_api_client/commands/unpause.py

@@ -0,0 +1,33 @@
+# Copyright 2026 Rackspace Technology, Inc. SPDX-License-Identifier: Apache-2.0
+"""Code related to the `unpause` sub-command."""
+
+from airflow_api_client.commands import Command
+from airflow_api_client.io import output
+from airflow_api_client.log import logger
+from airflow_api_client.models import DAG
+
+DESCRIPTION = """
+Command to unpause a DAG. This command allows you to unpause a DAG by its ID.
+"""
+
+
+class UnpauseDagCommand(Command):
+    """Command to unpause a DAG."""
+
+    name = "unpause"
+    help = DESCRIPTION
+
+    @classmethod
+    def add_arguments(cls, parser) -> None:
+        parser.add_argument("dag_id", help="DAG ID to unpause")
+
+    def execute(self, **kwargs) -> None:
+        """Unpause a DAG"""
+        logger.debug("Executing UnpauseDagCommand")
+
+        dag = DAG(self.app, self.app.args.dag_id)
+        response = dag.unpause()
+        logger.info(f"DAG '{self.app.args.dag_id}' has been unpaused.")
+
+        if self.app.args.output_format in ("json", "yaml"):
+            output(response.to_dict(), self.app.args)