Repository: apache/airflow
File: airflow-ctl/src/airflowctl/ctl/commands/pool_command.py (Python)

Code before:

import rich

from airflowctl.api.client import NEW_API_CLIENT, Client, ClientKind, provide_api_client
from airflowctl.api.datamodels.generated import (
    BulkActionOnExistence,
    BulkBodyPoolBody,
    BulkCreateActionPoolBody,
    PoolBody,
)


@provide_api_client(kind=ClientKind.CLI)
def import_(args, api_client: Client = NEW_API_CLIENT) -> None:
    """Import pools from file."""
    filepath = Path(args.file)
    if not filepath.exists():
        raise SystemExit(f"Missing pools file {args.file}")

    success, errors = _import_helper(api_client, filepath)
    if errors:
        raise SystemExit(f"Failed to update pool(s): {errors}")
    rich.print(success)


def _import_helper(api_client: Client, filepath: Path):
    """Help import pools from the json file."""
    try:
        with open(filepath) as f:
            pools_json = json.load(f)
    except JSONDecodeError as e:
        raise SystemExit(f"Invalid json file: {e}")

    if not isinstance(pools_json, list):
        raise SystemExit("Invalid format: Expected a list of pool objects")

    pools_to_update = []
    for pool_config in pools_json:
        if not isinstance(pool_config, dict) or "name" not in pool_config or "slots" not in pool_config:
            raise SystemExit(f"Invalid pool configuration: {pool_config}")

        pools_to_update.append(
            PoolBody(
                name=pool_config["name"],
                slots=pool_config["slots"],
                description=pool_config.get("description", ""),
                include_deferred=pool_config.get("include_deferred", False),
            )
        )

    bulk_body = BulkBodyPoolBody(

Review comment (security):
Ok. "all" was a bit too much. Pool export does bit contain sensitive data, so we should not remove it. 

Code after:

import rich

from airflowctl.api.client import NEW_API_CLIENT, Client, ClientKind, provide_api_client
from airflowctl.api.datamodels.generated import (
    BulkActionOnExistence,
    BulkBodyPoolBody,
    BulkCreateActionPoolBody,
    PoolBody,
)


@provide_api_client(kind=ClientKind.CLI)
def import_(args, api_client: Client = NEW_API_CLIENT) -> None:
    """Import pools from file."""
    filepath = Path(args.file)
    if not filepath.exists():
        raise SystemExit(f"Missing pools file {args.file}")

    success, errors = _import_helper(api_client, filepath)
    if errors:
        raise SystemExit(f"Failed to update pool(s): {errors}")
    rich.print(success)


@provide_api_client(kind=ClientKind.CLI)
def export(args, api_client: Client = NEW_API_CLIENT) -> None:
    """
    Export all pools.

    If output is json, write to file. Otherwise, print to console.
    """
    try:
        pools_response = api_client.pools.list()
        pools_list = [
            {
                "name": pool.name,
                "slots": pool.slots,
                "description": pool.description,
                "include_deferred": pool.include_deferred,
                "occupied_slots": pool.occupied_slots,
                "running_slots": pool.running_slots,
                "queued_slots": pool.queued_slots,
                "scheduled_slots": pool.scheduled_slots,
                "open_slots": pool.open_slots,
                "deferred_slots": pool.deferred_slots,
            }
            for pool in pools_response.pools
        ]

        if args.output == "json":