sqlsaber 0.1.0__py3-none-any.whl

sqlsaber/config/api_keys.py

@@ -0,0 +1,102 @@
+"""API Key management for SQLSaber."""
+
+import getpass
+import os
+from typing import Optional
+
+import keyring
+from rich.console import Console
+
+console = Console()
+
+
+class APIKeyManager:
+    """Manages API keys with cascading retrieval: env var -> keyring -> prompt."""
+
+    def __init__(self):
+        self.service_prefix = "sqlsaber"
+
+    def get_api_key(self, provider: str) -> Optional[str]:
+        """Get API key for the specified provider using cascading logic."""
+        env_var_name = self._get_env_var_name(provider)
+        service_name = self._get_service_name(provider)
+
+        # 1. Check environment variable first
+        api_key = os.getenv(env_var_name)
+        if api_key:
+            console.print(f"Using {env_var_name} from environment", style="dim")
+            return api_key
+
+        # 2. Check keyring storage
+        try:
+            api_key = keyring.get_password(service_name, provider)
+            if api_key:
+                console.print(
+                    f"Using stored {provider} API key from keyring", style="dim"
+                )
+                return api_key
+        except Exception as e:
+            # Keyring access failed, continue to prompt
+            console.print(f"Keyring access failed: {e}", style="dim yellow")
+
+        # 3. Prompt user for API key
+        return self._prompt_and_store_key(provider, env_var_name, service_name)
+
+    def _get_env_var_name(self, provider: str) -> str:
+        """Get the expected environment variable name for a provider."""
+        if provider == "openai":
+            return "OPENAI_API_KEY"
+        elif provider == "anthropic":
+            return "ANTHROPIC_API_KEY"
+        else:
+            return "AI_API_KEY"
+
+    def _get_service_name(self, provider: str) -> str:
+        """Get the keyring service name for a provider."""
+        return f"{self.service_prefix}-{provider}-api-key"
+
+    def _prompt_and_store_key(
+        self, provider: str, env_var_name: str, service_name: str
+    ) -> Optional[str]:
+        """Prompt user for API key and store it in keyring."""
+        try:
+            console.print(
+                f"\n{provider.title()} API key not found in environment or keyring."
+            )
+            console.print("You can either:")
+            console.print(f"  1. Set the {env_var_name} environment variable")
+            console.print(
+                "  2. Enter it now to securely store using your operating system's credentials store"
+            )
+
+            api_key = getpass.getpass(
+                f"\nEnter your {provider.title()} API key (or press Enter to skip): "
+            )
+
+            if not api_key.strip():
+                console.print(
+                    "No API key provided. Some functionality may not work.",
+                    style="yellow",
+                )
+                return None
+
+            # Store in keyring for future use
+            try:
+                keyring.set_password(service_name, provider, api_key.strip())
+                console.print("API key stored securely for future use", style="green")
+            except Exception as e:
+                console.print(
+                    f"Warning: Could not store API key in keyring: {e}", style="yellow"
+                )
+                console.print(
+                    "You may need to enter it again next time", style="yellow"
+                )
+
+            return api_key.strip()
+
+        except KeyboardInterrupt:
+            console.print("\nOperation cancelled", style="yellow")
+            return None
+        except Exception as e:
+            console.print(f"Error prompting for API key: {e}", style="red")
+            return None

sqlsaber/config/database.py

