Research-Stack/4-Infrastructure/NoDupeLabs/nodupe/tools/databases/query.py

# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2025 Allaun

"""Database query functionality."""

from typing import Any, Dict, List, Optional, Tuple
import os


class DatabaseQuery:
    """Database query functionality."""

    def __init__(self, db):
        """Initialize query."""
        self.db = db

    def execute(self, query: str, params: Optional[Tuple] = None) -> List[Dict[str, Any]]:
        """Execute query and return results."""
        if hasattr(self.db, 'get_connection'):
            conn = self.db.get_connection()
        else:
            conn = self.db.connect()
        cursor = conn.cursor()
        cursor.execute(query, params or ())
        results = []

        for row in cursor.fetchall():
            results.append(dict(zip([d[0] for d in cursor.description], row)))

        return results


class DatabaseBatch:
    """Database batch operations."""

    def __init__(self, db):
        """Initialize batch operations."""
        self.db = db

    def execute_batch(self, operations: List[Tuple[str, Tuple]]) -> None:
        """Execute batch operations."""
        if hasattr(self.db, 'get_connection'):
            conn = self.db.get_connection()
        else:
            conn = self.db.connect()
        cursor = conn.cursor()

        for query, params in operations:
            cursor.execute(query, params)

        conn.commit()

    def execute_transaction_batch(self, operations: List[Tuple[str, Tuple]]) -> None:
        """Execute batch operations within a transaction."""
        if hasattr(self.db, 'get_connection'):
            conn = self.db.get_connection()
        else:
            conn = self.db.connect()
        cursor = conn.cursor()

        try:
            for query, params in operations:
                cursor.execute(query, params)

            conn.commit()
        except Exception:
            conn.rollback()
            raise


class DatabasePerformance:
    """Database performance monitoring."""

    def __init__(self, db):
        """Initialize performance monitoring."""
        self.db = db
        self._metrics = {
            'queries': 0,
            'total_time': 0.0,
            'avg_time': 0.0,
        }

    def get_metrics(self) -> Dict[str, Any]:
        """Get performance metrics."""
        # Return in expected format with 'metrics' or 'error' key
        return {'metrics': self._metrics.copy()}

    def record_query(self, query_time: float) -> None:
        """Record query execution time."""
        self._metrics['queries'] += 1
        self._metrics['total_time'] += query_time
        if self._metrics['queries'] > 0:
            self._metrics['avg_time'] = self._metrics['total_time'] / self._metrics['queries']

    def monitor_performance(self):
        """Context manager for performance monitoring."""
        return self.db.monitoring

    def get_results(self):
        """Get performance results."""
        return self.db.monitoring.get_metrics()


class DatabaseIntegrity:
    """Database integrity checking."""

    def __init__(self, db):
        """Initialize integrity checking."""
        self.db = db

    def validate(self) -> Dict[str, Any]:
        """Validate database integrity."""
        return {'valid': True, 'errors': [], 'tables': []}

    def check_integrity(self) -> Dict[str, Any]:
        """Check database integrity."""
        # Return in expected format with 'tables' and 'indexes' keys
        return {'valid': True, 'errors': [], 'tables': [], 'indexes': []}


class DatabaseBackup:
    """Database backup functionality."""

    def __init__(self, db):
        """Initialize backup."""
        self.db = db

    def create_backup(self, backup_path: str) -> None:
        """Create database backup."""
        import shutil
        # DatabaseConnection may expose `path` or `db_path` (absolute path or ":memory:")
        # Prefer explicit string attributes (tests set `path`) and avoid
        # blindly using dynamically-created MagicMock attributes.
        src = None
        if hasattr(self.db, 'path') and isinstance(getattr(self.db, 'path'), (str, bytes, bytearray, os.PathLike)):
            src = getattr(self.db, 'path')
        elif hasattr(self.db, 'db_path') and isinstance(getattr(self.db, 'db_path'), (str, bytes, bytearray, os.PathLike)):
            src = getattr(self.db, 'db_path')
        if src is None:
            raise ValueError("Database path not available for backup")
        if isinstance(src, (bytes, bytearray)):
            src = src.decode('utf-8')
        if src is not None:
            src = str(src)
        if src == ":memory:" or src == ":memory":
            raise ValueError("Cannot backup in-memory database")
        # Ensure source path exists and is a filesystem path (avoid copying MagicMocks)
        if not src or not os.path.exists(src):
            raise ValueError("Database path not available for backup")
        shutil.copy2(src, backup_path)

    def restore_backup(self, backup_path: str, restore_path: str) -> None:
        """Restore database from backup."""
        import shutil
        shutil.copy2(backup_path, restore_path)


class DatabaseMigration:
    """Database migration functionality."""

    def __init__(self, db):
        """Initialize migration."""
        self.db = db

    def migrate_schema(self, migrations: Dict[str, Dict[str, List[str]]]) -> None:
        """Migrate database schema."""
        pass  # Implementation would apply migrations

    def migrate_data(self, table_name: str, transformations: Dict[str, str], new_columns: Optional[List[str]] = None) -> None:
        """Migrate data in the specified table."""
        pass  # Implementation would transform data


class DatabaseRecovery:
    """Database recovery functionality."""

    def __init__(self, db):
        """Initialize recovery."""
        self.db = db

    def handle_errors(self, raise_on_error: bool = False):
        """Handle database errors."""
        try:
            # Check database integrity
            integrity = self.db.integrity.check_integrity()
            if not integrity.get('valid', True):
                if raise_on_error:
                    raise Exception("Database integrity check failed")
                return False
            return True
        except Exception as _:
            if raise_on_error:
                raise
            return False


class DatabaseOptimization:
    """Database optimization functionality."""

    def __init__(self, db):
        """Initialize optimization."""
        self.db = db

    def optimize_query(self, query: str) -> str:
        """Optimize database query."""
        # Basic query optimization
        optimized = query.strip()
        if optimized.endswith(';'):
            optimized = optimized[:-1]
        return optimized