tuskdata 0.1.0__py3-none-any.whl

tusk/__init__.py

@@ -0,0 +1,3 @@
+"""Tusk - Modern data platform"""
+
+__version__ = "0.1.0"

tusk/admin/__init__.py

@@ -0,0 +1,42 @@
+"""PostgreSQL administration module"""
+
+from tusk.admin.stats import get_server_stats, ServerStats
+from tusk.admin.processes import get_active_queries, kill_query, ActiveQuery
+from tusk.admin.backup import create_backup, list_backups, get_backup_path
+from tusk.admin.extensions import (
+    get_extensions,
+    get_installed_extensions,
+    install_extension,
+    uninstall_extension,
+    Extension,
+)
+from tusk.admin.maintenance import (
+    get_locks,
+    get_all_locks,
+    get_table_bloat,
+    vacuum_table,
+    analyze_table,
+    reindex_table,
+)
+
+__all__ = [
+    "get_server_stats",
+    "ServerStats",
+    "get_active_queries",
+    "kill_query",
+    "ActiveQuery",
+    "create_backup",
+    "list_backups",
+    "get_backup_path",
+    "get_extensions",
+    "get_installed_extensions",
+    "install_extension",
+    "uninstall_extension",
+    "Extension",
+    "get_locks",
+    "get_all_locks",
+    "get_table_bloat",
+    "vacuum_table",
+    "analyze_table",
+    "reindex_table",
+]

tusk/admin/backup.py

