src/db/database.py - master - yakobsonsa/bi_agent

src/db/database.py
296 строк11 KB
Yakobsonsa
fix
10 фев 2026, 18:44
10 фев 2026, 18:442671a73
"""Database connection and query execution."""

import logging
from typing import List, Dict, Any, Optional, Tuple, Union
import asyncio
from src.settings import get_settings
# from contextlib import asynccontextmanager

logger = logging.getLogger(__name__)


class DatabaseError(Exception):
    """Database related errors."""
    pass


class AsyncDatabase:
    """Async database wrapper for multiple database types."""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """Initialize database from config or settings."""
        settings = get_settings()
        # Use provided config or get from settings
        if config:
            self.config = config
        else:
            self.config = {
                'type': settings.database.type,
                'host': settings.database.host,
                'port': settings.database.port,
                'user': settings.database.user,
                'password': settings.database.password,
                'database': settings.database.database,
                'sqlite_path': settings.database.sqlite_path,
            }
        
        self.db_type = self.config.get('type', 'postgresql')
        self.connection = None
        self.pool = None
        self.pool_min_size = settings.database.pool_min_size
        self.pool_max_size = settings.database.pool_max_size
        self.connection_timeout = settings.database.connection_timeout
        self.query_timeout = settings.database.query_timeout

    async def connect(self):
        """Establish database connection."""
        if self.db_type == 'postgresql':
            await self._connect_postgresql()
        elif self.db_type == 'mysql':
            await self._connect_mysql()
        elif self.db_type == 'sqlite':
            await self._connect_sqlite()
        else:
            raise DatabaseError(f"Unsupported database type: {self.db_type}")
        
        logger.info(f"Connected to {self.db_type} database (pool: {self.pool_min_size}-{self.pool_max_size})")

    async def _connect_postgresql(self):
        """Connect to PostgreSQL."""
        import asyncpg
        
        self.pool = await asyncpg.create_pool(
            host=self.config.get('host'),
            port=self.config.get('port', 5432),
            user=self.config.get('user'),
            password=self.config.get('password'),
            database=self.config.get('database'),
            min_size=self.pool_min_size,
            max_size=self.pool_max_size,
            command_timeout=self.connection_timeout
        )

    async def _connect_mysql(self):
        """Connect to MySQL."""
        import aiomysql
        
        self.pool = await aiomysql.create_pool(
            host=self.config.get('host'),
            port=self.config.get('port', 3306),
            user=self.config.get('user'),
            password=self.config.get('password'),
            db=self.config.get('database'),
            minsize=self.pool_min_size,
            maxsize=self.pool_max_size
        )

    async def _connect_sqlite(self):
        """Connect to SQLite."""
        import aiosqlite
        
        db_path = self.config.get('sqlite_path', './data/revenue.db')
        self.connection = await aiosqlite.connect(db_path, timeout=self.connection_timeout)
        logger.info(f"Connected to SQLite at {db_path}")

    async def disconnect(self):
        """Close database connection."""
        if self.db_type == 'sqlite':
            if self.connection:
                await self.connection.close()
        elif self.pool:
            await self.pool.close()
        logger.info("Database connection closed")

    async def execute_query(
        self,
        query: str,
        params: Optional[tuple] = None,
        timeout: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Execute SELECT query and return results."""
        query_timeout = timeout or self.query_timeout
        try:
            if self.db_type == 'postgresql':
                return await self._execute_postgresql(query, params, query_timeout)
            elif self.db_type == 'mysql':
                return await self._execute_mysql(query, params, query_timeout)
            elif self.db_type == 'sqlite':
                return await self._execute_sqlite(query, params, query_timeout)
        except Exception as e:
            logger.error(f"Query execution error: {e}\nQuery: {query}")
            raise DatabaseError(f"Query failed: {str(e)}") from e

    async def _execute_postgresql(
        self,
        query: str,
        params: Optional[tuple],
        timeout: int
    ) -> List[Dict[str, Any]]:
        """Execute PostgreSQL query."""
        async with self.pool.acquire() as connection:
            rows = await asyncio.wait_for(
                connection.fetch(query, *(params or [])),
                timeout=timeout
            )
            return [dict(row) for row in rows]

    async def _execute_mysql(
        self,
        query: str,
        params: Optional[tuple],
        timeout: int
    ) -> List[Dict[str, Any]]:
        """Execute MySQL query."""
        async with self.pool.acquire() as connection:
            async with connection.cursor() as cursor:
                await asyncio.wait_for(
                    cursor.execute(query, params or []),
                    timeout=timeout,
                )
                columns = [desc[0] for desc in cursor.description]
                rows = await asyncio.wait_for(
                    cursor.fetchall(),
                    timeout=timeout,
                )
                return [dict(zip(columns, row)) for row in rows]

    async def _execute_sqlite(
        self,
        query: str,
        params: Optional[tuple],
        timeout: int
    ) -> List[Dict[str, Any]]:
        """Execute SQLite query."""
        cursor = await asyncio.wait_for(
            self.connection.execute(query, params or []),
            timeout=timeout,
        )
        columns = [desc[0] for desc in cursor.description]
        rows = await asyncio.wait_for(cursor.fetchall(), timeout=timeout)
        return [dict(zip(columns, row)) for row in rows]

    async def get_schema(self) -> str:
        """Get database schema for context."""
        if self.db_type == 'postgresql':
            schema_query = """
            SELECT table_name, column_name, data_type
            FROM information_schema.columns
            WHERE table_schema = 'public'
            ORDER BY table_name, ordinal_position
            """
        elif self.db_type == 'mysql':
            schema_query = """
            SELECT TABLE_NAME, COLUMN_NAME, COLUMN_TYPE
            FROM INFORMATION_SCHEMA.COLUMNS
            WHERE TABLE_SCHEMA = DATABASE()
            ORDER BY TABLE_NAME, ORDINAL_POSITION
            """
        elif self.db_type == 'sqlite':
            # SQLite schema is more limited
            schema_query = (
                "SELECT name, type FROM sqlite_master WHERE type='table'"
            )
        else:
            return ""
        
        try:
            results = await self.execute_query(schema_query)
            
            schema_text = "Database Schema:\n\n"
            current_table = None
            
            for row in results:
                if self.db_type == 'sqlite':
                    table_name = row.get('name')
                    schema_text += f"TABLE: {table_name}\n"
                else:
                    table_name = row.get('table_name') or row.get('TABLE_NAME')
                    column_name = (
                        row.get('column_name') or row.get('COLUMN_NAME')
                    )
                    data_type = row.get('data_type') or row.get('COLUMN_TYPE')
                    
                    if table_name != current_table:
                        schema_text += f"\nTABLE: {table_name}\n"
                        current_table = table_name
                    
                    schema_text += f"  - {column_name}: {data_type}\n"
            
            return schema_text
        except Exception as e:
            logger.error(f"Error fetching schema: {e}")
            return ""

    async def find_similar_values(
        self,
        table_name: str,
        column_name: str,
        search_value: str,
        limit: int = 5,
        return_scores: bool = False
    ) -> Union[List[str], List[Tuple[str, int]]]:
        """
        Find similar values in a column to help correct typos/misspellings.
        
        Args:
            table_name: Table to search in
            column_name: Column to search
            search_value: Value to search for (can be partial or misspelled)
            limit: Maximum results to return
            return_scores: If True, return tuples of (value, score) instead of just values
            
        Returns:
            List of similar values from the column, or list of (value, score) tuples if return_scores=True
        """
        try:
            if not search_value:
                return []
            
            search_lower = search_value.lower()
            
            # Get all distinct values from the column
            query = f"SELECT DISTINCT {column_name} FROM {table_name} WHERE {column_name} IS NOT NULL LIMIT 1000"
            results = await self.execute_query(query)
            
            if not results:
                return []
            
            # Score matches: exact match, starts with, contains, similar
            matches = []
            for row in results:
                value = row.get(column_name, "")
                if not value:
                    continue
                
                value_lower = str(value).lower()
                
                # Exact match (highest priority)
                if value_lower == search_lower:
                    matches.insert(0, (value, 1000))
                # Starts with
                elif value_lower.startswith(search_lower):
                    matches.append((value, 100))
                # Contains
                elif search_lower in value_lower:
                    matches.append((value, 50))
                # Fuzzy match (Levenshtein-like: count common characters)
                else:
                    common = sum(1 for c in search_lower if c in value_lower)
                    if common >= len(search_lower) * 0.6:  # At least 60% match
                        matches.append((value, common))
            
            # Sort by score (descending) and return unique values
            matches.sort(key=lambda x: x[1], reverse=True)
            seen = set()
            result = []
            for value, score in matches:
                if value not in seen:
                    seen.add(value)
                    result.append((value, score) if return_scores else value)
                    if len(result) >= limit:
                        break
            
            return result
        except Exception as e:
            logger.error(f"Error finding similar values in {table_name}.{column_name}: {e}")
            return []