executor.py

"""SQL Executor - executes translated queries against Redis."""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

import redis

from sql_redis.schema import SchemaRegistry
from sql_redis.translator import Translator


@dataclass
class QueryResult:
    """Result of executing a SQL query."""

    rows: list[dict]
    count: int


class Executor:
    """Executes SQL queries against Redis."""

    def __init__(self, client: redis.Redis, schema_registry: SchemaRegistry):
        """Initialize executor with Redis client and schema registry."""
        self._client = client
        self._schema_registry = schema_registry
        self._translator = Translator(schema_registry)

    def _substitute_params(self, sql: str, params: dict[str, Any]) -> str:
        """Substitute parameter placeholders in SQL with actual values.

        Uses token-based approach: splits SQL on :param patterns, then rebuilds
        with substituted values. This approach solves two critical bugs:

        1. PARTIAL MATCHING BUG: Prevents :id from matching inside :product_id
           by treating each :identifier as a complete token

        2. QUOTE ESCAPING BUG: Properly escapes single quotes in string values
           using SQL standard (single quote -> double single quote)

        Args:
            sql: The SQL string with :param placeholders.
            params: Dictionary mapping parameter names to values.

        Returns:
            SQL string with parameters substituted.

        Implementation Details:
            - Uses regex to split on parameter patterns: :[a-zA-Z_][a-zA-Z0-9_]*
            - Keeps delimiters (the :param tokens) in the split result
            - Iterates through tokens, substituting matched parameters
            - String values are wrapped in single quotes with proper escaping
            - Numeric values are converted to strings
            - Bytes values (e.g., vectors) are NOT substituted here

        Known Limitations:
            - Colons in string literals: SQL like "WHERE x = 'test:value'" would
              theoretically match :value as a parameter. However, this is not a
              practical issue because:
              1. Users pass values via parameters, not hardcoded in SQL
              2. The translator has its own handling of string literals
              3. No real-world use cases have been identified
            - Parameter names are case-sensitive (:id != :ID)
            - Only handles int, float, str types; other types keep placeholder
        """
        if not params:
            return sql

        # Split SQL on :param patterns, keeping the delimiters
        # Pattern matches : followed by valid identifier:
        #   [a-zA-Z_]       - First char must be letter or underscore
        #   [a-zA-Z0-9_]*   - Subsequent chars can be alphanumeric or underscore
        # This prevents partial matching: :id and :product_id are separate tokens
        tokens = re.split(r"(:[a-zA-Z_][a-zA-Z0-9_]*)", sql)

        result = []
        for token in tokens:
            if token.startswith(":"):
                # This is a parameter placeholder
                key = token[1:]  # Remove leading :
                if key in params:
                    value = params[key]
                    if isinstance(value, (int, float)):
                        # Numeric values: convert to string
                        result.append(str(value))
                    elif isinstance(value, str):
                        # String values: wrap in quotes and escape single quotes
                        # SQL standard: ' -> '' (double single quote)
                        # This fixes the quote escaping bug
                        escaped = value.replace("'", "''")
                        result.append(f"'{escaped}'")
                    else:
                        # Other types (bytes, None, bool, list, etc.):
                        # Keep placeholder as-is (handled elsewhere or unsupported)
                        result.append(token)
                else:
                    # Parameter not provided: keep placeholder as-is
                    result.append(token)
            else:
                # Not a parameter: keep as-is
                result.append(token)

        return "".join(result)

    def execute(self, sql: str, *, params: dict | None = None) -> QueryResult:
        """Execute a SQL query and return results."""
        params = params or {}

        # Substitute non-bytes params in SQL using token-based approach
        sql = self._substitute_params(sql, params)

        # Translate SQL to Redis command
        translated = self._translator.translate(sql)

        # Build command list and substitute vector params
        # Use list[str | bytes] to allow bytes for vector params
        cmd: list[str | bytes] = list(translated.to_command_list())

        # Find any bytes params (vectors) to substitute
        vector_param: bytes | None = None
        for value in params.values():
            if isinstance(value, bytes):
                vector_param = value
                break

        # Replace $vector placeholder with actual bytes
        if vector_param:
            for i, arg in enumerate(cmd):
                if arg == "$vector":
                    cmd[i] = vector_param

        # Execute command
        raw_result = self._client.execute_command(*cmd)

        # Parse result based on command type
        count = raw_result[0] if raw_result else 0
        rows = []

        if translated.command == "FT.SEARCH":
            # FT.SEARCH format: [count, key1, [fields1], key2, [fields2], ...]
            # Skip document keys (odd indices), take field lists (even indices after count)
            for i in range(2, len(raw_result), 2):
                row_data = raw_result[i]
                row = dict(zip(row_data[::2], row_data[1::2]))
                rows.append(row)
        else:
            # FT.AGGREGATE format: [count, [fields1], [fields2], ...]
            for row_data in raw_result[1:]:
                row = dict(zip(row_data[::2], row_data[1::2]))
                rows.append(row)

        return QueryResult(rows=rows, count=count)