sqlsaber 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

sqlsaber/config/database.py

@@ -6,7 +6,7 @@ import platform
 import stat
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any
 from urllib.parse import quote_plus
 
 import keyring
@@ -19,16 +19,16 @@ class DatabaseConfig:
 
     name: str
     type: str  # postgresql, mysql, sqlite, csv
-    host: Optional[str]
-    port: Optional[int]
+    host: str | None
+    port: int | None
     database: str
-    username: Optional[str]
-    password: Optional[str] = None
-    ssl_mode: Optional[str] = None
-    ssl_ca: Optional[str] = None
-    ssl_cert: Optional[str] = None
-    ssl_key: Optional[str] = None
-    schema: Optional[str] = None
+    username: str | None
+    password: str | None = None
+    ssl_mode: str | None = None
+    ssl_ca: str | None = None
+    ssl_cert: str | None = None
+    ssl_key: str | None = None
+    schema: str | None = None
 
     def to_connection_string(self) -> str:
         """Convert config to database connection string."""
@@ -115,7 +115,7 @@ class DatabaseConfig:
         else:
             raise ValueError(f"Unsupported database type: {self.type}")
 
-    def _get_password_from_keyring(self) -> Optional[str]:
+    def _get_password_from_keyring(self) -> str | None:
         """Get password from OS keyring."""
         try:
             return keyring.get_password("sqlsaber", f"{self.name}_{self.username}")
