From e41ed31c9890ad7ee61c7116e18c674de79cfc5b Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 24 Aug 2026 21:16:57 -0400 Subject: [PATCH 1/4] feat: add OAuth auth command with login, status, refresh and auto-401 retry Add `dremio auth` command group with subcommands: - `login`: browser-based OAuth 2.0 PKCE flow that stores access + refresh tokens in ~/.config/dremioai/config.yaml - `status`: shows current auth method, token expiry (JWT decode), redacted token values - `refresh`: manually refreshes the access token using stored refresh token - `logout`: removes OAuth tokens from config OAuth configuration: - Client ID defaults to https://connectors.dremio.app/claude - Redirect URI: http://localhost:{port}/Callback - Scopes: dremio.all offline_access - Endpoints discovered via .well-known/oauth-authorization-server with hardcoded fallbacks for US/EU regions using /oauth/authorize path Automatic 401 token refresh in DremioClient: - On any 401 response, if OAuth refresh_token + client_id are available, the client transparently refreshes the access token and retries once - New tokens are persisted to disk and httpx headers updated in-place - Guard flag prevents infinite refresh loops (one attempt per client instance) Auth priority chain (unchanged): CLI --token > DREMIO_TOKEN env > oauth.access_token from config > pat field from config New modules: - src/drs/oauth.py: PKCE flow, well-known discovery, token refresh - src/drs/commands/auth.py: typer command group for auth subcommands Tests: 4 new tests covering 401 refresh trigger, no-refresh without OAuth config, single-attempt guard, and refresh failure propagation. --- src/drs/auth.py | 82 +++++++++- src/drs/cli.py | 2 + src/drs/client.py | 66 +++++++- src/drs/commands/auth.py | 325 +++++++++++++++++++++++++++++++++++++ src/drs/oauth.py | 272 +++++++++++++++++++++++++++++++ tests/test_client_retry.py | 107 ++++++++++++ 6 files changed, 850 insertions(+), 4 deletions(-) create mode 100644 src/drs/commands/auth.py create mode 100644 src/drs/oauth.py diff --git a/src/drs/auth.py b/src/drs/auth.py index 051f365..e85bada 100644 --- a/src/drs/auth.py +++ b/src/drs/auth.py @@ -17,6 +17,7 @@ from __future__ import annotations +import logging import os from pathlib import Path from typing import Any @@ -24,14 +25,25 @@ import yaml from pydantic import BaseModel +logger = logging.getLogger(__name__) + DEFAULT_CONFIG_PATH = Path.home() / ".config" / "dremioai" / "config.yaml" DEFAULT_URI = "https://api.dremio.cloud" +class OAuthConfig(BaseModel): + """OAuth token state stored in the config file.""" + + access_token: str | None = None + refresh_token: str | None = None + client_id: str | None = None + + class DrsConfig(BaseModel): uri: str = DEFAULT_URI pat: str project_id: str + oauth: OAuthConfig | None = None def load_config( @@ -46,10 +58,12 @@ def load_config( Authentication priority: 1. --token CLI arg 2. DREMIO_TOKEN / DREMIO_PAT env var - 3. Config file pat/token field + 3. Config file oauth.access_token (from OAuth login) + 4. Config file pat/token field """ # -- Config file (lowest priority) -- file_values: dict[str, Any] = {} + oauth_config: OAuthConfig | None = None path = config_path or DEFAULT_CONFIG_PATH if path.exists(): with path.open() as f: @@ -61,6 +75,17 @@ def load_config( } file_values = {k: v for k, v in file_values.items() if v is not None} + # Load OAuth section if present + if "oauth" in raw and isinstance(raw["oauth"], dict): + oauth_config = OAuthConfig( + access_token=raw["oauth"].get("access_token"), + refresh_token=raw["oauth"].get("refresh_token"), + client_id=raw["oauth"].get("client_id"), + ) + # Use OAuth access_token as the PAT if no explicit PAT is set + if oauth_config.access_token and "pat" not in file_values: + file_values["pat"] = oauth_config.access_token + # -- Env vars (override file) -- env_values: dict[str, Any] = {} if v := os.environ.get("DREMIO_URI"): @@ -83,4 +108,57 @@ def load_config( if cli_token: merged["pat"] = cli_token - return DrsConfig(**merged) + config = DrsConfig(**merged) + config.oauth = oauth_config + return config + + +def save_oauth_tokens( + access_token: str, + refresh_token: str | None, + client_id: str, + config_path: Path | None = None, +) -> None: + """Write OAuth tokens to the config file (preserves other fields).""" + path = config_path or DEFAULT_CONFIG_PATH + path.parent.mkdir(parents=True, exist_ok=True) + + # Read existing config + raw: dict[str, Any] = {} + if path.exists(): + with path.open() as f: + raw = yaml.safe_load(f) or {} + + # Update oauth section + raw["oauth"] = { + "access_token": access_token, + "refresh_token": refresh_token, + "client_id": client_id, + } + + # Write back + header = "# Dremio CLI config — generated by 'dremio setup' / 'dremio auth login'\n" + path.write_text(header + yaml.dump(raw, default_flow_style=False, sort_keys=False)) + path.chmod(0o600) + logger.debug("OAuth tokens saved to %s", path) + + +def clear_oauth_tokens(config_path: Path | None = None) -> None: + """Remove OAuth tokens from the config file.""" + path = config_path or DEFAULT_CONFIG_PATH + if not path.exists(): + return + + with path.open() as f: + raw = yaml.safe_load(f) or {} + + if "oauth" in raw: + del raw["oauth"] + header = "# Dremio CLI config — generated by 'dremio setup' / 'dremio auth login'\n" + path.write_text(header + yaml.dump(raw, default_flow_style=False, sort_keys=False)) + path.chmod(0o600) + + +def get_config_path_from_context(config_path: Path | None = None) -> Path: + """Return the effective config path.""" + return config_path or DEFAULT_CONFIG_PATH diff --git a/src/drs/cli.py b/src/drs/cli.py index 5ab21da..6e2f171 100644 --- a/src/drs/cli.py +++ b/src/drs/cli.py @@ -30,6 +30,7 @@ from drs.auth import DrsConfig, load_config from drs.client import DremioClient from drs.commands import ( + auth, chat, engine, folder, @@ -57,6 +58,7 @@ ) # Register command groups +app.add_typer(auth.app, name="auth") app.add_typer(query.app, name="query") app.add_typer(folder.app, name="folder") app.add_typer(schema.app, name="schema") diff --git a/src/drs/client.py b/src/drs/client.py index 5762169..6bd4b20 100644 --- a/src/drs/client.py +++ b/src/drs/client.py @@ -24,7 +24,7 @@ import httpx from drs import __version__ -from drs.auth import DrsConfig +from drs.auth import DrsConfig, save_oauth_tokens logger = logging.getLogger(__name__) @@ -41,6 +41,9 @@ class DremioClient: Transient failures (timeouts, 429, 502, 503, 504) are retried up to 3 times with exponential backoff (1s, 2s, 4s). + + 401 responses trigger an automatic token refresh using the stored OAuth + refresh token, then retry the request once with the new access token. """ def __init__(self, config: DrsConfig) -> None: @@ -53,10 +56,59 @@ def __init__(self, config: DrsConfig) -> None: }, timeout=120.0, ) + self._refreshed = False # guard against infinite refresh loops async def close(self) -> None: await self._client.aclose() + # -- OAuth 401 auto-refresh -- + + def _try_refresh_token(self) -> bool: + """Attempt to refresh the OAuth access token synchronously. + + Returns True if the token was refreshed and the client headers updated. + """ + if self._refreshed: + return False # already tried once this session + + oauth = self.config.oauth + if not oauth or not oauth.refresh_token or not oauth.client_id: + return False + + from drs.oauth import discover_oauth_metadata, do_token_refresh + + try: + metadata = discover_oauth_metadata(self.config.uri) + result = do_token_refresh( + metadata.token_endpoint, oauth.client_id, oauth.refresh_token + ) + except Exception as exc: + logger.warning("OAuth token refresh failed: %s", exc) + return False + + if result is None: + logger.warning("OAuth token refresh returned no tokens") + return False + + # Update in-memory state + self.config.pat = result.access_token + oauth.access_token = result.access_token + if result.refresh_token: + oauth.refresh_token = result.refresh_token + + # Persist to config file + save_oauth_tokens( + access_token=result.access_token, + refresh_token=result.refresh_token or oauth.refresh_token, + client_id=oauth.client_id, + ) + + # Update httpx client headers + self._client.headers["Authorization"] = f"Bearer {result.access_token}" + self._refreshed = True + logger.info("OAuth token refreshed successfully") + return True + # -- URL builders -- def _v0(self, path: str) -> str: @@ -73,11 +125,21 @@ def _v1(self, path: str) -> str: # -- HTTP helpers with retry -- async def _request_with_retry(self, method: str, url: str, **kwargs: Any) -> httpx.Response: - """Execute an HTTP request with retry on transient errors.""" + """Execute an HTTP request with retry on transient errors. + + Also handles 401 by attempting an OAuth token refresh once. + """ last_exc: Exception | None = None for attempt in range(_MAX_RETRIES): try: resp = await self._client.request(method, url, **kwargs) + + # 401 — attempt token refresh and retry once + if resp.status_code == 401 and self._try_refresh_token(): + logger.info("Retrying %s %s after token refresh", method, url) + resp = await self._client.request(method, url, **kwargs) + return resp + if resp.status_code in _RETRYABLE_STATUS_CODES and attempt < _MAX_RETRIES - 1: delay = _RETRY_BACKOFF[attempt] logger.warning( diff --git a/src/drs/commands/auth.py b/src/drs/commands/auth.py new file mode 100644 index 0000000..0eb573b --- /dev/null +++ b/src/drs/commands/auth.py @@ -0,0 +1,325 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""drs auth — OAuth login, status, and token refresh for Dremio Cloud.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import typer +from rich.console import Console +from rich.panel import Panel +from rich.table import Table + +from drs.auth import ( + DEFAULT_CONFIG_PATH, + DEFAULT_URI, + clear_oauth_tokens, + save_oauth_tokens, +) +from drs.oauth import ( + DEFAULT_CLIENT_ID, + DEFAULT_REDIRECT_PORT, + discover_oauth_metadata, + do_token_refresh, + run_oauth_flow, +) + +app = typer.Typer( + help="Authenticate with Dremio Cloud via OAuth.", + context_settings={"help_option_names": ["-h", "--help"]}, +) + +console = Console() +err_console = Console(stderr=True) + + +@app.command("login") +def login_command( + ctx: typer.Context, + client_id: str = typer.Option( + DEFAULT_CLIENT_ID, "--client-id", help="OAuth client ID." + ), + port: int = typer.Option( + DEFAULT_REDIRECT_PORT, "--port", help="Local port for OAuth redirect listener." + ), + uri: str | None = typer.Option( + None, + "--uri", + help="Dremio API base URI (overrides config). E.g. https://api.eu.dremio.cloud", + ), +) -> None: + """Log in to Dremio Cloud via browser-based OAuth (PKCE flow). + + Opens your browser to authenticate, then stores the access token and + refresh token in ~/.config/dremioai/config.yaml. Subsequent CLI commands + will use this token automatically and refresh it when expired. + """ + if not sys.stdin.isatty(): + err_console.print( + "[bold red]dremio auth login[/bold red] requires an interactive terminal." + ) + raise typer.Exit(1) + + # Determine config path and API URI + config_path = _get_config_path(ctx) + api_uri = _resolve_uri(uri, config_path) + + console.print() + console.print( + Panel( + f"[bold]OAuth Login[/bold]\n\n" + f" API: {api_uri}\n" + f" Client ID: {client_id}\n" + f" Port: {port}\n\n" + "A browser window will open for you to sign in to Dremio Cloud.", + title="[bold cyan]dremio auth login[/bold cyan]", + border_style="cyan", + ) + ) + + console.print("\n[dim]Opening browser...[/dim]") + + try: + tokens = run_oauth_flow( + api_uri=api_uri, + client_id=client_id, + redirect_port=port, + ) + except SystemExit as exc: + err_console.print(f"\n[red]✗ {exc}[/red]") + raise typer.Exit(1) + + # Save tokens to config + save_oauth_tokens( + access_token=tokens.access_token, + refresh_token=tokens.refresh_token, + client_id=client_id, + config_path=config_path, + ) + + console.print(f"\n[green]✓ Logged in successfully![/green]") + console.print(f" Tokens saved to [cyan]{config_path}[/cyan]") + if tokens.refresh_token: + console.print(" [dim]Refresh token stored — CLI will auto-refresh on expiry.[/dim]") + else: + console.print(" [yellow]No refresh token received — you'll need to re-login when the token expires.[/yellow]") + + +@app.command("status") +def status_command( + ctx: typer.Context, +) -> None: + """Show current authentication status — token type, expiry, and stored credentials.""" + import yaml + + config_path = _get_config_path(ctx) + + if not config_path.exists(): + console.print("[yellow]No config file found.[/yellow]") + console.print(f" Expected at: {config_path}") + console.print(" Run [bold cyan]dremio auth login[/bold cyan] or [bold cyan]dremio setup[/bold cyan] first.") + raise typer.Exit(1) + + with config_path.open() as f: + raw = yaml.safe_load(f) or {} + + tbl = Table(title="Authentication Status", show_header=True, header_style="bold") + tbl.add_column("Field", style="cyan", no_wrap=True) + tbl.add_column("Value") + + # Config path + tbl.add_row("Config file", str(config_path)) + + # URI + api_uri = raw.get("uri", raw.get("endpoint", DEFAULT_URI)) + tbl.add_row("API URI", api_uri) + + # Project ID + project_id = raw.get("project_id", raw.get("projectId", "")) + tbl.add_row("Project ID", project_id or "[dim]not set[/dim]") + + # Auth method + has_oauth = "oauth" in raw and isinstance(raw.get("oauth"), dict) + has_pat = bool(raw.get("pat") or raw.get("token")) + + if has_oauth: + oauth = raw["oauth"] + access_token = oauth.get("access_token", "") + refresh_token = oauth.get("refresh_token", "") + client_id = oauth.get("client_id", "") + + tbl.add_row("Auth method", "[bold green]OAuth[/bold green]") + tbl.add_row("Client ID", client_id or "[dim]not set[/dim]") + tbl.add_row("Access token", _redact(access_token)) + tbl.add_row("Refresh token", _redact(refresh_token) if refresh_token else "[dim]none[/dim]") + + # Try to decode JWT expiry + expiry_info = _decode_token_expiry(access_token) + if expiry_info: + tbl.add_row("Token expiry", expiry_info) + elif has_pat: + pat = raw.get("pat") or raw.get("token", "") + tbl.add_row("Auth method", "[bold]PAT (Personal Access Token)[/bold]") + tbl.add_row("Token", _redact(pat)) + else: + tbl.add_row("Auth method", "[red]Not configured[/red]") + + console.print() + console.print(tbl) + console.print() + + +@app.command("refresh") +def refresh_command( + ctx: typer.Context, + uri: str | None = typer.Option( + None, "--uri", help="Dremio API base URI (overrides config)." + ), +) -> None: + """Refresh the OAuth access token using the stored refresh token. + + On success, updates the stored access token in the config file. + """ + import yaml + + config_path = _get_config_path(ctx) + + if not config_path.exists(): + err_console.print("[red]No config file found.[/red]") + err_console.print(" Run [bold cyan]dremio auth login[/bold cyan] first.") + raise typer.Exit(1) + + with config_path.open() as f: + raw = yaml.safe_load(f) or {} + + oauth = raw.get("oauth", {}) + if not isinstance(oauth, dict): + oauth = {} + + refresh_token = oauth.get("refresh_token") + client_id = oauth.get("client_id") + + if not refresh_token: + err_console.print("[red]No refresh token found in config.[/red]") + err_console.print(" Run [bold cyan]dremio auth login[/bold cyan] to re-authenticate.") + raise typer.Exit(1) + + if not client_id: + err_console.print("[red]No client_id found in config.[/red]") + err_console.print(" Run [bold cyan]dremio auth login[/bold cyan] to re-authenticate.") + raise typer.Exit(1) + + api_uri = _resolve_uri(uri, config_path) + + console.print("[dim]Refreshing token...[/dim]") + metadata = discover_oauth_metadata(api_uri) + result = do_token_refresh(metadata.token_endpoint, client_id, refresh_token) + + if result is None: + err_console.print("[red]✗ Token refresh failed.[/red]") + err_console.print(" The refresh token may have expired. Run [bold cyan]dremio auth login[/bold cyan] again.") + raise typer.Exit(1) + + # Save the new tokens + save_oauth_tokens( + access_token=result.access_token, + refresh_token=result.refresh_token or refresh_token, + client_id=client_id, + config_path=config_path, + ) + + console.print("[green]✓ Token refreshed successfully![/green]") + console.print(f" Updated tokens in [cyan]{config_path}[/cyan]") + + +@app.command("logout") +def logout_command( + ctx: typer.Context, +) -> None: + """Remove stored OAuth tokens from the config file.""" + config_path = _get_config_path(ctx) + clear_oauth_tokens(config_path) + console.print("[green]✓ OAuth tokens removed.[/green]") + + +# -- Helpers -- + + +def _get_config_path(ctx: typer.Context) -> Path: + """Extract the config path from the typer context.""" + if ctx.obj and ctx.obj.get("config_path"): + return ctx.obj["config_path"] + return DEFAULT_CONFIG_PATH + + +def _resolve_uri(explicit_uri: str | None, config_path: Path) -> str: + """Determine the API URI from explicit flag, config, or default.""" + if explicit_uri: + return explicit_uri + + if config_path.exists(): + import yaml + + with config_path.open() as f: + raw = yaml.safe_load(f) or {} + return raw.get("uri", raw.get("endpoint", DEFAULT_URI)) + + return DEFAULT_URI + + +def _redact(value: str | None, keep: int = 12) -> str: + """Redact a token for display.""" + if not value: + return "[dim]empty[/dim]" + if len(value) <= keep: + return value + return f"{value[:keep]}..." + + +def _decode_token_expiry(token: str | None) -> str | None: + """Try to decode JWT expiry without verifying signature.""" + if not token: + return None + try: + import base64 + import json + from datetime import datetime, timezone + + # Decode JWT payload (second segment) + parts = token.split(".") + if len(parts) != 3: + return None + payload = parts[1] + # Fix padding + padding = 4 - len(payload) % 4 + if padding != 4: + payload += "=" * padding + decoded = json.loads(base64.urlsafe_b64decode(payload)) + exp = decoded.get("exp") + if exp: + exp_dt = datetime.fromtimestamp(exp, tz=timezone.utc) + now = datetime.now(tz=timezone.utc) + if exp_dt < now: + return f"[red]{exp_dt.isoformat()} (EXPIRED)[/red]" + delta = exp_dt - now + hours = int(delta.total_seconds() // 3600) + minutes = int((delta.total_seconds() % 3600) // 60) + return f"{exp_dt.isoformat()} (expires in {hours}h {minutes}m)" + except Exception: + pass + return None diff --git a/src/drs/oauth.py b/src/drs/oauth.py new file mode 100644 index 0000000..b0a0781 --- /dev/null +++ b/src/drs/oauth.py @@ -0,0 +1,272 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""OAuth 2.0 PKCE flow for Dremio Cloud — device/browser-based login.""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import secrets +import threading +import webbrowser +from base64 import urlsafe_b64encode +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any +from urllib.parse import parse_qs, urlencode, urlparse + +import httpx +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + +# Dremio Cloud OAuth well-known endpoints by region +OAUTH_ENDPOINTS = { + "https://api.dremio.cloud": { + "authorization_endpoint": "https://login.dremio.cloud/oauth/authorize", + "token_endpoint": "https://login.dremio.cloud/oauth/token", + }, + "https://api.eu.dremio.cloud": { + "authorization_endpoint": "https://login.eu.dremio.cloud/oauth/authorize", + "token_endpoint": "https://login.eu.dremio.cloud/oauth/token", + }, +} + +# Default OAuth client ID for Dremio CLI (public client) +DEFAULT_CLIENT_ID = "https://connectors.dremio.app/claude" +DEFAULT_REDIRECT_PORT = 8976 +DEFAULT_SCOPES = "dremio.all offline_access" + + +class OAuthTokens(BaseModel): + """Result of an OAuth token exchange.""" + + access_token: str + refresh_token: str | None = None + expires_in: int | None = None + token_type: str = "Bearer" + + +class OAuthMetadata(BaseModel): + """OAuth authorization server metadata.""" + + authorization_endpoint: str + token_endpoint: str + + +def get_pkce_pair() -> tuple[str, str]: + """Generate a PKCE code_verifier + code_challenge pair.""" + code_verifier = secrets.token_urlsafe(96)[:128] + digest = hashlib.sha256(code_verifier.encode("ascii")).digest() + code_challenge = urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + return code_verifier, code_challenge + + +def discover_oauth_metadata(api_uri: str) -> OAuthMetadata: + """Discover OAuth endpoints for a Dremio Cloud API URI. + + Tries the well-known endpoint first, falls back to hardcoded mappings. + """ + # Try well-known discovery + parsed = urlparse(api_uri) + well_known_url = f"{parsed.scheme}://{parsed.hostname}/.well-known/oauth-authorization-server" + try: + resp = httpx.get(well_known_url, timeout=10, follow_redirects=True) + if resp.status_code == 200: + data = resp.json() + return OAuthMetadata( + authorization_endpoint=data["authorization_endpoint"], + token_endpoint=data["token_endpoint"], + ) + except Exception: + logger.debug("Well-known discovery failed for %s, using fallback", api_uri) + + # Fallback to hardcoded endpoints + normalized = api_uri.rstrip("/") + if normalized in OAUTH_ENDPOINTS: + return OAuthMetadata(**OAUTH_ENDPOINTS[normalized]) + + # Default: derive from the API URI hostname + # api.X.dremio.cloud -> login.X.dremio.cloud + hostname = parsed.hostname or "" + if hostname.startswith("api."): + login_host = "login." + hostname[4:] + else: + login_host = hostname + + return OAuthMetadata( + authorization_endpoint=f"{parsed.scheme}://{login_host}/authorize", + token_endpoint=f"{parsed.scheme}://{login_host}/oauth/token", + ) + + +def do_token_refresh( + token_endpoint: str, + client_id: str, + refresh_token: str, +) -> OAuthTokens | None: + """Exchange a refresh_token for a new access_token. + + Returns None if the refresh fails. + """ + try: + resp = httpx.post( + token_endpoint, + data={ + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": client_id, + }, + timeout=15, + ) + if resp.status_code == 200: + body = resp.json() + return OAuthTokens( + access_token=body["access_token"], + refresh_token=body.get("refresh_token", refresh_token), + expires_in=body.get("expires_in"), + token_type=body.get("token_type", "Bearer"), + ) + logger.warning("Token refresh returned HTTP %d: %s", resp.status_code, resp.text[:200]) + except Exception as exc: + logger.warning("Token refresh request failed: %s", exc) + return None + + +class _OAuthCallbackHandler(BaseHTTPRequestHandler): + """HTTP request handler for the OAuth redirect callback.""" + + auth_code: str | None = None + error: str | None = None + server_instance: "_OAuthCallbackServer | None" = None + + def do_GET(self) -> None: # noqa: N802 + parsed = urlparse(self.path) + params = parse_qs(parsed.query) + + if "code" in params: + _OAuthCallbackHandler.auth_code = params["code"][0] + self._respond_html( + "

