sql-redis 0.1.0__py3-none-any.whl

sql_redis/__init__.py

@@ -0,0 +1,5 @@
+"""SQL to Redis command translation utility."""
+
+from sql_redis.translator import TranslatedQuery, Translator
+
+__all__ = ["Translator", "TranslatedQuery"]

sql_redis/analyzer.py

@@ -0,0 +1,133 @@
+"""SQL analyzer component - resolves field types from schema."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from sql_redis.parser import AggregationSpec, ComputedField, Condition, ParsedQuery
+
+
+@dataclass
+class VectorSearchAnalysis:
+    """Analyzed vector search details."""
+
+    field: str
+    k: int
+    alias: str
+
+
+@dataclass
+class AnalyzedQuery:
+    """Result of analyzing a parsed SQL query with schema context."""
+
+    parsed: ParsedQuery = field(default_factory=ParsedQuery)
+    field_types: dict[str, str] = field(default_factory=dict)
+    aggregations: list[AggregationSpec] = field(default_factory=list)
+    computed_fields: list[ComputedField] = field(default_factory=list)
+    groupby_fields: list[str] = field(default_factory=list)
+    is_global_aggregation: bool = False
+    vector_search: VectorSearchAnalysis | None = None
+    has_prefilter: bool = False
+
+    def get_field_type(self, field_name: str) -> str | None:
+        """Get the type of a field."""
+        return self.field_types.get(field_name)
+
+    def get_conditions_by_type(self, field_type: str) -> list[Condition]:
+        """Get conditions for fields of a specific type."""
+        return [
+            c
+            for c in self.parsed.conditions
+            if self.field_types.get(c.field) == field_type
+        ]
+
+
+class Analyzer:
+    """Analyzes parsed SQL queries with schema context."""
+
+    def __init__(self, schemas: dict[str, dict[str, str]]):
+        """Initialize analyzer with schema registry data.
+
+        Args:
+            schemas: Dict mapping index names to field->type dicts.
+        """
+        self._schemas = schemas
+
+    def analyze(self, parsed: ParsedQuery) -> AnalyzedQuery:
+        """Analyze a parsed query, resolving field types.
+
+        Args:
+            parsed: The parsed SQL query.
+
+        Returns:
+            An AnalyzedQuery with field types resolved.
+
+        Raises:
+            ValueError: If the index or a field is unknown.
+        """
+        # Validate index exists
+        if parsed.index not in self._schemas:
+            raise ValueError(f"Unknown index: {parsed.index}")
+
+        schema = self._schemas[parsed.index]
+        result = AnalyzedQuery(parsed=parsed)
+
+        # Collect all fields referenced in the query
+        referenced_fields: set[str] = set()
+
+        # Fields from SELECT
+        for field_name in parsed.fields:
+            if field_name != "*":
+                referenced_fields.add(field_name)
+
+        # Fields from conditions
+        for condition in parsed.conditions:
+            referenced_fields.add(condition.field)
+
+        # Fields from aggregations
+        for agg in parsed.aggregations:
+            if agg.field:
+                referenced_fields.add(agg.field)
+
+        # Fields from computed fields (extract field references from expressions)
+        for computed in parsed.computed_fields:
+            # Simple extraction - look for field names in the expression
+            for field_name in schema.keys():
+                if field_name in computed.expression:
+                    referenced_fields.add(field_name)
+
+        # Fields from vector search
+        if parsed.vector_search:
+            referenced_fields.add(parsed.vector_search.field)
+
+        # Fields from GROUP BY
+        for field_name in parsed.groupby_fields:
+            referenced_fields.add(field_name)
+
+        # Resolve field types
+        for field_name in referenced_fields:
+            if field_name not in schema:
+                raise ValueError(f"Unknown field: {field_name}")
+            result.field_types[field_name] = schema[field_name]
+
+        # Copy aggregations and computed fields
+        result.aggregations = parsed.aggregations
+        result.computed_fields = parsed.computed_fields
+        result.groupby_fields = parsed.groupby_fields
+
+        # Determine if this is a global aggregation
+        result.is_global_aggregation = (
+            len(parsed.aggregations) > 0 and len(parsed.groupby_fields) == 0
+        )
+
+        # Analyze vector search
+        if parsed.vector_search:
+            result.vector_search = VectorSearchAnalysis(
+                field=parsed.vector_search.field,
+                k=parsed.limit or parsed.vector_search.k or 10,
+                alias=parsed.vector_search.alias,
+            )
+            # Has prefilter if there are conditions
+            result.has_prefilter = len(parsed.conditions) > 0
+
+        return result

sql_redis/executor.py

@@ -0,0 +1,83 @@
+"""SQL Executor - executes translated queries against Redis."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+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 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
+        for key, value in params.items():
+            placeholder = f":{key}"
+            if isinstance(value, (int, float)):
+                sql = sql.replace(placeholder, str(value))
+            elif isinstance(value, str):
+                sql = sql.replace(placeholder, f"'{value}'")
+            # bytes (vectors) are handled via Redis 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)

