diff --git a/src/drs/auth.py b/src/drs/auth.py index 051f365..9da0e16 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,27 @@ 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 + auth_source: str = "file" # "oauth" | "file" | "env" | "cli" + config_path: Path = DEFAULT_CONFIG_PATH def load_config( @@ -46,10 +60,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,7 +77,21 @@ def load_config( } file_values = {k: v for k, v in file_values.items() if v is not None} - # -- Env vars (override file) -- + # 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"), + ) + + # -- 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 @@ -73,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) -- @@ -82,5 +114,59 @@ def load_config( merged["project_id"] = cli_project_id if cli_token: merged["pat"] = cli_token + auth_source = "cli" + + config = DrsConfig(**merged, auth_source=auth_source, config_path=path) + 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) + - return DrsConfig(**merged) +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..5cfc19a 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,64 @@ 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. + 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 + + 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 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 + 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 +130,23 @@ 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) + # 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] logger.warning( diff --git a/src/drs/commands/auth.py b/src/drs/commands/auth.py new file mode 100644 index 0000000..2937a7e --- /dev/null +++ b/src/drs/commands/auth.py @@ -0,0 +1,318 @@ +# +# 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 datetime import UTC +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("\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 + + # 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=UTC) + now = datetime.now(tz=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..095eadf --- /dev/null +++ b/src/drs/oauth.py @@ -0,0 +1,267 @@ +# +# 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 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 "" + login_host = "login." + hostname[4:] if hostname.startswith("api.") else 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: + 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: + # 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_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 485229b..b4f48c4 100644 --- a/tests/test_client_retry.py +++ b/tests/test_client_retry.py @@ -154,3 +154,182 @@ 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", + ) + config.auth_source = "oauth" + 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", + config_path=config.config_path, + ) + # 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 auth_source 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", + ) + config.auth_source = "oauth" + 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", + ) + config.auth_source = "oauth" + 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 + + +@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