✓ Authentication successful!

" + "

You can close this tab and return to the terminal.

" + ) + elif "error" in params: + _OAuthCallbackHandler.error = params.get("error_description", params["error"])[0] + self._respond_html(f"

✗ Authentication failed

{_OAuthCallbackHandler.error}

") + else: + self._respond_html("

Unexpected response

") + + # Signal the server to stop + if self.server_instance: + threading.Thread(target=self.server_instance.shutdown, daemon=True).start() + + def _respond_html(self, body: str) -> None: + html = f""" +Dremio CLI Auth + +{body}""" + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write(html.encode()) + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + # Suppress default HTTP server access logs + logger.debug(format, *args) + + +class _OAuthCallbackServer(HTTPServer): + """Minimal HTTP server to receive the OAuth callback.""" + + def __init__(self, port: int) -> None: + super().__init__(("127.0.0.1", port), _OAuthCallbackHandler) + _OAuthCallbackHandler.server_instance = self + _OAuthCallbackHandler.auth_code = None + _OAuthCallbackHandler.error = None + + +def run_oauth_flow( + api_uri: str, + client_id: str = DEFAULT_CLIENT_ID, + redirect_port: int = DEFAULT_REDIRECT_PORT, + scopes: str = DEFAULT_SCOPES, +) -> OAuthTokens: + """Run the full OAuth 2.0 PKCE authorization code flow. + + 1. Start local HTTP server for redirect + 2. Open browser to authorization URL + 3. Wait for callback with auth code + 4. Exchange auth code for tokens + + Raises SystemExit on failure. + """ + metadata = discover_oauth_metadata(api_uri) + code_verifier, code_challenge = get_pkce_pair() + redirect_uri = f"http://localhost:{redirect_port}/Callback" + + # Build authorization URL + auth_params = { + "response_type": "code", + "client_id": client_id, + "redirect_uri": redirect_uri, + "scope": scopes, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + auth_url = f"{metadata.authorization_endpoint}?{urlencode(auth_params)}" + + # Start local server in background thread + server = _OAuthCallbackServer(redirect_port) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + + try: + # Open browser + logger.debug("Opening browser to: %s", auth_url) + webbrowser.open(auth_url) + + # Wait for callback (blocks until the GET handler calls shutdown) + server_thread.join(timeout=300) + + if _OAuthCallbackHandler.error: + raise SystemExit(f"OAuth error: {_OAuthCallbackHandler.error}") + + auth_code = _OAuthCallbackHandler.auth_code + if not auth_code: + raise SystemExit("OAuth flow timed out — no authorization code received.") + + # Exchange code for tokens + token_data = { + "grant_type": "authorization_code", + "client_id": client_id, + "code": auth_code, + "redirect_uri": redirect_uri, + "code_verifier": code_verifier, + } + + resp = httpx.post(metadata.token_endpoint, data=token_data, timeout=15) + if resp.status_code != 200: + raise SystemExit(f"Token exchange failed (HTTP {resp.status_code}): {resp.text[:200]}") + + body = resp.json() + return OAuthTokens( + access_token=body["access_token"], + refresh_token=body.get("refresh_token"), + expires_in=body.get("expires_in"), + token_type=body.get("token_type", "Bearer"), + ) + finally: + server.shutdown() diff --git a/tests/test_client_retry.py b/tests/test_client_retry.py index 485229b..a35e577 100644 --- a/tests/test_client_retry.py +++ b/tests/test_client_retry.py @@ -154,3 +154,110 @@ async def test_retry_backoff_delays(config) -> None: assert mock_sleep.call_count == 2 mock_sleep.assert_any_call(1.0) mock_sleep.assert_any_call(2.0) + + +@pytest.mark.asyncio +async def test_401_triggers_token_refresh(config) -> None: + """401 should trigger OAuth refresh and retry with the new token.""" + from drs.auth import OAuthConfig + + config.oauth = OAuthConfig( + access_token="expired-token", + refresh_token="my-refresh-token", + client_id="my-client-id", + ) + client = DremioClient(config) + + unauthorized = httpx.Response(401, request=httpx.Request("GET", "https://example.com/test")) + ok_response = httpx.Response(200, json={"ok": True}, request=httpx.Request("GET", "https://example.com/test")) + + # First call returns 401, after refresh the second call succeeds + client._client.request = AsyncMock(side_effect=[unauthorized, ok_response]) + + from drs.oauth import OAuthTokens + + mock_tokens = OAuthTokens(access_token="new-access-token", refresh_token="new-refresh-token") + + with ( + patch("drs.oauth.discover_oauth_metadata") as mock_discover, + patch("drs.oauth.do_token_refresh", return_value=mock_tokens) as mock_refresh, + patch("drs.client.save_oauth_tokens") as mock_save, + ): + mock_discover.return_value = type("Meta", (), {"token_endpoint": "https://login.dremio.cloud/oauth/token"})() + result = await client._get("https://example.com/test") + + assert result == {"ok": True} + assert client._client.request.call_count == 2 + mock_refresh.assert_called_once_with("https://login.dremio.cloud/oauth/token", "my-client-id", "my-refresh-token") + mock_save.assert_called_once_with( + access_token="new-access-token", + refresh_token="new-refresh-token", + client_id="my-client-id", + ) + # Verify the client header was updated + assert client._client.headers["Authorization"] == "Bearer new-access-token" + + +@pytest.mark.asyncio +async def test_401_no_refresh_without_oauth_config(config) -> None: + """401 without OAuth config should NOT attempt refresh — just raise.""" + client = DremioClient(config) + + unauthorized = httpx.Response(401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test")) + client._client.request = AsyncMock(return_value=unauthorized) + + with pytest.raises(httpx.HTTPStatusError): + await client._get("https://example.com/test") + + # Only one request, no retry + assert client._client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_401_refresh_only_once(config) -> None: + """Should only attempt refresh once per client instance — no infinite loops.""" + from drs.auth import OAuthConfig + + config.oauth = OAuthConfig( + access_token="expired-token", + refresh_token="my-refresh-token", + client_id="my-client-id", + ) + client = DremioClient(config) + # Simulate: first refresh succeeds but token is still invalid (401 again) + client._refreshed = True # already refreshed once + + unauthorized = httpx.Response(401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test")) + client._client.request = AsyncMock(return_value=unauthorized) + + with pytest.raises(httpx.HTTPStatusError): + await client._get("https://example.com/test") + + assert client._client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_401_refresh_fails_raises_original(config) -> None: + """If token refresh fails, the original 401 should propagate.""" + from drs.auth import OAuthConfig + + config.oauth = OAuthConfig( + access_token="expired-token", + refresh_token="my-refresh-token", + client_id="my-client-id", + ) + client = DremioClient(config) + + unauthorized = httpx.Response(401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test")) + client._client.request = AsyncMock(return_value=unauthorized) + + with ( + patch("drs.oauth.discover_oauth_metadata") as mock_discover, + patch("drs.oauth.do_token_refresh", return_value=None), + pytest.raises(httpx.HTTPStatusError), + ): + mock_discover.return_value = type("Meta", (), {"token_endpoint": "https://login.dremio.cloud/oauth/token"})() + await client._get("https://example.com/test") + + # Only one request attempt — refresh failed so no retry + assert client._client.request.call_count == 1 From 351c69f8390fadda6bfafa822be8ae8ba878a1a6 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 24 Aug 2026 21:21:58 -0400 Subject: [PATCH 2/4] style: fix ruff lint errors in oauth modules --- src/drs/commands/auth.py | 9 +++++---- src/drs/oauth.py | 12 ++++-------- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/src/drs/commands/auth.py b/src/drs/commands/auth.py index 0eb573b..49721f4 100644 --- a/src/drs/commands/auth.py +++ b/src/drs/commands/auth.py @@ -18,6 +18,7 @@ from __future__ import annotations import sys +from datetime import UTC from pathlib import Path import typer @@ -112,7 +113,7 @@ def login_command( config_path=config_path, ) - console.print(f"\n[green]✓ Logged in successfully![/green]") + console.print("\n[green]✓ Logged in successfully![/green]") console.print(f" Tokens saved to [cyan]{config_path}[/cyan]") if tokens.refresh_token: console.print(" [dim]Refresh token stored — CLI will auto-refresh on expiry.[/dim]") @@ -298,7 +299,7 @@ def _decode_token_expiry(token: str | None) -> str | None: try: import base64 import json - from datetime import datetime, timezone + from datetime import datetime # Decode JWT payload (second segment) parts = token.split(".") @@ -312,8 +313,8 @@ def _decode_token_expiry(token: str | None) -> str | None: decoded = json.loads(base64.urlsafe_b64decode(payload)) exp = decoded.get("exp") if exp: - exp_dt = datetime.fromtimestamp(exp, tz=timezone.utc) - now = datetime.now(tz=timezone.utc) + exp_dt = datetime.fromtimestamp(exp, tz=UTC) + now = datetime.now(tz=UTC) if exp_dt < now: return f"[red]{exp_dt.isoformat()} (EXPIRED)[/red]" delta = exp_dt - now diff --git a/src/drs/oauth.py b/src/drs/oauth.py index b0a0781..7ffa674 100644 --- a/src/drs/oauth.py +++ b/src/drs/oauth.py @@ -17,7 +17,6 @@ from __future__ import annotations -import asyncio import hashlib import logging import secrets @@ -102,10 +101,7 @@ def discover_oauth_metadata(api_uri: str) -> OAuthMetadata: # Default: derive from the API URI hostname # api.X.dremio.cloud -> login.X.dremio.cloud hostname = parsed.hostname or "" - if hostname.startswith("api."): - login_host = "login." + hostname[4:] - else: - login_host = hostname + login_host = "login." + hostname[4:] if hostname.startswith("api.") else hostname return OAuthMetadata( authorization_endpoint=f"{parsed.scheme}://{login_host}/authorize", @@ -151,9 +147,9 @@ class _OAuthCallbackHandler(BaseHTTPRequestHandler): auth_code: str | None = None error: str | None = None - server_instance: "_OAuthCallbackServer | None" = None + server_instance: _OAuthCallbackServer | None = None - def do_GET(self) -> None: # noqa: N802 + def do_GET(self) -> None: parsed = urlparse(self.path) params = parse_qs(parsed.query) @@ -183,7 +179,7 @@ def _respond_html(self, body: str) -> None: self.end_headers() self.wfile.write(html.encode()) - def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + def log_message(self, format: str, *args: Any) -> None: # Suppress default HTTP server access logs logger.debug(format, *args) From 1de108ef39ef32f06c9602afd634c134fdcf3510 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 24 Aug 2026 21:23:58 -0400 Subject: [PATCH 3/4] style: apply ruff format --- src/drs/client.py | 4 +--- src/drs/commands/auth.py | 16 ++++------------ src/drs/oauth.py | 3 +-- tests/test_client_retry.py | 12 +++++++++--- 4 files changed, 15 insertions(+), 20 deletions(-) diff --git a/src/drs/client.py b/src/drs/client.py index 6bd4b20..0b5da32 100644 --- a/src/drs/client.py +++ b/src/drs/client.py @@ -79,9 +79,7 @@ def _try_refresh_token(self) -> bool: try: metadata = discover_oauth_metadata(self.config.uri) - result = do_token_refresh( - metadata.token_endpoint, oauth.client_id, oauth.refresh_token - ) + result = do_token_refresh(metadata.token_endpoint, oauth.client_id, oauth.refresh_token) except Exception as exc: logger.warning("OAuth token refresh failed: %s", exc) return False diff --git a/src/drs/commands/auth.py b/src/drs/commands/auth.py index 49721f4..2937a7e 100644 --- a/src/drs/commands/auth.py +++ b/src/drs/commands/auth.py @@ -52,12 +52,8 @@ @app.command("login") def login_command( ctx: typer.Context, - client_id: str = typer.Option( - DEFAULT_CLIENT_ID, "--client-id", help="OAuth client ID." - ), - port: int = typer.Option( - DEFAULT_REDIRECT_PORT, "--port", help="Local port for OAuth redirect listener." - ), + client_id: str = typer.Option(DEFAULT_CLIENT_ID, "--client-id", help="OAuth client ID."), + port: int = typer.Option(DEFAULT_REDIRECT_PORT, "--port", help="Local port for OAuth redirect listener."), uri: str | None = typer.Option( None, "--uri", @@ -71,9 +67,7 @@ def login_command( will use this token automatically and refresh it when expired. """ if not sys.stdin.isatty(): - err_console.print( - "[bold red]dremio auth login[/bold red] requires an interactive terminal." - ) + err_console.print("[bold red]dremio auth login[/bold red] requires an interactive terminal.") raise typer.Exit(1) # Determine config path and API URI @@ -188,9 +182,7 @@ def status_command( @app.command("refresh") def refresh_command( ctx: typer.Context, - uri: str | None = typer.Option( - None, "--uri", help="Dremio API base URI (overrides config)." - ), + uri: str | None = typer.Option(None, "--uri", help="Dremio API base URI (overrides config)."), ) -> None: """Refresh the OAuth access token using the stored refresh token. diff --git a/src/drs/oauth.py b/src/drs/oauth.py index 7ffa674..095eadf 100644 --- a/src/drs/oauth.py +++ b/src/drs/oauth.py @@ -156,8 +156,7 @@ def do_GET(self) -> None: if "code" in params: _OAuthCallbackHandler.auth_code = params["code"][0] self._respond_html( - "