@@ -0,0 +1,252 @@
+"""Database configuration management."""
+
+import json
+import os
+import platform
+import stat
+import keyring
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+from urllib.parse import quote_plus
+
+import platformdirs
+
+
+@dataclass
+class DatabaseConfig:
+    """Database connection configuration."""
+
+    name: str
+    type: str  # postgresql, mysql, sqlite
+    host: Optional[str]
+    port: Optional[int]
+    database: str
+    username: Optional[str]
+    password: Optional[str] = None
+    ssl_mode: Optional[str] = None
+    schema: Optional[str] = None
+
+    def to_connection_string(self) -> str:
+        """Convert config to database connection string."""
+        password = self.password or self._get_password_from_keyring()
+
+        if self.type == "postgresql":
+            if not all([self.host, self.port, self.username]):
+                raise ValueError("Host, port, and username are required for PostgreSQL")
+            if password:
+                encoded_password = quote_plus(password)
+                return f"postgresql://{self.username}:{encoded_password}@{self.host}:{self.port}/{self.database}"
+            else:
+                return f"postgresql://{self.username}@{self.host}:{self.port}/{self.database}"
+        elif self.type == "mysql":
+            if not all([self.host, self.port, self.username]):
+                raise ValueError("Host, port, and username are required for MySQL")
+            if password:
+                encoded_password = quote_plus(password)
+                return f"mysql://{self.username}:{encoded_password}@{self.host}:{self.port}/{self.database}"
+            else:
+                return (
+                    f"mysql://{self.username}@{self.host}:{self.port}/{self.database}"
+                )
+        elif self.type == "sqlite":
+            return f"sqlite:///{self.database}"
+        else:
+            raise ValueError(f"Unsupported database type: {self.type}")
+
+    def _get_password_from_keyring(self) -> Optional[str]:
+        """Get password from OS keyring."""
+        try:
+            return keyring.get_password("sqlsaber", f"{self.name}_{self.username}")
+        except Exception:
+            return None
+
+    def store_password_in_keyring(self, password: str) -> None:
+        """Store password in OS keyring."""
+        keyring.set_password("sqlsaber", f"{self.name}_{self.username}", password)
+
+    def delete_password_from_keyring(self) -> None:
+        """Delete password from OS keyring."""
+        try:
+            keyring.delete_password("sqlsaber", f"{self.name}_{self.username}")
+        except Exception:
+            pass
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "name": self.name,
+            "type": self.type,
+            "host": self.host,
+            "port": self.port,
+            "database": self.database,
+            "username": self.username,
+            "ssl_mode": self.ssl_mode,
+            "schema": self.schema,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatabaseConfig":
+        """Create from dictionary."""
+        return cls(
+            name=data["name"],
+            type=data["type"],
+            host=data["host"],
+            port=data["port"],
+            database=data["database"],
+            username=data["username"],
+            ssl_mode=data.get("ssl_mode"),
+            schema=data.get("schema"),
+        )
+
+
+class DatabaseConfigManager:
+    """Manages database configurations."""
+
+    def __init__(self):
+        self.config_dir = Path(platformdirs.user_config_dir("sqlsaber", "sqlsaber"))
+        self.config_file = self.config_dir / "database_config.json"
+        self._ensure_config_dir()
+
+    def _ensure_config_dir(self) -> None:
+        """Ensure config directory exists with proper permissions."""
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        self._set_secure_permissions(self.config_dir, is_directory=True)
+
+    def _set_secure_permissions(self, path: Path, is_directory: bool = False) -> None:
+        """Set secure permissions cross-platform."""
+        try:
+            if platform.system() == "Windows":
+                # On Windows, rely on NTFS permissions and avoid chmod
+                # The default permissions are usually sufficient for user-only access
+                return
+            else:
+                # Unix-like systems (Linux, macOS)
+                if is_directory:
+                    os.chmod(
+                        path, stat.S_IRWXU
+                    )  # 0o700 - owner read/write/execute only
+                else:
+                    os.chmod(
+                        path, stat.S_IRUSR | stat.S_IWUSR
+                    )  # 0o600 - owner read/write only
+        except (OSError, PermissionError):
+            # If we can't set permissions, continue anyway
+            # The directory/file creation should still work
+            pass
+
+    def _load_config(self) -> Dict[str, Any]:
+        """Load configuration from file."""
+        if not self.config_file.exists():
+            return {"default": None, "connections": {}}
+
+        try:
+            with open(self.config_file, "r") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, IOError):
+            return {"default": None, "connections": {}}
+
+    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)
+
+        # Set secure permissions cross-platform
+        self._set_secure_permissions(self.config_file, is_directory=False)
+
+    def add_database(
+        self, db_config: DatabaseConfig, password: Optional[str] = None
+    ) -> None:
+        """Add a database configuration."""
+        config = self._load_config()
+
+        # Check if database with this name already exists
+        if db_config.name in config["connections"]:
+            raise ValueError(f"Database '{db_config.name}' already exists")
+
+        # Store password in keyring if provided
+        if password:
+            db_config.store_password_in_keyring(password)
+
+        # Add to config
+        config["connections"][db_config.name] = db_config.to_dict()
+
+        # Set as default if it's the first one
+        if not config["default"]:
+            config["default"] = db_config.name
+
+        self._save_config(config)
+
+    def get_database(self, name: str) -> Optional[DatabaseConfig]:
+        """Get a database configuration by name."""
+        config = self._load_config()
+
+        if name not in config["connections"]:
+            return None
+
+        return DatabaseConfig.from_dict(config["connections"][name])
+
+    def get_default_database(self) -> Optional[DatabaseConfig]:
+        """Get the default database configuration."""
+        config = self._load_config()
+
+        default_name = config.get("default")
+        if not default_name:
+            return None
+
+        return self.get_database(default_name)
+
+    def list_databases(self) -> List[DatabaseConfig]:
+        """List all database configurations."""
+        config = self._load_config()
+
+        databases = []
+        for name, db_data in config["connections"].items():
+            databases.append(DatabaseConfig.from_dict(db_data))
+
+        return databases
+
+    def remove_database(self, name: str) -> bool:
+        """Remove a database configuration."""
+        config = self._load_config()
+
+        if name not in config["connections"]:
+            return False
+
+        # Remove password from keyring
+        db_config = DatabaseConfig.from_dict(config["connections"][name])
+        db_config.delete_password_from_keyring()
+
+        # Remove from config
+        del config["connections"][name]
+
+        # Update default if this was the default
+        if config["default"] == name:
+            remaining_connections = list(config["connections"].keys())
+            config["default"] = (
+                remaining_connections[0] if remaining_connections else None
+            )
+
+        self._save_config(config)
+        return True
+
+    def set_default_database(self, name: str) -> bool:
+        """Set the default database."""
+        config = self._load_config()
+
+        if name not in config["connections"]:
+            return False
+
+        config["default"] = name
+        self._save_config(config)
+        return True
+
+    def has_databases(self) -> bool:
+        """Check if any databases are configured."""
+        config = self._load_config()
+        return len(config["connections"]) > 0
+
+    def get_default_name(self) -> Optional[str]:
+        """Get the name of the default database."""
+        config = self._load_config()
+        return config.get("default")