@@ -133,7 +133,7 @@ class DatabaseConfig:
         except Exception:
             pass
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Convert to dictionary for JSON serialization."""
         return {
             "name": self.name,
@@ -150,7 +150,7 @@ class DatabaseConfig:
         }
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "DatabaseConfig":
+    def from_dict(cls, data: dict[str, Any]) -> "DatabaseConfig":
         """Create from dictionary."""
         return cls(
             name=data["name"],
@@ -202,7 +202,7 @@ class DatabaseConfigManager:
             # The directory/file creation should still work
             pass
 
-    def _load_config(self) -> Dict[str, Any]:
+    def _load_config(self) -> dict[str, Any]:
         """Load configuration from file."""
         if not self.config_file.exists():
             return {"default": None, "connections": {}}
@@ -213,7 +213,7 @@ class DatabaseConfigManager:
         except (json.JSONDecodeError, IOError):
             return {"default": None, "connections": {}}
 
-    def _save_config(self, config: Dict[str, Any]) -> None:
+    def _save_config(self, config: dict[str, Any]) -> None:
         """Save configuration to file."""
         with open(self.config_file, "w") as f:
             json.dump(config, f, indent=2)
@@ -222,7 +222,7 @@ class DatabaseConfigManager:
         self._set_secure_permissions(self.config_file, is_directory=False)
 
     def add_database(
-        self, db_config: DatabaseConfig, password: Optional[str] = None
+        self, db_config: DatabaseConfig, password: str | None = None
     ) -> None:
         """Add a database configuration."""
         config = self._load_config()
@@ -244,7 +244,7 @@ class DatabaseConfigManager:
 
         self._save_config(config)
 
-    def get_database(self, name: str) -> Optional[DatabaseConfig]:
+    def get_database(self, name: str) -> DatabaseConfig | None:
         """Get a database configuration by name."""
         config = self._load_config()
 
@@ -253,7 +253,7 @@ class DatabaseConfigManager:
 
         return DatabaseConfig.from_dict(config["connections"][name])
 
-    def get_default_database(self) -> Optional[DatabaseConfig]:
+    def get_default_database(self) -> DatabaseConfig | None:
         """Get the default database configuration."""
         config = self._load_config()
 
@@ -263,7 +263,7 @@ class DatabaseConfigManager:
 
         return self.get_database(default_name)
 
-    def list_databases(self) -> List[DatabaseConfig]:
+    def list_databases(self) -> list[DatabaseConfig]:
         """List all database configurations."""
         config = self._load_config()
 
@@ -313,7 +313,7 @@ class DatabaseConfigManager:
         config = self._load_config()
         return len(config["connections"]) > 0
 
-    def get_default_name(self) -> Optional[str]:
+    def get_default_name(self) -> str | None:
         """Get the name of the default database."""
         config = self._load_config()
         return config.get("default")

sqlsaber/config/oauth_flow.py

@@ -0,0 +1,274 @@
+"""Synchronous OAuth flow management for Anthropic Claude Pro authentication."""
+
+import base64
+import hashlib
+import logging
+import secrets
+import urllib.parse
+import webbrowser
+from datetime import datetime, timezone
+
+import httpx
+import questionary
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+from .oauth_tokens import OAuthToken, OAuthTokenManager
+
+console = Console()
+logger = logging.getLogger(__name__)
+
+
+CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
+
+
+class AnthropicOAuthFlow:
+    """Handles the complete OAuth flow for Anthropic Claude Pro authentication."""
+
+    def __init__(self):
+        self.client_id = CLIENT_ID
+        self.token_manager = OAuthTokenManager()
+
+    def _generate_pkce(self) -> tuple[str, str]:
+        """Generate PKCE code verifier and challenge."""
+        verifier = (
+            base64.urlsafe_b64encode(secrets.token_bytes(32))
+            .decode("utf-8")
+            .rstrip("=")
+        )
+        challenge = (
+            base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("utf-8")).digest())
+            .decode("utf-8")
+            .rstrip("=")
+        )
+        return verifier, challenge
+
+    def _create_authorization_url(self) -> tuple[str, str]:
+        """Create OAuth authorization URL with PKCE."""
+        verifier, challenge = self._generate_pkce()
+
+        params = {
+            "code": "true",
+            "client_id": self.client_id,
+            "response_type": "code",
+            "redirect_uri": "https://console.anthropic.com/oauth/code/callback",
+            "scope": "org:create_api_key user:profile user:inference",
+            "code_challenge": challenge,
+            "code_challenge_method": "S256",
+            "state": verifier,
+        }
+
+        url = "https://claude.ai/oauth/authorize?" + urllib.parse.urlencode(params)
+        return url, verifier
+
+    def _exchange_code_for_tokens(self, code: str, verifier: str) -> dict[str, str]:
+        """Exchange authorization code for access and refresh tokens."""
+        # Handle the code format (may have # separator for state)
+        code_parts = code.split("#")
+        auth_code = code_parts[0]
+        state = code_parts[1] if len(code_parts) > 1 else verifier
+
+        data = {
+            "code": auth_code,
+            "state": state,
+            "grant_type": "authorization_code",
+            "client_id": self.client_id,
+            "redirect_uri": "https://console.anthropic.com/oauth/code/callback",
+            "code_verifier": verifier,
+        }
+
+        with httpx.Client() as client:
+            response = client.post(
+                "https://console.anthropic.com/v1/oauth/token",
+                headers={"Content-Type": "application/json"},
+                json=data,
+            )
+
+            if not response.is_success:
+                error_msg = (
+                    f"Token exchange failed: {response.status_code} {response.text}"
+                )
+                logger.error(error_msg)
+                raise Exception(error_msg)
+
+            return response.json()
+
+    def refresh_access_token(self, refresh_token: str) -> dict[str, str]:
+        """Refresh access token using refresh token."""
+        data = {
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+            "client_id": self.client_id,
+        }
+
+        with httpx.Client() as client:
+            response = client.post(
+                "https://console.anthropic.com/v1/oauth/token",
+                headers={"Content-Type": "application/json"},
+                json=data,
+            )
+
+            if not response.is_success:
+                error_msg = (
+                    f"Token refresh failed: {response.status_code} {response.text}"
+                )
+                logger.error(error_msg)
+                raise Exception(error_msg)
+
+            return response.json()
+
+    def authenticate(self) -> bool:
+        """Complete OAuth authentication flow."""
+        console.print(
+            "\n[bold blue]Claude Pro/Max Subscription Authentication[/bold blue]"
+        )
+        console.print(
+            "This will open your web browser to authenticate with your Claude subscription.\n"
+        )
+
+        # Check if user wants to proceed
+        if not questionary.confirm(
+            "Continue with browser-based authentication?", default=True
+        ).ask():
+            console.print("[yellow]Authentication cancelled.[/yellow]")
+            return False
+
+        try:
+            # Step 1: Create authorization URL
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                console=console,
+            ) as progress:
+                task = progress.add_task("Preparing authentication...", total=None)
+
+                auth_url, verifier = self._create_authorization_url()
+                progress.update(task, description="Opening browser...")
+
+                # Open browser for user authorization
+                webbrowser.open(auth_url)
+
+            console.print("\n[green]✓[/green] Browser opened for authentication")
+            console.print(
+                "[dim]If your browser didn't open automatically, visit this URL:[/dim]"
+            )
+            console.print(f"[dim]{auth_url}[/dim]\n")
+
+            # Get authorization code from user
+            console.print("After authorizing, you'll be redirected to a callback URL.")
+            console.print(
+                "Copy the 'code' that shows up on your screen and paste it here."
+            )
+
+            auth_code = questionary.text(
+                "Enter the authorization code:",
+                validate=lambda x: len(x.strip()) > 0
+                or "Authorization code is required",
+            ).ask()
+
+            if not auth_code:
+                console.print("[yellow]Authentication cancelled.[/yellow]")
+                return False
+
+            # Step 2: Exchange code for tokens
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                console=console,
+            ) as progress:
+                task = progress.add_task("Exchanging code for tokens...", total=None)
+
+                tokens = self._exchange_code_for_tokens(auth_code.strip(), verifier)
+
+                # Calculate expiration time if provided
+                expires_at = None
+                if "expires_in" in tokens:
+                    expires_in = int(tokens["expires_in"])
+                    expires_dt = datetime.now(timezone.utc).timestamp() + expires_in
+                    expires_at = datetime.fromtimestamp(
+                        expires_dt, timezone.utc
+                    ).isoformat()
+
+                # Store tokens
+                oauth_token = OAuthToken(
+                    access_token=tokens["access_token"],
+                    refresh_token=tokens["refresh_token"],
+                    expires_at=expires_at,
+                )
+
+                if self.token_manager.store_oauth_token("anthropic", oauth_token):
+                    console.print(
+                        "\n[bold green]✓ Authentication successful![/bold green]"
+                    )
+                    console.print(
+                        "Your Claude Pro/Max subscription is now configured for SQLSaber."
+                    )
+                    return True
+                else:
+                    console.print("[red]✗ Failed to store authentication tokens.[/red]")
+                    return False
+
+        except KeyboardInterrupt:
+            console.print("\n[yellow]Authentication cancelled by user.[/yellow]")
+            return False
+        except Exception as e:
+            logger.error(f"OAuth authentication failed: {e}")
+            console.print(f"[red]✗ Authentication failed: {str(e)}[/red]")
+            return False
+
+    def refresh_token_if_needed(self) -> OAuthToken | None:
+        """Refresh OAuth token if it's expired or expiring soon."""
+        current_token = self.token_manager.get_oauth_token("anthropic")
+        if not current_token:
+            return None
+
+        # If token is not expired and not expiring soon, return it as-is
+        if not current_token.is_expired() and not current_token.expires_soon():
+            return current_token
+
+        # Attempt to refresh
+        try:
+            console.print("Refreshing OAuth token...", style="dim")
+            new_tokens = self.refresh_access_token(current_token.refresh_token)
+
+            # Calculate new expiration time
+            expires_at = None
+            if "expires_in" in new_tokens:
+                expires_in = int(new_tokens["expires_in"])
+                expires_dt = datetime.now(timezone.utc).timestamp() + expires_in
+                expires_at = datetime.fromtimestamp(
+                    expires_dt, timezone.utc
+                ).isoformat()
+
+            # Create new token object
+            refreshed_token = OAuthToken(
+                access_token=new_tokens["access_token"],
+                refresh_token=new_tokens.get(
+                    "refresh_token", current_token.refresh_token
+                ),
+                expires_at=expires_at,
+            )
+
+            # Store the refreshed token
+            if self.token_manager.store_oauth_token("anthropic", refreshed_token):
+                console.print("OAuth token refreshed successfully", style="green")
+                return refreshed_token
+            else:
+                console.print("Failed to store refreshed token", style="yellow")
+                return current_token
+
+        except Exception as e:
+            logger.warning(f"Token refresh failed: {e}")
+            console.print(
+                "Token refresh failed. You may need to re-authenticate.", style="yellow"
+            )
+            return current_token
+
+    def remove_authentication(self) -> bool:
+        """Remove stored OAuth authentication."""
+        return self.token_manager.remove_oauth_token("anthropic")
+
+    def has_valid_authentication(self) -> bool:
+        """Check if valid OAuth authentication exists."""
+        token = self.token_manager.get_oauth_token("anthropic")
+        return token is not None and not token.is_expired()

