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
44 changes: 31 additions & 13 deletions airflow-ctl/src/airflowctl/api/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ class ClientKind(enum.Enum):

CLI = "cli"
AUTH = "auth"
NO_AUTH = "no_auth"


def add_correlation_id(request: httpx.Request):
Expand Down Expand Up @@ -253,6 +254,8 @@ def load(self) -> Credentials:
with open(config_path) as f:
credentials = json.load(f)
self.api_url = credentials["api_url"]
if self.client_kind == ClientKind.NO_AUTH:
return self
if self.api_token is not None:
return self
if os.getenv("AIRFLOW_CLI_DEBUG_MODE") == "true":
Expand Down Expand Up @@ -294,16 +297,17 @@ def load(self) -> Credentials:
raise AirflowCtlKeyringException("Keyring backend is not available") from e
self.api_token = None
except FileNotFoundError:
# This is expected during the auth login command
if self.client_kind != ClientKind.AUTH:
# This is expected during the auth login command.
# Also allow token-only usage without local config (for commands like `version --remote`).
if self.client_kind not in (ClientKind.AUTH, ClientKind.NO_AUTH) and self.api_token is None:
raise AirflowCtlCredentialNotFoundException("No credentials file found. Please login first.")

return self


class BearerAuth(httpx.Auth):
def __init__(self, token: str):
self.token: str = token
def __init__(self, token: str | None):
self.token: str | None = token