sqlsaber/config/settings.py

@@ -0,0 +1,115 @@
+"""Configuration management for SQLSaber SQL Agent."""
+
+import json
+import os
+import platform
+import stat
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import platformdirs
+
+from sqlsaber.config.api_keys import APIKeyManager
+
+
+class ModelConfigManager:
+    """Manages model configuration persistence."""
+
+    DEFAULT_MODEL = "anthropic:claude-sonnet-4-20250514"
+
+    def __init__(self):
+        self.config_dir = Path(platformdirs.user_config_dir("sqlsaber", "sqlsaber"))
+        self.config_file = self.config_dir / "model_config.json"
+        self._ensure_config_dir()
+
+    def _ensure_config_dir(self) -> None:
+        """Ensure config directory exists with proper permissions."""
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        self._set_secure_permissions(self.config_dir, is_directory=True)
+
+    def _set_secure_permissions(self, path: Path, is_directory: bool = False) -> None:
+        """Set secure permissions cross-platform."""
+        try:
+            if platform.system() == "Windows":
+                return
+            else:
+                if is_directory:
+                    os.chmod(path, stat.S_IRWXU)  # 0o700
+                else:
+                    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)  # 0o600
+        except (OSError, PermissionError):
+            pass
+
+    def _load_config(self) -> Dict[str, Any]:
+        """Load configuration from file."""
+        if not self.config_file.exists():
+            return {"model": self.DEFAULT_MODEL}
+
+        try:
+            with open(self.config_file, "r") as f:
+                config = json.load(f)
+                # Ensure we have a model set
+                if "model" not in config:
+                    config["model"] = self.DEFAULT_MODEL
+                return config
+        except (json.JSONDecodeError, IOError):
+            return {"model": self.DEFAULT_MODEL}
+
+    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)
+
+        self._set_secure_permissions(self.config_file, is_directory=False)
+
+    def get_model(self) -> str:
+        """Get the configured model."""
+        config = self._load_config()
+        return config.get("model", self.DEFAULT_MODEL)
+
+    def set_model(self, model: str) -> None:
+        """Set the model configuration."""
+        config = self._load_config()
+        config["model"] = model
+        self._save_config(config)
+
+
+class Config:
+    """Configuration class for SQLSaber."""
+
+    def __init__(self):
+        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()
+
+    def _get_api_key(self) -> Optional[str]:
+        """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:"):
+            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."""
+        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.")

sqlsaber/database/__init__.py

@@ -0,0 +1,9 @@
+"""Database module for SQLSaber."""
+
+from .connection import DatabaseConnection
+from .schema import SchemaManager
+
+__all__ = [
+    "DatabaseConnection",
+    "SchemaManager",
+]

sqlsaber/database/connection.py

