src/agent/sql_normalizer.py - master - yakobsonsa/bi_agent

src/agent/sql_normalizer.py
158 строк6 KB
Yakobsonsa
vers 2.0
02 фев 2026, 00:06
02 фев 2026, 00:0672f33ab
"""SQL dialect converter - normalizes MySQL/PostgreSQL SQL to SQLite."""

import re
import logging
from typing import Dict

logger = logging.getLogger(__name__)


class SQLNormalizer:
    """Converts SQL from MySQL/PostgreSQL dialect to SQLite compatible syntax."""
    
    # Mapping of MySQL/PostgreSQL functions to SQLite equivalents
    FUNCTION_MAPPINGS: Dict[str, str] = {
        # Date functions
        r'\bMONTH\s*\(\s*([^)]+)\s*\)': r"strftime('%m', \1)",
        r'\bYEAR\s*\(\s*([^)]+)\s*\)': r"strftime('%Y', \1)",
        r'\bDAY\s*\(\s*([^)]+)\s*\)': r"strftime('%d', \1)",
        r'\bQUARTER\s*\(\s*([^)]+)\s*\)': r"(CAST(strftime('%m', \1) AS INTEGER) - 1) / 3 + 1",
        r'\bWEEK\s*\(\s*([^)]+)\s*\)': r"strftime('%W', \1)",
        r'\bWEEKDAY\s*\(\s*([^)]+)\s*\)': r"(CAST(strftime('%w', \1) AS INTEGER) + 6) % 7",
        r'\bDAYOFWEEK\s*\(\s*([^)]+)\s*\)': r"strftime('%w', \1)",
        r'\bDAYOFMONTH\s*\(\s*([^)]+)\s*\)': r"strftime('%d', \1)",
        r'\bDAYOFYEAR\s*\(\s*([^)]+)\s*\)': r"strftime('%j', \1)",
        
        # Time functions
        r'\bNOW\(\)': r"datetime('now')",
        r'\bCURRENT_TIMESTAMP': r"datetime('now')",
        r'\bCURDATE\(\)': r"date('now')",
        r'\bCURTIME\(\)': r"time('now')",
        r'\bHOUR\s*\(\s*([^)]+)\s*\)': r"strftime('%H', \1)",
        r'\bMINUTE\s*\(\s*([^)]+)\s*\)': r"strftime('%M', \1)",
        r'\bSECOND\s*\(\s*([^)]+)\s*\)': r"strftime('%S', \1)",
        
        # String functions - these exist in SQLite but may have different names
        r'\bCONCAT_WS\s*\(\s*([^,]+)\s*,': r"GROUP_CONCAT(",  # Basic - needs manual adjustment
        
        # Math functions (usually same, but adding for completeness)
        r'\bABS\s*\(': r"ABS(",
        r'\bROUND\s*\(': r"ROUND(",
        r'\bSQRT\s*\(': r"SQRT(",
    }
    
    @classmethod
    def normalize(cls, sql: str) -> str:
        """
        Normalize SQL from MySQL/PostgreSQL to SQLite syntax.
        
        Args:
            sql: SQL query string potentially containing MySQL/PostgreSQL syntax
            
        Returns:
            SQLite-compatible SQL query string
        """
        normalized = sql
        
        # Apply all function mappings
        for mysql_pattern, sqlite_replacement in cls.FUNCTION_MAPPINGS.items():
            normalized = re.sub(mysql_pattern, sqlite_replacement, normalized, flags=re.IGNORECASE)
        
        # Handle DATE_FORMAT -> STRFTIME conversion
        # This is complex so we do it separately
        normalized = cls._convert_date_format(normalized)
        
        # Handle DATE_ADD/DATE_SUB -> DATE modification
        normalized = cls._convert_date_arithmetic(normalized)
        
        logger.debug(f"SQL normalized. Original:\n{sql[:200]}...\nNormalized:\n{normalized[:200]}...")
        
        return normalized
    
    @classmethod
    def _convert_date_format(cls, sql: str) -> str:
        """
        Convert DATE_FORMAT(column, format) to strftime(format, column).
        
        Examples:
            DATE_FORMAT(sale_date, '%Y-%m') -> strftime('%Y-%m', sale_date)
            DATE_FORMAT(order_date, '%Y') -> strftime('%Y', order_date)
        """
        # Pattern: DATE_FORMAT(column, 'format_string')
        pattern = r"DATE_FORMAT\s*\(\s*([^,]+?)\s*,\s*['\"]([^'\"]+)['\"]\s*\)"
        
        def replace_date_format(match):
            column = match.group(1).strip()
            format_str = match.group(2)
            # Ensure format string uses double quotes for SQLite strftime
            return f"strftime('{format_str}', {column})"
        
        return re.sub(pattern, replace_date_format, sql, flags=re.IGNORECASE)
    
    @classmethod
    def _convert_date_arithmetic(cls, sql: str) -> str:
        """
        Convert DATE_ADD/DATE_SUB to SQLite date() function.
        
        Examples:
            DATE_ADD(date_col, INTERVAL 1 MONTH) -> date(date_col, '+1 month')
            DATE_SUB(date_col, INTERVAL 7 DAY) -> date(date_col, '-7 days')
        """
        # Pattern: DATE_ADD(column, INTERVAL n UNIT) or DATE_SUB(column, INTERVAL n UNIT)
        pattern = r"DATE_(ADD|SUB)\s*\(\s*([^,]+?)\s*,\s*INTERVAL\s+(\d+)\s+(\w+)\s*\)"
        
        def replace_date_arithmetic(match):
            operation = match.group(1).lower()  # 'add' or 'sub'
            column = match.group(2).strip()
            interval_value = match.group(3)
            interval_unit = match.group(4).lower()
            
            # Determine sign and unit
            sign = '+' if operation == 'add' else '-'
            
            # Map MySQL interval units to SQLite units
            unit_map = {
                'day': 'days',
                'month': 'months',
                'year': 'years',
                'hour': 'hours',
                'minute': 'minutes',
                'second': 'seconds',
                'week': 'days',  # Convert weeks to days
            }
            
            sqlite_unit = unit_map.get(interval_unit, interval_unit + 's')
            
            # If weeks, multiply by 7
            if interval_unit.lower() == 'week':
                interval_value = str(int(interval_value) * 7)
            
            return f"date({column}, '{sign}{interval_value} {sqlite_unit}')"
        
        return re.sub(pattern, replace_date_arithmetic, sql, flags=re.IGNORECASE)
    
    @classmethod
    def normalize_query_for_sqlite(cls, sql: str) -> str:
        """
        Full normalization for SQLite including data type conversions.
        
        Args:
            sql: SQL query string
            
        Returns:
            SQLite-compatible SQL query string
        """
        normalized = cls.normalize(sql)
        
        # Additional SQLite-specific adjustments
        # Convert LIMIT/OFFSET (usually same, but good practice)
        # Remove ENGINE and other MySQL-specific table options
        normalized = re.sub(r'\s+ENGINE\s*=\s*[A-Za-z0-9]+', '', normalized, flags=re.IGNORECASE)
        normalized = re.sub(r'\s+COLLATE\s+[A-Za-z0-9_]+', '', normalized, flags=re.IGNORECASE)
        
        return normalized


def normalize_sql(sql: str) -> str:
    """Convenience function to normalize SQL for SQLite."""
    return SQLNormalizer.normalize(sql)