sqlsaber/config/oauth_tokens.py

@@ -0,0 +1,175 @@
+"""OAuth token management for SQLSaber."""
+
+import json
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+import keyring
+from rich.console import Console
+
+console = Console()
+logger = logging.getLogger(__name__)
+
+
+class OAuthToken:
+    """Represents an OAuth token with metadata."""
+
+    def __init__(
+        self,
+        access_token: str,
+        refresh_token: str,
+        expires_at: str | None = None,
+        token_type: str = "Bearer",
+    ):
+        self.access_token = access_token
+        self.refresh_token = refresh_token
+        self.expires_at = expires_at
+        self.token_type = token_type
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "OAuthToken":
+        """Create token from dictionary."""
+        return cls(
+            access_token=data["access_token"],
+            refresh_token=data["refresh_token"],
+            expires_at=data.get("expires_at"),
+            token_type=data.get("token_type", "Bearer"),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert token to dictionary."""
+        return {
+            "access_token": self.access_token,
+            "refresh_token": self.refresh_token,
+            "expires_at": self.expires_at,
+            "token_type": self.token_type,
+        }
+
+    def is_expired(self) -> bool:
+        """Check if the token is expired."""
+        if not self.expires_at:
+            return False
+
+        try:
+            expires_dt = datetime.fromisoformat(self.expires_at.replace("Z", "+00:00"))
+            return datetime.now(timezone.utc) >= expires_dt
+        except (ValueError, AttributeError):
+            # If we can't parse the expiration, assume expired for safety
+            return True
+
+    def expires_soon(self, buffer_seconds: int = 300) -> bool:
+        """Check if token expires within buffer_seconds."""
+        if not self.expires_at:
+            return False
+
+        try:
+            expires_dt = datetime.fromisoformat(self.expires_at.replace("Z", "+00:00"))
+
+            return (
+                datetime.now(timezone.utc) + timedelta(seconds=buffer_seconds)
+            ) >= expires_dt
+        except (ValueError, AttributeError):
+            return True
+
+
+class OAuthTokenManager:
+    """Manages OAuth tokens with secure storage and refresh logic."""
+
+    def __init__(self):
+        self.service_prefix = "sqlsaber"
+
+    def get_oauth_token(self, provider: str) -> OAuthToken | None:
+        """Get OAuth token for the specified provider."""
+        service_name = self._get_service_name(provider)
+
+        try:
+            token_data = keyring.get_password(service_name, provider)
+            if not token_data:
+                return None
+
+            # Parse the stored JSON
+            data = json.loads(token_data)
+            token = OAuthToken.from_dict(data)
+
+            # Check if token is expired
+            if token.is_expired():
+                console.print(
+                    f"OAuth token for {provider} has expired and needs refresh",
+                    style="dim yellow",
+                )
+                return token  # Return anyway for refresh attempt
+
+            if token.expires_soon():
+                console.print(
+                    f"OAuth token for {provider} expires soon, consider refreshing",
+                    style="dim yellow",
+                )
+
+            return token
+
+        except Exception as e:
+            logger.warning(f"Failed to retrieve OAuth token for {provider}: {e}")
+            return None
+
+    def store_oauth_token(self, provider: str, token: OAuthToken) -> bool:
+        """Store OAuth token securely."""
+        service_name = self._get_service_name(provider)
+
+        try:
+            token_data = json.dumps(token.to_dict())
+            keyring.set_password(service_name, provider, token_data)
+            console.print(f"OAuth token for {provider} stored securely", style="green")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to store OAuth token for {provider}: {e}")
+            console.print(
+                f"Warning: Could not store OAuth token in keyring: {e}",
+                style="yellow",
+            )
+            return False
+
+    def update_oauth_token(
+        self, provider: str, access_token: str, expires_at: str | None = None
+    ) -> bool:
+        """Update only the access token (keep refresh token)."""
+        existing_token = self.get_oauth_token(provider)
+        if not existing_token:
+            console.print(
+                f"No existing OAuth token found for {provider}", style="yellow"
+            )
+            return False
+
+        # Update the access token while preserving refresh token
+        updated_token = OAuthToken(
+            access_token=access_token,
+            refresh_token=existing_token.refresh_token,
+            expires_at=expires_at,
+            token_type=existing_token.token_type,
+        )
+
+        return self.store_oauth_token(provider, updated_token)
+
+    def remove_oauth_token(self, provider: str) -> bool:
+        """Remove OAuth token from storage."""
+        service_name = self._get_service_name(provider)
+
+        try:
+            keyring.delete_password(service_name, provider)
+            console.print(f"OAuth token for {provider} removed", style="green")
+            return True
+        except keyring.errors.PasswordDeleteError:
+            # Token doesn't exist
+            return True
+        except Exception as e:
+            logger.error(f"Failed to remove OAuth token for {provider}: {e}")
+            console.print(f"Warning: Could not remove OAuth token: {e}", style="yellow")
+            return False
+
+    def has_oauth_token(self, provider: str) -> bool:
+        """Check if OAuth token exists for provider."""
+        return self.get_oauth_token(provider) is not None
+
+    def _get_service_name(self, provider: str) -> str:
+        """Get the keyring service name for OAuth tokens."""
+        return f"{self.service_prefix}-{provider}-oauth"

sqlsaber/config/settings.py

@@ -5,11 +5,13 @@ import os
 import platform
 import stat
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any
 
 import platformdirs
 
 from sqlsaber.config.api_keys import APIKeyManager
+from sqlsaber.config.auth import AuthConfigManager, AuthMethod
+from sqlsaber.config.oauth_flow import AnthropicOAuthFlow
 
 
 class ModelConfigManager:
@@ -40,7 +42,7 @@ class ModelConfigManager:
         except (OSError, PermissionError):
             pass
 
-    def _load_config(self) -> Dict[str, Any]:
+    def _load_config(self) -> dict[str, Any]:
         """Load configuration from file."""
         if not self.config_file.exists():
             return {"model": self.DEFAULT_MODEL}
@@ -55,7 +57,7 @@ class ModelConfigManager:
         except (json.JSONDecodeError, IOError):
             return {"model": self.DEFAULT_MODEL}
 
-    def _save_config(self, config: Dict[str, Any]) -> None:
+    def _save_config(self, config: dict[str, Any]) -> None:
         """Save configuration to file."""
         with open(self.config_file, "w") as f:
             json.dump(config, f, indent=2)
@@ -81,35 +83,48 @@ class Config:
         self.model_config_manager = ModelConfigManager()
         self.model_name = self.model_config_manager.get_model()
         self.api_key_manager = APIKeyManager()
-        self.api_key = self._get_api_key()
+        self.auth_config_manager = AuthConfigManager()
+        self.oauth_flow = AnthropicOAuthFlow()
+
+        # Get authentication credentials based on configured method
+        self.auth_method = self.auth_config_manager.get_auth_method()
+        self.api_key = None
+        self.oauth_token = None
+
+        if self.auth_method == AuthMethod.CLAUDE_PRO:
+            # Try to get OAuth token and refresh if needed
+            try:
+                token = self.oauth_flow.refresh_token_if_needed()
+                if token:
+                    self.oauth_token = token.access_token
+            except Exception:
+                # OAuth token unavailable, will need to re-authenticate
+                pass
+        else:
+            # Use API key authentication (default or explicitly configured)
+            self.api_key = self._get_api_key()
 
-    def _get_api_key(self) -> Optional[str]:
+    def _get_api_key(self) -> str | None:
         """Get API key for the model provider using cascading logic."""
         model = self.model_name
-
-        if model.startswith("openai:"):
-            return self.api_key_manager.get_api_key("openai")
-        elif model.startswith("anthropic:"):
+        if model.startswith("anthropic:"):
             return self.api_key_manager.get_api_key("anthropic")
-        else:
-            # For other providers, use generic key
-            return self.api_key_manager.get_api_key("generic")
 
     def set_model(self, model: str) -> None:
         """Set the model and update configuration."""
         self.model_config_manager.set_model(model)
         self.model_name = model
-        # Update API key for new model
-        self.api_key = self._get_api_key()
 
     def validate(self):
         """Validate that necessary configuration is present."""
+        # 1. Claude-Pro flow → require OAuth token only
+        if self.auth_method == AuthMethod.CLAUDE_PRO:
+            if not self.oauth_token:
+                raise ValueError(
+                    "OAuth token not available. Run 'saber auth setup' to authenticate with Claude Pro."
+                )
+            return  # OAuth path satisfied – nothing more to check
+
+        # 2. Default / API-key flow → require API key
         if not self.api_key:
-            model = self.model_name
-            provider = "generic"
-            if model.startswith("openai:"):
-                provider = "OpenAI"
-            elif model.startswith("anthropic:"):
-                provider = "Anthropic"
-
-            raise ValueError(f"{provider} API key not found.")
+            raise ValueError("Anthropic API key not found.")