@@ -0,0 +1,331 @@
+"""Backup and restore functionality for PostgreSQL"""
+
+import subprocess
+import os
+from pathlib import Path
+from datetime import datetime
+
+from tusk.core.connection import ConnectionConfig, TUSK_DIR
+from tusk.core.config import get_config
+
+BACKUP_DIR = TUSK_DIR / "backups"
+
+
+def _get_pg_bin_search_paths() -> list[Path]:
+    """Get list of paths to search for PostgreSQL binaries"""
+    paths = []
+    seen_resolved = set()
+
+    def add_path(p: Path):
+        """Add path if not already seen (resolves symlinks to avoid duplicates)"""
+        try:
+            resolved = p.resolve()
+            if resolved not in seen_resolved and resolved.exists():
+                seen_resolved.add(resolved)
+                paths.append(p)
+        except Exception:
+            pass
+
+    # Postgres.app (macOS) - check latest first
+    add_path(Path("/Applications/Postgres.app/Contents/Versions/latest/bin"))
+
+    # Add versioned Postgres.app paths
+    pg_app_versions = Path("/Applications/Postgres.app/Contents/Versions")
+    if pg_app_versions.exists():
+        for version_dir in sorted(pg_app_versions.iterdir(), reverse=True):
+            if version_dir.name != "latest" and version_dir.is_dir():
+                add_path(version_dir / "bin")
+
+    # Homebrew (macOS)
+    for p in [
+        Path("/opt/homebrew/opt/postgresql/bin"),
+        Path("/opt/homebrew/bin"),
+        Path("/usr/local/opt/postgresql/bin"),
+        Path("/usr/local/bin"),
+    ]:
+        add_path(p)
+
+    # Linux common paths
+    for p in [
+        Path("/usr/lib/postgresql/16/bin"),
+        Path("/usr/lib/postgresql/15/bin"),
+        Path("/usr/lib/postgresql/14/bin"),
+        Path("/usr/pgsql-16/bin"),
+        Path("/usr/pgsql-15/bin"),
+        Path("/usr/bin"),
+    ]:
+        add_path(p)
+
+    return paths
+
+
+def _find_pg_binary(name: str) -> str:
+    """Find a PostgreSQL binary (pg_dump, psql, etc.)
+
+    Priority:
+    1. User-configured pg_bin_path in ~/.tusk/config.toml
+    2. System PATH
+    3. Common installation locations (Postgres.app, Homebrew, etc.)
+    """
+    import shutil
+
+    # First check user-configured path
+    config = get_config()
+    if config.pg_bin_path:
+        configured_path = Path(config.pg_bin_path) / name
+        if configured_path.exists() and configured_path.is_file():
+            return str(configured_path)
+
+    # Then check if it's in PATH
+    if shutil.which(name):
+        return name
+
+    # Search in common locations
+    for search_path in _get_pg_bin_search_paths():
+        binary_path = search_path / name
+        if binary_path.exists() and binary_path.is_file():
+            return str(binary_path)
+
+    # Fallback to just the name (will fail with clear error if not found)
+    return name
+
+
+def get_pg_dump_path() -> str:
+    """Get the path to pg_dump binary"""
+    return _find_pg_binary("pg_dump")
+
+
+def get_psql_path() -> str:
+    """Get the path to psql binary"""
+    return _find_pg_binary("psql")
+
+
+def create_backup(config: ConnectionConfig) -> tuple[bool, str, Path | None]:
+    """Create a pg_dump backup of the database
+
+    Returns: (success, message, filepath)
+    """
+    BACKUP_DIR.mkdir(parents=True, exist_ok=True)
+
+    timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
+    filename = f"{config.database}_{timestamp}.sql.gz"
+    filepath = BACKUP_DIR / filename
+
+    pg_dump = get_pg_dump_path()
+
+    cmd = [
+        pg_dump,
+        "-h", config.host or "localhost",
+        "-p", str(config.port),
+        "-U", config.user or "postgres",
+        "-d", config.database or "postgres",
+        "--format=plain",
+    ]
+
+    env = os.environ.copy()
+    if config.password:
+        env["PGPASSWORD"] = config.password
+
+    try:
+        # pg_dump | gzip > file
+        with open(filepath, "wb") as f:
+            dump_proc = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                env=env
+            )
+            gzip_proc = subprocess.Popen(
+                ["gzip"],
+                stdin=dump_proc.stdout,
+                stdout=f,
+                stderr=subprocess.PIPE
+            )
+
+            dump_proc.stdout.close()
+            gzip_proc.communicate()
+            dump_proc.wait()
+
+            if dump_proc.returncode != 0:
+                stderr = dump_proc.stderr.read().decode() if dump_proc.stderr else "Unknown error"
+                filepath.unlink(missing_ok=True)
+                return False, f"pg_dump failed: {stderr}", None
+
+        size = filepath.stat().st_size
+        size_human = f"{size / 1024 / 1024:.2f} MB" if size > 1024 * 1024 else f"{size / 1024:.1f} KB"
+        return True, f"Backup created: {filename} ({size_human})", filepath
+
+    except FileNotFoundError:
+        return False, f"pg_dump not found at '{pg_dump}'. Install PostgreSQL client tools or check Postgres.app installation.", None
+    except Exception as e:
+        filepath.unlink(missing_ok=True)
+        return False, f"Backup failed: {str(e)}", None
+
+
+def list_backups() -> list[dict]:
+    """List all available backups"""
+    if not BACKUP_DIR.exists():
+        return []
+
+    backups = []
+    for f in sorted(BACKUP_DIR.glob("*.sql.gz"), reverse=True):
+        stat = f.stat()
+        backups.append({
+            "filename": f.name,
+            "size_bytes": stat.st_size,
+            "size_human": f"{stat.st_size / 1024 / 1024:.2f} MB" if stat.st_size > 1024 * 1024 else f"{stat.st_size / 1024:.1f} KB",
+            "created": datetime.fromtimestamp(stat.st_mtime).isoformat(),
+        })
+
+    return backups
+
+
+def get_backup_path(filename: str) -> Path | None:
+    """Get full path to a backup file (for download)"""
+    filepath = BACKUP_DIR / filename
+    if filepath.exists() and filepath.is_file():
+        return filepath
+    return None
+
+
+def restore_backup(config: ConnectionConfig, filename: str) -> tuple[bool, str]:
+    """Restore a database from backup
+
+    WARNING: This will overwrite the target database!
+    """
+    filepath = get_backup_path(filename)
+    if not filepath:
+        return False, f"Backup file not found: {filename}"
+
+    psql = get_psql_path()
+
+    cmd = [
+        psql,
+        "-h", config.host or "localhost",
+        "-p", str(config.port),
+        "-U", config.user or "postgres",
+        "-d", config.database or "postgres",
+    ]
+
+    env = os.environ.copy()
+    if config.password:
+        env["PGPASSWORD"] = config.password
+
+    try:
+        # gunzip < file | psql
+        with open(filepath, "rb") as f:
+            gunzip_proc = subprocess.Popen(
+                ["gunzip", "-c"],
+                stdin=f,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE
+            )
+            psql_proc = subprocess.Popen(
+                cmd,
+                stdin=gunzip_proc.stdout,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                env=env
+            )
+
+            gunzip_proc.stdout.close()
+            stdout, stderr = psql_proc.communicate()
+
+            if psql_proc.returncode != 0:
+                return False, f"Restore failed: {stderr.decode()}"
+
+        return True, f"Database restored from {filename}"
+
+    except FileNotFoundError:
+        return False, f"psql not found at '{psql}'. Install PostgreSQL client tools or check Postgres.app installation."
+    except Exception as e:
+        return False, f"Restore failed: {str(e)}"
+
+
+def get_createdb_path() -> str:
+    """Get the path to createdb binary"""
+    return _find_pg_binary("createdb")
+
+
+def create_database(config: ConnectionConfig, db_name: str, owner: str | None = None) -> tuple[bool, str]:
+    """Create a new database on the PostgreSQL server
+
+    Args:
+        config: Connection config (uses host, port, user, password)
+        db_name: Name for the new database
+        owner: Optional owner for the database
+    """
+    createdb = get_createdb_path()
+
+    cmd = [
+        createdb,
+        "-h", config.host or "localhost",
+        "-p", str(config.port),
+        "-U", config.user or "postgres",
+    ]
+
+    if owner:
+        cmd.extend(["-O", owner])
+
+    cmd.append(db_name)
+
+    env = os.environ.copy()
+    if config.password:
+        env["PGPASSWORD"] = config.password
+
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            env=env
+        )
+
+        if result.returncode != 0:
+            return False, f"Create database failed: {result.stderr}"
+
+        return True, f"Database '{db_name}' created successfully"
+
+    except FileNotFoundError:
+        return False, f"createdb not found at '{createdb}'. Install PostgreSQL client tools."
+    except Exception as e:
+        return False, f"Create database failed: {str(e)}"
+
+
+def create_database_from_backup(
+    config: ConnectionConfig,
+    filename: str,
+    new_db_name: str,
+    owner: str | None = None
+) -> tuple[bool, str]:
+    """Create a new database and restore from backup
+
+    Args:
+        config: Connection config for the PostgreSQL server
+        filename: Backup file to restore from
+        new_db_name: Name for the new database
+        owner: Optional owner for the new database
+    """
+    # First create the new database
+    success, message = create_database(config, new_db_name, owner)
+    if not success:
+        return False, message
+
+    # Create a modified config pointing to the new database
+    new_config = ConnectionConfig(
+        id=config.id,
+        name=config.name,
+        type=config.type,
+        host=config.host,
+        port=config.port,
+        user=config.user,
+        password=config.password,
+        database=new_db_name,
+    )
+
+    # Restore the backup to the new database
+    success, message = restore_backup(new_config, filename)
+    if not success:
+        return False, f"Database created but restore failed: {message}"
+
+    return True, f"Database '{new_db_name}' created and restored from {filename}"