def auth_flow(self, request: httpx.Request):
if self.token:
Expand Down Expand Up @@ -332,11 +336,13 @@ def __init__(
self,
*,
base_url: str,
token: str,
kind: Literal[ClientKind.CLI, ClientKind.AUTH] = ClientKind.CLI,
token: str | None = None,
kind: Literal[ClientKind.CLI, ClientKind.AUTH, ClientKind.NO_AUTH] = ClientKind.CLI,
**kwargs: Any,
) -> None:
auth = BearerAuth(token)
auth: httpx.Auth | None = None
if kind != ClientKind.NO_AUTH:
auth = BearerAuth(token)
kwargs["base_url"] = self._get_base_url(base_url=base_url, kind=kind)
pyver = f"{'.'.join(map(str, sys.version_info[:3]))}"
super().__init__(
Expand All @@ -347,14 +353,18 @@ def __init__(
)

def refresh_base_url(
self, base_url: str, kind: Literal[ClientKind.AUTH, ClientKind.CLI] = ClientKind.CLI
self,
base_url: str,
kind: Literal[ClientKind.AUTH, ClientKind.CLI, ClientKind.NO_AUTH] = ClientKind.CLI,
):
"""Refresh the base URL of the client."""
self.base_url = URL(self._get_base_url(base_url=base_url, kind=kind))

@classmethod
def _get_base_url(
cls, base_url: str, kind: Literal[ClientKind.AUTH, ClientKind.CLI] = ClientKind.CLI
cls,
base_url: str,
kind: Literal[ClientKind.AUTH, ClientKind.CLI, ClientKind.NO_AUTH] = ClientKind.CLI,
) -> str:
"""Get the base URL of the client."""
base_url = base_url.rstrip("/")
Expand Down Expand Up @@ -466,7 +476,10 @@ def plugins(self):

# API Client Decorator for CLI Actions
@contextlib.contextmanager
def get_client(kind: Literal[ClientKind.CLI, ClientKind.AUTH] = ClientKind.CLI, api_token: str | None = None):
def get_client(
kind: Literal[ClientKind.CLI, ClientKind.AUTH, ClientKind.NO_AUTH] = ClientKind.CLI,
api_token: str | None = None,
):
"""
Get CLI API client.

Expand All @@ -476,11 +489,16 @@ def get_client(kind: Literal[ClientKind.CLI, ClientKind.AUTH] = ClientKind.CLI,
api_token = api_token or os.getenv("AIRFLOW_CLI_TOKEN", None)
try:
# API URL always loaded from the config file, please save with it if you are using other than ClientKind.CLI
credentials = Credentials(client_kind=kind, api_token=api_token).load()
if kind == ClientKind.NO_AUTH:
credentials = Credentials(client_kind=kind).load()
resolved_token = None
else:
credentials = Credentials(client_kind=kind, api_token=api_token).load()
resolved_token = api_token or credentials.api_token
api_client = Client(
base_url=credentials.api_url or "http://localhost:8080",
limits=httpx.Limits(max_keepalive_connections=1, max_connections=1),
token=str(api_token or credentials.api_token),
token=resolved_token,
kind=kind,
)
yield api_client
Expand All @@ -492,7 +510,7 @@ def get_client(kind: Literal[ClientKind.CLI, ClientKind.AUTH] = ClientKind.CLI,


def provide_api_client(
kind: Literal[ClientKind.CLI, ClientKind.AUTH] = ClientKind.CLI,
kind: Literal[ClientKind.CLI, ClientKind.AUTH, ClientKind.NO_AUTH] = ClientKind.CLI,
) -> Callable[[Callable[PS, RT]], Callable[PS, RT]]:
"""
Provide a CLI API Client to the decorated function.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def version_info(arg):
"""Get version information."""
version_dict = {"airflowctl_version": airflowctl_version}
if arg.remote:
with get_client(kind=ClientKind.CLI, api_token=getattr(arg, "api_token", None)) as api_client:
with get_client(kind=ClientKind.NO_AUTH) as api_client:
version_response = api_client.version.get()
version_dict.update(version_response.model_dump())
rich.print(version_dict)
40 changes: 39 additions & 1 deletion airflow-ctl/tests/airflow_ctl/api/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
import time_machine
from httpx import URL

from airflowctl.api.client import Client, ClientKind, Credentials, _bounded_get_new_password
from airflowctl.api.client import Client, ClientKind, Credentials, _bounded_get_new_password, get_client
from airflowctl.api.operations import ServerResponseError
from airflowctl.exceptions import (
AirflowCtlCredentialNotFoundException,
Expand Down Expand Up @@ -105,12 +105,16 @@ def handle_request(request: httpx.Request) -> httpx.Response:
[
("http://localhost:8080", ClientKind.CLI, "http://localhost:8080/api/v2/"),
("http://localhost:8080", ClientKind.AUTH, "http://localhost:8080/auth/"),
("http://localhost:8080", ClientKind.NO_AUTH, "http://localhost:8080/api/v2/"),
("https://example.com", ClientKind.CLI, "https://example.com/api/v2/"),
("https://example.com", ClientKind.AUTH, "https://example.com/auth/"),
("https://example.com", ClientKind.NO_AUTH, "https://example.com/api/v2/"),
("http://localhost:8080/", ClientKind.CLI, "http://localhost:8080/api/v2/"),
("http://localhost:8080/", ClientKind.AUTH, "http://localhost:8080/auth/"),
("http://localhost:8080/", ClientKind.NO_AUTH, "http://localhost:8080/api/v2/"),
("https://example.com/", ClientKind.CLI, "https://example.com/api/v2/"),
("https://example.com/", ClientKind.AUTH, "https://example.com/auth/"),
("https://example.com/", ClientKind.NO_AUTH, "https://example.com/api/v2/"),
],
)
def test_refresh_base_url(self, base_url, client_kind, expected_base_url):
Expand Down Expand Up @@ -214,6 +218,25 @@ def test_load_no_credentials(self, mock_keyring):

assert not os.path.exists(config_dir)

@patch.dict(os.environ, {"AIRFLOW_CLI_ENVIRONMENT": "TEST_NO_CONFIG_WITH_EXPLICIT_TOKEN"})
@patch("airflowctl.api.client.keyring")
def test_load_no_config_with_explicit_token(self, mock_keyring):
cli_client = ClientKind.CLI
credentials = Credentials(client_kind=cli_client, api_token="TEST_TOKEN").load()

assert credentials.api_token == "TEST_TOKEN"
assert credentials.api_url is None
mock_keyring.get_password.assert_not_called()

@patch.dict(os.environ, {"AIRFLOW_CLI_ENVIRONMENT": "TEST_NO_CONFIG_NO_AUTH"})
@patch("airflowctl.api.client.keyring")
def test_load_no_config_no_auth_kind(self, mock_keyring):
credentials = Credentials(client_kind=ClientKind.NO_AUTH).load()

assert credentials.api_token is None
assert credentials.api_url is None
mock_keyring.get_password.assert_not_called()

@patch.dict(os.environ, {"AIRFLOW_CLI_ENVIRONMENT": "TEST_KEYRING_VALUE_ERROR"})
@patch("airflowctl.api.client.keyring")
def test_load_incorrect_keyring_password(self, mock_keyring):
Expand Down Expand Up @@ -400,6 +423,21 @@ def test_credentials_accepts_safe_env():
assert creds.api_environment == "prod-us_1"


@patch.dict(os.environ, {"AIRFLOW_CLI_ENVIRONMENT": "TEST_GET_CLIENT_WITH_TOKEN_ONLY"})
def test_get_client_allows_explicit_token_without_config():
with get_client(kind=ClientKind.CLI, api_token="TEST_TOKEN") as client:
assert str(client.base_url) == "http://localhost:8080/api/v2/"


@patch.dict(os.environ, {"AIRFLOW_CLI_ENVIRONMENT": "TEST_GET_CLIENT_NO_AUTH_WITHOUT_CONFIG"})
@patch("airflowctl.api.client.keyring")
def test_get_client_no_auth_without_config(mock_keyring):
with get_client(kind=ClientKind.NO_AUTH) as client:
assert str(client.base_url) == "http://localhost:8080/api/v2/"

mock_keyring.get_password.assert_not_called()


@pytest.mark.parametrize("api_environment", ["../evil", "..\\evil", "a/b", "a\\b"])
def test_credentials_rejects_unsafe_env_argument(api_environment):
with pytest.raises(AirflowCtlException, match="environment"):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@

import pytest

from airflowctl.api.client import Client
from airflowctl.api.client import Client, ClientKind
from airflowctl.ctl import cli_parser
from airflowctl.ctl.commands.version_command import version_info

Expand Down Expand Up @@ -52,6 +52,13 @@ def test_ctl_version_remote(self, mock_client):
assert "version" in stdout.getvalue()
assert "git_version" in stdout.getvalue()
assert "airflowctl_version" in stdout.getvalue()
mock_get_client.assert_called_once_with(kind=ClientKind.NO_AUTH)

def test_ctl_version_remote_with_api_token(self, mock_client):
with mock.patch("airflowctl.ctl.commands.version_command.get_client") as mock_get_client:
mock_get_client.return_value.__enter__.return_value = mock_client
version_info(self.parser.parse_args(["version", "--remote", "--api-token", "TOKEN"]))
mock_get_client.assert_called_once_with(kind=ClientKind.NO_AUTH)

def test_ctl_version_only_local_version(self, mock_client):
"""Test the version command without --remote does not touch credentials."""
Expand Down
Loading