@@ -0,0 +1,187 @@
+"""Database connection management."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+from urllib.parse import urlparse
+
+import aiomysql
+import aiosqlite
+import asyncpg
+
+
+class BaseDatabaseConnection(ABC):
+    """Abstract base class for database connections."""
+
+    def __init__(self, connection_string: str):
+        self.connection_string = connection_string
+        self._pool = None
+
+    @abstractmethod
+    async def get_pool(self):
+        """Get or create connection pool."""
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """Close the connection pool."""
+        pass
+
+    @abstractmethod
+    async def execute_query(self, query: str, *args) -> List[Dict[str, Any]]:
+        """Execute a query and return results as list of dicts.
+
+        All queries run in a transaction that is rolled back at the end,
+        ensuring no changes are persisted to the database.
+        """
+        pass
+
+
+class PostgreSQLConnection(BaseDatabaseConnection):
+    """PostgreSQL database connection using asyncpg."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+        self._pool: Optional[asyncpg.Pool] = None
+
+    async def get_pool(self) -> asyncpg.Pool:
+        """Get or create connection pool."""
+        if self._pool is None:
+            self._pool = await asyncpg.create_pool(
+                self.connection_string, min_size=1, max_size=10
+            )
+        return self._pool
+
+    async def close(self):
+        """Close the connection pool."""
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+
+    async def execute_query(self, query: str, *args) -> List[Dict[str, Any]]:
+        """Execute a query and return results as list of dicts.
+
+        All queries run in a transaction that is rolled back at the end,
+        ensuring no changes are persisted to the database.
+        """
+        pool = await self.get_pool()
+        async with pool.acquire() as conn:
+            # Start a transaction that we'll always rollback
+            transaction = conn.transaction()
+            await transaction.start()
+
+            try:
+                rows = await conn.fetch(query, *args)
+                return [dict(row) for row in rows]
+            finally:
+                # Always rollback to ensure no changes are committed
+                await transaction.rollback()
+
+
+class MySQLConnection(BaseDatabaseConnection):
+    """MySQL database connection using aiomysql."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+        self._pool: Optional[aiomysql.Pool] = None
+        self._parse_connection_string()
+
+    def _parse_connection_string(self):
+        """Parse MySQL connection string into components."""
+        parsed = urlparse(self.connection_string)
+        self.host = parsed.hostname or "localhost"
+        self.port = parsed.port or 3306
+        self.database = parsed.path.lstrip("/") if parsed.path else ""
+        self.user = parsed.username or ""
+        self.password = parsed.password or ""
+
+    async def get_pool(self) -> aiomysql.Pool:
+        """Get or create connection pool."""
+        if self._pool is None:
+            self._pool = await aiomysql.create_pool(
+                host=self.host,
+                port=self.port,
+                user=self.user,
+                password=self.password,
+                db=self.database,
+                minsize=1,
+                maxsize=10,
+                autocommit=False,
+            )
+        return self._pool
+
+    async def close(self):
+        """Close the connection pool."""
+        if self._pool:
+            self._pool.close()
+            await self._pool.wait_closed()
+            self._pool = None
+
+    async def execute_query(self, query: str, *args) -> List[Dict[str, Any]]:
+        """Execute a query and return results as list of dicts.
+
+        All queries run in a transaction that is rolled back at the end,
+        ensuring no changes are persisted to the database.
+        """
+        pool = await self.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor(aiomysql.DictCursor) as cursor:
+                # Start transaction
+                await conn.begin()
+                try:
+                    await cursor.execute(query, args if args else None)
+                    rows = await cursor.fetchall()
+                    return [dict(row) for row in rows]
+                finally:
+                    # Always rollback to ensure no changes are committed
+                    await conn.rollback()
+
+
+class SQLiteConnection(BaseDatabaseConnection):
+    """SQLite database connection using aiosqlite."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+        # Extract database path from sqlite:///path format
+        self.database_path = connection_string.replace("sqlite:///", "")
+
+    async def get_pool(self):
+        """SQLite doesn't use connection pooling, return database path."""
+        return self.database_path
+
+    async def close(self):
+        """SQLite connections are created per query, no persistent pool to close."""
+        pass
+
+    async def execute_query(self, query: str, *args) -> List[Dict[str, Any]]:
+        """Execute a query and return results as list of dicts.
+
+        All queries run in a transaction that is rolled back at the end,
+        ensuring no changes are persisted to the database.
+        """
+        async with aiosqlite.connect(self.database_path) as conn:
+            # Enable row factory for dict-like access
+            conn.row_factory = aiosqlite.Row
+
+            # Start transaction
+            await conn.execute("BEGIN")
+            try:
+                cursor = await conn.execute(query, args if args else ())
+                rows = await cursor.fetchall()
+                return [dict(row) for row in rows]
+            finally:
+                # Always rollback to ensure no changes are committed
+                await conn.rollback()
+
+
+def DatabaseConnection(connection_string: str) -> BaseDatabaseConnection:
+    """Factory function to create appropriate database connection based on connection string."""
+    if connection_string.startswith("postgresql://"):
+        return PostgreSQLConnection(connection_string)
+    elif connection_string.startswith("mysql://"):
+        return MySQLConnection(connection_string)
+    elif connection_string.startswith("sqlite:///"):
+        return SQLiteConnection(connection_string)
+    else:
+        raise ValueError(
+            f"Unsupported database type in connection string: {connection_string}"
+        )