feat: add OAuth auth command with login, status, refresh and auto-401 retry

src/drs/auth.py

@@ -17,21 +17,35 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from pathlib import Path
 from typing import Any
 
 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

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")

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(