dataface 0.1.2__py3-none-any.whl

dataface/core/compile/filter_injection.py

@@ -0,0 +1,319 @@
+"""Declarative SQL filter injection via sqlglot AST manipulation.
+
+Stage: COMPILE / EXECUTE bridge
+Purpose: Parse SQL with sqlglot, inject WHERE clauses from a declarative
+    filters spec, and return parameterized SQL.
+
+Users write pure SQL and declare variable-to-column bindings in YAML:
+
+    queries:
+      sales:
+        sql: SELECT region, SUM(revenue) FROM orders GROUP BY region
+        filters:
+          region: region              # equality (auto-IN for arrays)
+          revenue: { gte: min_revenue }  # comparison
+          order_date: { between: date_range }  # date range
+
+The engine parses the SQL with sqlglot, builds predicate AST nodes,
+and injects them into the WHERE clause. Null variables produce no
+predicate — no 1=1 hacks.
+
+Entry Point:
+    inject_filters(sql, filters, variables, dialect, param_offset) → (sql, params)
+
+Dependencies:
+    - sqlglot
+    - dataface.core.execute.dialects (SQLDialect for placeholder format)
+"""
+
+from typing import Any
+
+import sqlglot
+from sqlglot import exp
+
+from dataface.core.execute.dialects.base import SQLDialect
+
+# Map YAML operator keys to sqlglot binary expression types.
+_OP_MAP: dict[str, type[exp.Binary]] = {
+    "eq": exp.EQ,
+    "neq": exp.NEQ,
+    "gt": exp.GT,
+    "gte": exp.GTE,
+    "lt": exp.LT,
+    "lte": exp.LTE,
+    "like": exp.Like,
+    "ilike": exp.ILike,
+}
+
+# Map our dialect names to sqlglot dialect names.
+_SQLGLOT_DIALECT: dict[str, str] = {
+    "postgres": "postgres",
+    "postgresql": "postgres",
+    "duckdb": "duckdb",
+    "mysql": "mysql",
+    "mariadb": "mysql",
+    "snowflake": "snowflake",
+    "bigquery": "bigquery",
+    "redshift": "redshift",
+    "sqlserver": "tsql",
+    "mssql": "tsql",
+    "databricks": "databricks",
+    "spark": "spark",
+    "athena": "presto",
+    "presto": "presto",
+    "trino": "trino",
+}
+
+
+def inject_filters(
+    sql: str,
+    filters: dict[str, Any] | None,
+    variables: dict[str, Any],
+    dialect: SQLDialect,
+    param_offset: int = 0,
+) -> tuple[str, list[Any]]:
+    """Inject WHERE clauses into SQL from a declarative filters spec.
+
+    Parses the SQL with sqlglot, builds predicate nodes for each filter
+    where the bound variable has a non-null value, and injects them into
+    the appropriate WHERE clause (outermost SELECT or targeted CTE).
+
+    Args:
+        sql: SQL string (may already contain parameterized placeholders
+            from Jinja rendering).
+        filters: Declarative filter spec from YAML. Keys are column names
+            (or CTE names for CTE targeting). Values are variable names
+            (str) or operator dicts ({op: var_name}).
+        variables: Current variable values.
+        dialect: SQLDialect instance for placeholder format.
+        param_offset: Number of existing params (from prior Jinja rendering).
+            New placeholders start numbering after this offset.
+
+    Returns:
+        Tuple of (modified_sql, new_params). The new_params list contains
+        only the params added by filter injection — caller should append
+        these to any existing Jinja params.
+    """
+    if not filters:
+        return sql, []
+
+    sg_dialect = _SQLGLOT_DIALECT.get(dialect.name, dialect.name)
+    tree = sqlglot.parse_one(sql, read=sg_dialect)
+
+    # Identify CTE names so we can distinguish CTE targeting from column filters.
+    cte_names = {cte.alias: cte for cte in tree.find_all(exp.CTE) if cte.alias}
+
+    all_params: list[Any] = []
+    idx = param_offset
+
+    # Separate CTE-targeted filters from column-level filters.
+    # A filter key is CTE-targeted when: (a) its name matches a CTE in the SQL,
+    # (b) its value is a dict, and (c) that dict's keys are NOT recognized
+    # operators (i.e., it contains column→binding pairs, not an operator→var pair).
+    column_filters: dict[str, str | dict[str, str]] = {}
+    cte_filters: dict[str, dict[str, str | dict[str, str]]] = {}
+
+    for key, binding in filters.items():
+        if (
+            key in cte_names
+            and isinstance(binding, dict)
+            and not _is_operator_dict(binding)
+        ):
+            # Validate each value in the CTE binding is a valid filter binding
+            for col, col_binding in binding.items():
+                if not isinstance(col_binding, (str, dict)):
+                    raise ValueError(
+                        f"Invalid filter binding for CTE '{key}', column '{col}': "
+                        f"expected string or operator dict, got {type(col_binding).__name__}"
+                    )
+            cte_filters[key] = binding
+        else:
+            column_filters[key] = binding
+
+    # Inject column-level filters against the query result shape rather than the
+    # innermost SELECT body. This preserves support for filters that target
+    # projected aliases (common in generated SQL) because aliases are not visible
+    # inside the same SELECT's WHERE clause in BigQuery and other dialects.
+    column_conditions, params, idx = _build_conditions(
+        column_filters, variables, dialect, idx
+    )
+    all_params.extend(params)
+
+    # Inject CTE-targeted filters.
+    for cte_name, cte_bindings in cte_filters.items():
+        cte_node = cte_names[cte_name]
+        cte_select = cte_node.find(exp.Select)
+        cte_conditions, params, idx = _build_conditions(
+            cte_bindings, variables, dialect, idx
+        )
+        all_params.extend(params)
+        if cte_conditions:
+            if cte_select is None:
+                raise ValueError(f"Could not find SELECT node for CTE '{cte_name}'")
+            _inject_where(cte_select, cte_conditions)
+
+    if column_conditions:
+        alias = exp.TableAlias(this=exp.to_identifier("_df_filter_base"))
+        wrapped = exp.select("*").from_(exp.Subquery(this=tree, alias=alias))
+        _inject_where(wrapped, column_conditions)
+        return wrapped.sql(dialect=sg_dialect), all_params
+
+    return tree.sql(dialect=sg_dialect), all_params
+
+
+def _is_operator_dict(d: dict) -> bool:
+    """Check if a dict's keys are all recognized filter operators."""
+    operators = set(_OP_MAP) | {"in", "not_in", "between"}
+    return all(k in operators for k in d)
+
+
+def _build_conditions(
+    filters: dict[str, Any],
+    variables: dict[str, Any],
+    dialect: SQLDialect,
+    start_idx: int,
+) -> tuple[list[exp.Expression], list[Any], int]:
+    """Build sqlglot condition expressions from filter bindings.
+
+    Returns (conditions, params, next_idx).
+    """
+    conditions: list[exp.Expression] = []
+    params: list[Any] = []
+    idx = start_idx
+
+    for col_name, binding in filters.items():
+        op_key, var_name = _parse_binding(binding)
+        value = variables.get(var_name)
+
+        # Skip null / missing / empty string — the whole point of declarative filters.
+        if value is None or value == "":
+            continue
+
+        col = exp.column(col_name)
+        cond, cond_params, idx = _build_condition(col, op_key, value, dialect, idx)
+        if cond is not None:
+            conditions.append(cond)
+            params.extend(cond_params)
+
+    return conditions, params, idx
+
+
+def _parse_binding(binding: Any) -> tuple[str, str]:
+    """Parse a filter binding into (operator_key, variable_name).
+
+    Simple binding:  "region"        → ("eq", "region")
+    Operator binding: {"gte": "min"} → ("gte", "min")
+    """
+    if isinstance(binding, str):
+        return "eq", binding
+    if isinstance(binding, dict):
+        op_key, var_name = next(iter(binding.items()))
+        return op_key, var_name
+    raise ValueError(f"Invalid filter binding: {binding!r}")
+
+
+def _build_condition(
+    col: exp.Column,
+    op_key: str,
+    value: Any,
+    dialect: SQLDialect,
+    start_idx: int,
+) -> tuple[exp.Expression | None, list[Any], int]:
+    """Build a single sqlglot predicate + params.
+
+    Returns (condition, params, next_param_idx).
+    """
+    # BETWEEN
+    if op_key == "between":
+        if not isinstance(value, (list, tuple)) or len(value) != 2:
+            return None, [], start_idx
+        start_val, end_val = value
+        if start_val is None or end_val is None:
+            return None, [], start_idx
+        p1 = _param_node(dialect, start_idx + 1)
+        p2 = _param_node(dialect, start_idx + 2)
+        return (
+            exp.Between(this=col, low=p1, high=p2),
+            [start_val, end_val],
+            start_idx + 2,
+        )
+
+    # IN / NOT IN (explicit operator or implicit from list value)
+    if op_key in ("in", "not_in") or (op_key == "eq" and isinstance(value, list)):
+        if not value:
+            return None, [], start_idx
+        placeholders = []
+        for _v in value:
+            start_idx += 1
+            placeholders.append(_param_node(dialect, start_idx))
+        in_node = exp.In(this=col, expressions=placeholders)
+        if op_key == "not_in":
+            return exp.Not(this=in_node), list(value), start_idx
+        return in_node, list(value), start_idx
+
+    # Standard comparison (eq, neq, gt, gte, lt, lte, like, ilike)
+    cls = _OP_MAP.get(op_key, exp.EQ)
+    start_idx += 1
+    return cls(this=col, expression=_param_node(dialect, start_idx)), [value], start_idx
+
+
+def _param_node(dialect: SQLDialect, idx: int) -> exp.Expression:
+    """Create a sqlglot node that renders as the dialect's placeholder.
+
+    Uses exp.Literal(is_string=False) so the placeholder text is emitted
+    verbatim without quotes.
+    """
+    return exp.Literal(this=dialect.param(idx), is_string=False)
+
+
+def _inject_where(
+    select: exp.Select,
+    conditions: list[exp.Expression],
+) -> None:
+    """Inject conditions into a SELECT's WHERE clause.
+
+    If the SELECT already has a WHERE, the new conditions are ANDed.
+    If not, a new WHERE is created.
+    """
+    combined = conditions[0]
+    for c in conditions[1:]:
+        combined = exp.And(this=combined, expression=c)
+
+    existing_where = select.args.get("where")
+    if existing_where:
+        select.set(
+            "where",
+            exp.Where(this=exp.And(this=existing_where.this, expression=combined)),
+        )
+    else:
+        select.set("where", exp.Where(this=combined))
+
+
+def extract_filter_variable_deps(filters: dict[str, Any]) -> set[str]:
+    """Extract variable names referenced in a declarative filters spec.
+
+    Used during normalization to compute variable_dependencies for a query.
+    Handles both flat column filters and CTE-targeted nested filters.
+
+    Args:
+        filters: Declarative filter spec (same format as inject_filters).
+
+    Returns:
+        Set of variable names referenced.
+    """
+    if not filters:
+        return set()
+
+    deps: set[str] = set()
+    for binding in filters.values():
+        _collect_binding_vars(binding, deps)
+    return deps
+
+
+def _collect_binding_vars(binding: Any, deps: set[str]) -> None:
+    """Recursively collect variable names from a filter binding."""
+    if isinstance(binding, str):
+        deps.add(binding)
+    elif isinstance(binding, dict):
+        for v in binding.values():
+            _collect_binding_vars(v, deps)