tusk/admin/extensions.py

@@ -0,0 +1,148 @@
+"""PostgreSQL extension management"""
+
+import msgspec
+from ..core.connection import ConnectionConfig
+from ..engines.postgres import execute_query
+
+
+class Extension(msgspec.Struct):
+    """PostgreSQL extension info"""
+    name: str
+    installed_version: str | None = None
+    default_version: str | None = None
+    description: str | None = None
+    is_installed: bool = False
+
+
+async def get_extensions(config: ConnectionConfig) -> list[Extension]:
+    """Get all available extensions with their installation status"""
+
+    # Query to get all available extensions and their installation status
+    sql = """
+    SELECT
+        a.name,
+        i.extversion as installed_version,
+        a.default_version,
+        a.comment as description,
+        CASE WHEN i.extname IS NOT NULL THEN true ELSE false END as is_installed
+    FROM pg_available_extensions a
+    LEFT JOIN pg_extension i ON a.name = i.extname
+    ORDER BY
+        CASE WHEN i.extname IS NOT NULL THEN 0 ELSE 1 END,
+        a.name
+    """
+
+    result = await execute_query(config, sql)
+
+    extensions = []
+    for row in result.rows:
+        extensions.append(Extension(
+            name=row[0],
+            installed_version=row[1],
+            default_version=row[2],
+            description=row[3],
+            is_installed=row[4]
+        ))
+
+    return extensions
+
+
+async def get_installed_extensions(config: ConnectionConfig) -> list[Extension]:
+    """Get only installed extensions"""
+
+    sql = """
+    SELECT
+        e.extname as name,
+        e.extversion as installed_version,
+        a.default_version,
+        a.comment as description,
+        true as is_installed
+    FROM pg_extension e
+    LEFT JOIN pg_available_extensions a ON e.extname = a.name
+    ORDER BY e.extname
+    """
+
+    result = await execute_query(config, sql)
+
+    extensions = []
+    for row in result.rows:
+        extensions.append(Extension(
+            name=row[0],
+            installed_version=row[1],
+            default_version=row[2],
+            description=row[3],
+            is_installed=True
+        ))
+
+    return extensions
+
+
+async def install_extension(config: ConnectionConfig, name: str, schema: str | None = None) -> bool:
+    """Install an extension"""
+
+    # Validate extension name (prevent SQL injection)
+    if not name.replace("_", "").replace("-", "").isalnum():
+        raise ValueError(f"Invalid extension name: {name}")
+
+    sql = f'CREATE EXTENSION IF NOT EXISTS "{name}"'
+    if schema:
+        # Validate schema name
+        if not schema.replace("_", "").isalnum():
+            raise ValueError(f"Invalid schema name: {schema}")
+        sql += f' SCHEMA "{schema}"'
+
+    await execute_query(config, sql)
+    return True
+
+
+async def uninstall_extension(config: ConnectionConfig, name: str, cascade: bool = False) -> bool:
+    """Uninstall an extension"""
+
+    # Validate extension name (prevent SQL injection)
+    if not name.replace("_", "").replace("-", "").isalnum():
+        raise ValueError(f"Invalid extension name: {name}")
+
+    sql = f'DROP EXTENSION IF EXISTS "{name}"'
+    if cascade:
+        sql += " CASCADE"
+
+    await execute_query(config, sql)
+    return True
+
+
+async def get_extension_details(config: ConnectionConfig, name: str) -> dict:
+    """Get detailed information about an extension"""
+
+    # Validate extension name
+    if not name.replace("_", "").replace("-", "").isalnum():
+        raise ValueError(f"Invalid extension name: {name}")
+
+    sql = f"""
+    SELECT
+        a.name,
+        a.default_version,
+        a.comment as description,
+        e.extversion as installed_version,
+        e.extrelocatable as relocatable,
+        n.nspname as schema
+    FROM pg_available_extensions a
+    LEFT JOIN pg_extension e ON a.name = e.extname
+    LEFT JOIN pg_namespace n ON e.extnamespace = n.oid
+    WHERE a.name = '{name}'
+    """
+
+    result = await execute_query(config, sql)
+
+    if not result.rows:
+        return {"error": f"Extension '{name}' not found"}
+
+    row = result.rows[0]
+    return {
+        "name": row[0],
+        "default_version": row[1],
+        "description": row[2],
+        "installed_version": row[3],
+        "relocatable": row[4],
+        "schema": row[5],
+        "is_installed": row[3] is not None
+    }