sql_redis/parser.py

@@ -0,0 +1,440 @@
+"""SQL parser component using sqlglot."""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+
+import sqlglot
+from sqlglot import exp
+
+
+@dataclass
+class AggregationSpec:
+    """Specification for an aggregation function."""
+
+    function: str
+    field: str | None = None
+    alias: str | None = None
+    extra_args: list[str] = dataclasses.field(
+        default_factory=list
+    )  # For reducers like QUANTILE
+
+
+@dataclass
+class ComputedField:
+    """Specification for a computed/APPLY field."""
+
+    expression: str
+    alias: str
+
+
+@dataclass
+class VectorSearchSpec:
+    """Specification for vector search."""
+
+    field: str
+    alias: str
+    k: int | None = None
+
+
+@dataclass
+class Condition:
+    """A WHERE condition."""
+
+    field: str
+    operator: str
+    value: object
+    negated: bool = False
+
+
+@dataclass
+class ParsedQuery:
+    """Result of parsing a SQL query."""
+
+    index: str = ""
+    fields: list[str] = dataclasses.field(default_factory=list)
+    conditions: list[Condition] = dataclasses.field(default_factory=list)
+    boolean_operator: str = "AND"
+    aggregations: list[AggregationSpec] = dataclasses.field(default_factory=list)
+    computed_fields: list[ComputedField] = dataclasses.field(default_factory=list)
+    vector_search: VectorSearchSpec | None = None
+    groupby_fields: list[str] = dataclasses.field(default_factory=list)
+    orderby_fields: list[tuple[str, str]] = dataclasses.field(
+        default_factory=list
+    )  # (field, ASC|DESC)
+    limit: int | None = None
+    offset: int | None = None
+
+
+class SQLParser:
+    """Parses SQL into a ParsedQuery structure."""
+
+    def parse(self, sql: str) -> ParsedQuery:
+        """Parse a SQL statement into a ParsedQuery.
+
+        Args:
+            sql: The SQL statement to parse.
+
+        Returns:
+            A ParsedQuery containing the extracted components.
+        """
+        ast = sqlglot.parse_one(sql)
+        result = ParsedQuery()
+
+        # Extract FROM clause (index name)
+        from_clause = ast.find(exp.From)
+        if from_clause:
+            table = from_clause.find(exp.Table)
+            if table:
+                result.index = table.name
+
+        # Extract SELECT fields and aggregations
+        select = ast.find(exp.Select)
+        if select:
+            for expression in select.expressions:
+                self._process_select_expression(expression, result)
+
+        # Extract WHERE clause conditions
+        where = ast.find(exp.Where)
+        if where:
+            self._process_where_clause(where.this, result)
+
+        # Extract GROUP BY clause
+        group = ast.find(exp.Group)
+        if group:
+            for expr in group.expressions:
+                if isinstance(expr, exp.Column):
+                    result.groupby_fields.append(expr.name)
+
+        # Extract ORDER BY clause
+        order = ast.find(exp.Order)
+        if order:
+            for ordered in order.expressions:
+                col = ordered.this
+                if isinstance(col, exp.Column):
+                    direction = "DESC" if ordered.args.get("desc") else "ASC"
+                    result.orderby_fields.append((col.name, direction))
+                elif isinstance(col, (exp.CosineDistance, exp.Distance)):
+                    # ORDER BY vector distance - handled by KNN, don't add to orderby
+                    # The vector_search should already be set from SELECT clause
+                    pass
+
+        # Extract LIMIT clause
+        limit = ast.find(exp.Limit)
+        if limit:
+            limit_expr = limit.args.get("expression") or limit.this
+            if isinstance(limit_expr, exp.Literal):
+                result.limit = int(limit_expr.this)
+
+        # Extract OFFSET clause
+        offset = ast.find(exp.Offset)
+        if offset:
+            offset_expr = offset.args.get("expression") or offset.this
+            if isinstance(offset_expr, exp.Literal):
+                result.offset = int(offset_expr.this)
+
+        return result
+
+    def _process_select_expression(self, expression, result: ParsedQuery) -> None:
+        """Process a single SELECT expression."""
+        # Handle aliased expressions (e.g., COUNT(*) AS count)
+        if isinstance(expression, exp.Alias):
+            alias = expression.alias
+            inner = expression.this
+            self._process_select_expression_inner(inner, result, alias)
+        else:
+            self._process_select_expression_inner(expression, result, None)
+
+    def _process_select_expression_inner(
+        self, expression, result: ParsedQuery, alias: str | None
+    ) -> None:
+        """Process the inner part of a SELECT expression."""
+        if isinstance(expression, exp.Column):
+            result.fields.append(expression.name)
+        elif isinstance(expression, exp.Star):
+            result.fields.append("*")
+        elif isinstance(
+            expression,
+            (
+                exp.Count,
+                exp.Sum,
+                exp.Avg,
+                exp.Min,
+                exp.Max,
+                exp.Stddev,
+                exp.Variance,
+                exp.FirstValue,
+                exp.ArrayAgg,
+            ),
+        ):
+            # Aggregation function
+            # Map sqlglot function names to Redis reducer names
+            func_name = expression.key.upper()
+            redis_func_map = {
+                "FIRSTVALUE": "FIRST_VALUE",
+                "ARRAYAGG": "TOLIST",
+            }
+            func_name = redis_func_map.get(func_name, func_name)
+            field_name = None
+            # Get the field being aggregated (if any)
+            if expression.this:
+                if isinstance(expression.this, exp.Column):
+                    field_name = expression.this.name
+                elif isinstance(expression.this, exp.Star):
+                    field_name = None  # COUNT(*)
+            result.aggregations.append(
+                AggregationSpec(function=func_name, field=field_name, alias=alias)
+            )
+        elif isinstance(expression, exp.Paren):
+            # Parenthesized expression - computed field
+            inner_expr = expression.this.sql()
+            # Use alias if provided, otherwise generate one from expression
+            field_alias = alias if alias else inner_expr
+            result.computed_fields.append(
+                ComputedField(expression=inner_expr, alias=field_alias)
+            )
+        elif isinstance(expression, (exp.Mul, exp.Div, exp.Add, exp.Sub)):
+            # Arithmetic expression without parentheses - computed field
+            expr_str = expression.sql()
+            # Use alias if provided, otherwise generate one from expression
+            field_alias = alias if alias else expr_str
+            result.computed_fields.append(
+                ComputedField(expression=expr_str, alias=field_alias)
+            )
+        elif isinstance(expression, (exp.Distance, exp.CosineDistance)):
+            # Vector distance functions:
+            # - Distance: L2/Euclidean distance
+            # - CosineDistance: cosine_distance() function
+            self._process_vector_distance(expression, result, alias)
+        elif isinstance(expression, exp.Quantile):
+            # QUANTILE(field, quantile_value) -> REDUCE QUANTILE 2 @field quantile_value
+            field_name = None
+            if expression.this and isinstance(expression.this, exp.Column):
+                field_name = expression.this.name
+            quantile_value = None
+            if expression.args.get("quantile"):
+                quantile_value = str(expression.args["quantile"].this)
+            extra_args = [quantile_value] if quantile_value else []
+            result.aggregations.append(
+                AggregationSpec(
+                    function="QUANTILE",
+                    field=field_name,
+                    alias=alias,
+                    extra_args=extra_args,
+                )
+            )
+        elif isinstance(expression, exp.Anonymous):
+            # Custom function call (e.g., vector_distance) - check before exp.Func
+            # since Anonymous is a subclass of Func
+            func_name = expression.name.lower()
+            # Redis-specific reducer functions that sqlglot doesn't recognize
+            redis_reducers = {
+                "count_distinct",
+                "count_distinctish",
+                "quantile",
+                "random_sample",
+            }
+            if func_name == "vector_distance":
+                # Extract the vector field name from first argument
+                if expression.expressions:
+                    first_arg = expression.expressions[0]
+                    if isinstance(first_arg, exp.Column):
+                        field_name = first_arg.name
+                        result.vector_search = VectorSearchSpec(
+                            field=field_name,
+                            alias=alias or func_name,
+                        )
+            elif func_name in redis_reducers:
+                # Redis-specific reducer functions
+                field_name = None
+                reducer_extra_args: list[str] = []
+                if expression.expressions:
+                    first_arg = expression.expressions[0]
+                    if isinstance(first_arg, exp.Column):
+                        field_name = first_arg.name
+                    # Extract additional arguments (e.g., quantile value for QUANTILE)
+                    for arg in expression.expressions[1:]:
+                        if isinstance(arg, exp.Literal):
+                            reducer_extra_args.append(str(arg.this))
+                result.aggregations.append(
+                    AggregationSpec(
+                        function=func_name.upper(),
+                        field=field_name,
+                        alias=alias,
+                        extra_args=reducer_extra_args,
+                    )
+                )
+            else:
+                # Other custom functions - treat as computed field
+                expr_str = expression.sql()
+                field_alias = alias if alias else expr_str
+                result.computed_fields.append(
+                    ComputedField(expression=expr_str, alias=field_alias)
+                )
+        elif isinstance(expression, exp.Func):
+            # Built-in function call (e.g., UPPER, LOWER, etc.) - treat as computed field
+            expr_str = expression.sql()
+            field_alias = alias if alias else expr_str
+            result.computed_fields.append(
+                ComputedField(expression=expr_str, alias=field_alias)
+            )
+
+    def _process_vector_distance(
+        self, expression, result: ParsedQuery, alias: str | None
+    ) -> None:
+        """Process a vector distance expression (cosine_distance, etc.)."""
+        field_name = None
+
+        # Extract field from the expression
+        # Both Distance and CosineDistance have 'this' as the first argument
+        if expression.this and isinstance(expression.this, exp.Column):
+            field_name = expression.this.name
+
+        if field_name:
+            result.vector_search = VectorSearchSpec(
+                field=field_name,
+                alias=alias or "vector_distance",
+            )
+
+    def _process_where_clause(
+        self, expression, result: ParsedQuery, negated: bool = False
+    ) -> None:
+        """Process WHERE clause expression recursively."""
+        if isinstance(expression, exp.EQ):
+            self._add_condition(expression, "=", result, negated)
+        elif isinstance(expression, exp.GT):
+            self._add_condition(expression, ">", result, negated)
+        elif isinstance(expression, exp.GTE):
+            self._add_condition(expression, ">=", result, negated)
+        elif isinstance(expression, exp.LT):
+            self._add_condition(expression, "<", result, negated)
+        elif isinstance(expression, exp.LTE):
+            self._add_condition(expression, "<=", result, negated)
+        elif isinstance(expression, exp.NEQ):
+            self._add_condition(expression, "!=", result, negated)
+        elif isinstance(expression, exp.Between):
+            self._add_between_condition(expression, result, negated)
+        elif isinstance(expression, exp.In):
+            self._add_in_condition(expression, result, negated)
+        elif isinstance(expression, exp.And):
+            result.boolean_operator = "AND"
+            self._process_where_clause(expression.this, result, negated)
+            self._process_where_clause(expression.expression, result, negated)
+        elif isinstance(expression, exp.Or):
+            result.boolean_operator = "OR"
+            self._process_where_clause(expression.this, result, negated)
+            self._process_where_clause(expression.expression, result, negated)
+        elif isinstance(expression, exp.Not):
+            self._process_where_clause(expression.this, result, negated=True)
+        elif isinstance(expression, exp.Anonymous):
+            # Custom function like MATCH(field, value)
+            self._add_function_condition(expression, result, negated)
+
+    def _add_condition(
+        self, expression, operator: str, result: ParsedQuery, negated: bool
+    ) -> None:
+        """Add a condition from a comparison expression."""
+        field_name = None
+        value = None
+
+        # Get field name from left side
+        if isinstance(expression.this, exp.Column):
+            field_name = expression.this.name
+        elif isinstance(expression.this, exp.Anonymous):
+            # Function call like DISTANCE(location, POINT(...))
+            # Extract field from first argument
+            func_name = expression.this.name.upper()
+            if expression.this.expressions:
+                first_arg = expression.this.expressions[0]
+                if isinstance(first_arg, exp.Column):
+                    field_name = first_arg.name
+                    # Use function name as operator prefix
+                    operator = f"{func_name}_{operator}"
+
+        # Get value from right side
+        if isinstance(expression.expression, exp.Literal):
+            value = expression.expression.this
+            # Convert numeric strings to numbers
+            if expression.expression.is_number:
+                value = int(value) if "." not in str(value) else float(value)
+
+        if field_name is not None:
+            result.conditions.append(
+                Condition(
+                    field=field_name, operator=operator, value=value, negated=negated
+                )
+            )
+
+    def _add_between_condition(
+        self, expression, result: ParsedQuery, negated: bool
+    ) -> None:
+        """Add a BETWEEN condition."""
+        field_name = None
+        if isinstance(expression.this, exp.Column):
+            field_name = expression.this.name
+
+        low = expression.args.get("low")
+        high = expression.args.get("high")
+
+        low_val = self._extract_literal_value(low)
+        high_val = self._extract_literal_value(high)
+
+        if field_name is not None:
+            result.conditions.append(
+                Condition(
+                    field=field_name,
+                    operator="BETWEEN",
+                    value=(low_val, high_val),
+                    negated=negated,
+                )
+            )
+
+    def _add_in_condition(self, expression, result: ParsedQuery, negated: bool) -> None:
+        """Add an IN condition."""
+        field_name = None
+        if isinstance(expression.this, exp.Column):
+            field_name = expression.this.name
+
+        values = [self._extract_literal_value(e) for e in expression.expressions]
+
+        if field_name is not None:
+            result.conditions.append(
+                Condition(
+                    field=field_name, operator="IN", value=values, negated=negated
+                )
+            )
+
+    def _add_function_condition(
+        self, expression, result: ParsedQuery, negated: bool
+    ) -> None:
+        """Add a condition from a function call like fulltext(field, value)."""
+        func_name = expression.name.upper()
+        if func_name == "FULLTEXT" and len(expression.expressions) >= 2:
+            first_arg = expression.expressions[0]
+            second_arg = expression.expressions[1]
+
+            field_name = None
+            if isinstance(first_arg, exp.Column):
+                field_name = first_arg.name
+
+            value = self._extract_literal_value(second_arg)
+
+            if field_name is not None:
+                result.conditions.append(
+                    Condition(
+                        field=field_name,
+                        operator="FULLTEXT",
+                        value=value,
+                        negated=negated,
+                    )
+                )
+
+    def _extract_literal_value(self, expression):
+        """Extract a Python value from a sqlglot Literal."""
+        if isinstance(expression, exp.Literal):
+            value = expression.this
+            if expression.is_number:
+                return int(value) if "." not in str(value) else float(value)
+            return value
+        return None