✓ Authentication successful!

" - "

You can close this tab and return to the terminal.

" + "

✓ Authentication successful!

You can close this tab and return to the terminal.

" ) elif "error" in params: _OAuthCallbackHandler.error = params.get("error_description", params["error"])[0] diff --git a/tests/test_client_retry.py b/tests/test_client_retry.py index a35e577..841f0e6 100644 --- a/tests/test_client_retry.py +++ b/tests/test_client_retry.py @@ -203,7 +203,9 @@ async def test_401_no_refresh_without_oauth_config(config) -> None: """401 without OAuth config should NOT attempt refresh — just raise.""" client = DremioClient(config) - unauthorized = httpx.Response(401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test")) + unauthorized = httpx.Response( + 401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test") + ) client._client.request = AsyncMock(return_value=unauthorized) with pytest.raises(httpx.HTTPStatusError): @@ -227,7 +229,9 @@ async def test_401_refresh_only_once(config) -> None: # Simulate: first refresh succeeds but token is still invalid (401 again) client._refreshed = True # already refreshed once - unauthorized = httpx.Response(401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test")) + unauthorized = httpx.Response( + 401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test") + ) client._client.request = AsyncMock(return_value=unauthorized) with pytest.raises(httpx.HTTPStatusError): @@ -248,7 +252,9 @@ async def test_401_refresh_fails_raises_original(config) -> None: ) client = DremioClient(config) - unauthorized = httpx.Response(401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test")) + unauthorized = httpx.Response( + 401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test") + ) client._client.request = AsyncMock(return_value=unauthorized) with ( From c5dd22feb5e1c21445a9ac2a4e89d2d0ac391235 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 24 Aug 2026 21:29:00 -0400 Subject: [PATCH 4/4] =?UTF-8?q?fix:=20address=20PR=20review=20=E2=80=94=20?= =?UTF-8?q?OAuth=20priority,=20auth=5Fsource=20gating,=20config=5Fpath,=20?= =?UTF-8?q?transient=20retry?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. OAuth access_token now overrides file PAT (not just fills in when missing), so 'auth login' is effective even with an existing PAT from 'setup'. (review comment #1) 2. Track auth_source on DrsConfig ('oauth'|'file'|'env'|'cli'). Auto- refresh only activates when auth_source == 'oauth', preventing the client from silently switching identity when --token or DREMIO_TOKEN supplies the active credential. (review comment #2) 3. Store the effective config_path on DrsConfig and pass it through to save_oauth_tokens during auto-refresh, so custom --config paths get their tokens updated correctly. (review comment #3) 4. After a successful 401 refresh, the retried response now falls through to the transient-status retry loop instead of returning immediately, so a post-refresh 429/503 is still retried. (review comment #4) Tests: 6 new tests (229 total) covering all four fixes. --- src/drs/auth.py | 18 +++++--- src/drs/client.py | 13 +++++- tests/test_auth.py | 91 ++++++++++++++++++++++++++++++++++++++ tests/test_client_retry.py | 68 +++++++++++++++++++++++++++- 4 files changed, 182 insertions(+), 8 deletions(-) diff --git a/src/drs/auth.py b/src/drs/auth.py index e85bada..9da0e16 100644 --- a/src/drs/auth.py +++ b/src/drs/auth.py @@ -44,6 +44,8 @@ class DrsConfig(BaseModel): pat: str project_id: str oauth: OAuthConfig | None = None + auth_source: str = "file" # "oauth" | "file" | "env" | "cli" + config_path: Path = DEFAULT_CONFIG_PATH def load_config( @@ -82,11 +84,14 @@ def load_config( refresh_token=raw["oauth"].get("refresh_token"), client_id=raw["oauth"].get("client_id"), ) - # Use OAuth access_token as the PAT if no explicit PAT is set - if oauth_config.access_token and "pat" not in file_values: - file_values["pat"] = oauth_config.access_token - # -- Env vars (override file) -- + # -- Determine auth source: OAuth access_token overrides file PAT -- + auth_source = "file" + if oauth_config and oauth_config.access_token: + file_values["pat"] = oauth_config.access_token + auth_source = "oauth" + + # -- Env vars (override file + oauth) -- env_values: dict[str, Any] = {} if v := os.environ.get("DREMIO_URI"): env_values["uri"] = v @@ -98,6 +103,8 @@ def load_config( # -- Merge: defaults < file < env -- merged: dict[str, Any] = {"uri": DEFAULT_URI} merged.update(file_values) + if "pat" in env_values: + auth_source = "env" merged.update(env_values) # -- CLI args (highest priority, override everything) -- @@ -107,8 +114,9 @@ def load_config( merged["project_id"] = cli_project_id if cli_token: merged["pat"] = cli_token + auth_source = "cli" - config = DrsConfig(**merged) + config = DrsConfig(**merged, auth_source=auth_source, config_path=path) config.oauth = oauth_config return config diff --git a/src/drs/client.py b/src/drs/client.py index 0b5da32..5cfc19a 100644 --- a/src/drs/client.py +++ b/src/drs/client.py @@ -67,10 +67,16 @@ def _try_refresh_token(self) -> bool: """Attempt to refresh the OAuth access token synchronously. Returns True if the token was refreshed and the client headers updated. + Only activates when the active credential is the OAuth access token + (auth_source == "oauth"), not when env vars or CLI --token are in use. """ if self._refreshed: return False # already tried once this session + # Only refresh when OAuth is the active credential source + if self.config.auth_source != "oauth": + return False + oauth = self.config.oauth if not oauth or not oauth.refresh_token or not oauth.client_id: return False @@ -94,11 +100,12 @@ def _try_refresh_token(self) -> bool: if result.refresh_token: oauth.refresh_token = result.refresh_token - # Persist to config file + # Persist to the config file that was actually loaded save_oauth_tokens( access_token=result.access_token, refresh_token=result.refresh_token or oauth.refresh_token, client_id=oauth.client_id, + config_path=self.config.config_path, ) # Update httpx client headers @@ -136,7 +143,9 @@ async def _request_with_retry(self, method: str, url: str, **kwargs: Any) -> htt if resp.status_code == 401 and self._try_refresh_token(): logger.info("Retrying %s %s after token refresh", method, url) resp = await self._client.request(method, url, **kwargs) - return resp + # Fall through to retryable-status handling below so a + # transient error (429/503) on the refreshed request is + # still retried instead of returned immediately. if resp.status_code in _RETRYABLE_STATUS_CODES and attempt < _MAX_RETRIES - 1: delay = _RETRY_BACKOFF[attempt] diff --git a/tests/test_auth.py b/tests/test_auth.py index 9187f81..e066ae3 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -162,3 +162,94 @@ def test_cli_args_override_env(tmp_path: Path) -> None: assert config.pat == "cli-token" assert config.project_id == "cli-project" assert config.uri == "https://api.eu.dremio.cloud" + + +def test_oauth_overrides_file_pat(tmp_path: Path) -> None: + """OAuth access_token should override the file PAT field.""" + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump( + { + "pat": "old-file-pat", + "project_id": "proj", + "oauth": { + "access_token": "oauth-access-token", + "refresh_token": "oauth-refresh-token", + "client_id": "my-client-id", + }, + } + ) + ) + + with patch.dict(os.environ, {}, clear=False): + for k in ["DREMIO_TOKEN", "DREMIO_PAT", "DREMIO_PROJECT_ID", "DREMIO_URI"]: + os.environ.pop(k, None) + config = load_config(config_file) + + assert config.pat == "oauth-access-token" + assert config.auth_source == "oauth" + + +def test_env_token_overrides_oauth(tmp_path: Path) -> None: + """DREMIO_TOKEN env var should override OAuth access_token.""" + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump( + { + "project_id": "proj", + "oauth": { + "access_token": "oauth-access-token", + "refresh_token": "oauth-refresh-token", + "client_id": "my-client-id", + }, + } + ) + ) + + with patch.dict(os.environ, {"DREMIO_TOKEN": "env-token"}, clear=False): + os.environ.pop("DREMIO_PAT", None) + os.environ.pop("DREMIO_PROJECT_ID", None) + os.environ.pop("DREMIO_URI", None) + config = load_config(config_file) + + assert config.pat == "env-token" + assert config.auth_source == "env" + + +def test_cli_token_auth_source(tmp_path: Path) -> None: + """--token CLI arg should set auth_source to 'cli'.""" + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump( + { + "project_id": "proj", + "oauth": { + "access_token": "oauth-access-token", + "refresh_token": "oauth-refresh-token", + "client_id": "my-client-id", + }, + } + ) + ) + + with patch.dict(os.environ, {}, clear=False): + for k in ["DREMIO_TOKEN", "DREMIO_PAT", "DREMIO_PROJECT_ID", "DREMIO_URI"]: + os.environ.pop(k, None) + config = load_config(config_file, cli_token="cli-token") + + assert config.pat == "cli-token" + assert config.auth_source == "cli" + + +def test_config_path_preserved(tmp_path: Path) -> None: + """The effective config path should be stored on DrsConfig.""" + config_file = tmp_path / "custom" / "config.yaml" + config_file.parent.mkdir(parents=True) + config_file.write_text(yaml.dump({"pat": "tok", "project_id": "proj"})) + + with patch.dict(os.environ, {}, clear=False): + for k in ["DREMIO_TOKEN", "DREMIO_PAT", "DREMIO_PROJECT_ID", "DREMIO_URI"]: + os.environ.pop(k, None) + config = load_config(config_file) + + assert config.config_path == config_file diff --git a/tests/test_client_retry.py b/tests/test_client_retry.py index 841f0e6..b4f48c4 100644 --- a/tests/test_client_retry.py +++ b/tests/test_client_retry.py @@ -166,6 +166,7 @@ async def test_401_triggers_token_refresh(config) -> None: refresh_token="my-refresh-token", client_id="my-client-id", ) + config.auth_source = "oauth" client = DremioClient(config) unauthorized = httpx.Response(401, request=httpx.Request("GET", "https://example.com/test")) @@ -193,6 +194,7 @@ async def test_401_triggers_token_refresh(config) -> None: access_token="new-access-token", refresh_token="new-refresh-token", client_id="my-client-id", + config_path=config.config_path, ) # Verify the client header was updated assert client._client.headers["Authorization"] == "Bearer new-access-token" @@ -200,7 +202,7 @@ async def test_401_triggers_token_refresh(config) -> None: @pytest.mark.asyncio async def test_401_no_refresh_without_oauth_config(config) -> None: - """401 without OAuth config should NOT attempt refresh — just raise.""" + """401 without OAuth auth_source should NOT attempt refresh — just raise.""" client = DremioClient(config) unauthorized = httpx.Response( @@ -225,6 +227,7 @@ async def test_401_refresh_only_once(config) -> None: refresh_token="my-refresh-token", client_id="my-client-id", ) + config.auth_source = "oauth" client = DremioClient(config) # Simulate: first refresh succeeds but token is still invalid (401 again) client._refreshed = True # already refreshed once @@ -250,6 +253,7 @@ async def test_401_refresh_fails_raises_original(config) -> None: refresh_token="my-refresh-token", client_id="my-client-id", ) + config.auth_source = "oauth" client = DremioClient(config) unauthorized = httpx.Response( @@ -267,3 +271,65 @@ async def test_401_refresh_fails_raises_original(config) -> None: # Only one request attempt — refresh failed so no retry assert client._client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_401_no_refresh_when_env_token_overrides_oauth(config) -> None: + """When DREMIO_TOKEN overrides OAuth, 401 should NOT try OAuth refresh.""" + from drs.auth import OAuthConfig + + config.oauth = OAuthConfig( + access_token="oauth-token", + refresh_token="my-refresh-token", + client_id="my-client-id", + ) + config.auth_source = "env" # env var won the priority chain + client = DremioClient(config) + + unauthorized = httpx.Response( + 401, json={"error": "unauthorized"}, request=httpx.Request("GET", "https://example.com/test") + ) + client._client.request = AsyncMock(return_value=unauthorized) + + with pytest.raises(httpx.HTTPStatusError): + await client._get("https://example.com/test") + + assert client._client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_401_refresh_then_transient_retry(config) -> None: + """After refresh, a 503 on the retry should still enter the transient-retry loop.""" + from drs.auth import OAuthConfig + + config.oauth = OAuthConfig( + access_token="expired-token", + refresh_token="my-refresh-token", + client_id="my-client-id", + ) + config.auth_source = "oauth" + client = DremioClient(config) + + unauthorized = httpx.Response(401, request=httpx.Request("GET", "https://example.com/test")) + transient = httpx.Response(503, request=httpx.Request("GET", "https://example.com/test")) + ok_response = httpx.Response(200, json={"ok": True}, request=httpx.Request("GET", "https://example.com/test")) + + # 401 → refresh → 503 (transient) → retry → 200 + client._client.request = AsyncMock(side_effect=[unauthorized, transient, ok_response]) + + from drs.oauth import OAuthTokens + + mock_tokens = OAuthTokens(access_token="new-access-token", refresh_token="new-refresh-token") + + with ( + patch("drs.oauth.discover_oauth_metadata") as mock_discover, + patch("drs.oauth.do_token_refresh", return_value=mock_tokens), + patch("drs.client.save_oauth_tokens"), + patch("drs.client.asyncio.sleep", new_callable=AsyncMock), + ): + mock_discover.return_value = type("Meta", (), {"token_endpoint": "https://login.dremio.cloud/oauth/token"})() + result = await client._get("https://example.com/test") + + assert result == {"ok": True} + # 3 requests: original 401, post-refresh 503, transient-retry 200 + assert client._client.request.call_count == 3