dataface 0.1.2__py3-none-any.whl

dataface/core/execute/batch.py

@@ -0,0 +1,744 @@
+"""Batch query execution with temp table optimization.
+
+Stage: EXECUTE
+Purpose: Generate batch SQL with temp tables for queries that are referenced
+by multiple other queries via {{ queries.X }} syntax.
+
+Entry Points:
+    - extract_query_refs(sql: str) -> Set[str]
+    - build_dependency_graph(queries: Dict[str, Query]) -> DependencyGraph
+    - generate_batch_sql(queries, graph, dialect) -> List[BatchStatement]
+    - collect_face_queries(face) -> Dict[str, AnyQuery]
+    - group_by_source(queries) -> Dict[str, Dict[str, AnyQuery]]
+
+When multiple queries reference the same base query:
+    - The base query is executed ONCE and stored in a temp table
+    - Dependent queries SELECT from the temp table
+
+Example:
+    queries:
+      orders:
+        sql: SELECT * FROM orders WHERE status = 'completed'
+      high_value:
+        sql: SELECT * FROM {{ queries.orders }} WHERE amount > 1000
+      recent:
+        sql: SELECT * FROM {{ queries.orders }} WHERE date >= CURRENT_DATE - 7
+
+    Generated batch:
+      1. CREATE TEMP TABLE _df_temp_orders AS SELECT * FROM orders WHERE status = 'completed'
+      2. SELECT * FROM _df_temp_orders WHERE amount > 1000  (for high_value)
+      3. SELECT * FROM _df_temp_orders WHERE date >= CURRENT_DATE - 7  (for recent)
+
+Dependencies:
+    - dataface.execute.dialects (for temp table syntax)
+    - dataface.core.compile.models.query.compiled (AnyQuery, is_sql_query)
+
+Errors:
+    - CrossSourceReferenceError: When {{ queries.X }} references a query
+      on a different database source
+"""
+
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+from dataface.core.compile.models.face.compiled import Face
+from dataface.core.compile.models.query.compiled import (
+    AnyQuery,
+    is_sql_query,
+)
+from dataface.core.execute.dialects import SQLDialect, get_dialect
+
+
+class CrossSourceReferenceError(Exception):
+    """Error raised when a query references another query on a different source.
+
+    Cross-database references are not supported because temp tables are
+    connection-local and cannot be shared between different database connections.
+
+    Example:
+        queries:
+          orders:
+            source: postgres
+          analysis:
+            sql: SELECT * FROM {{ queries.orders }}
+            source: mysql  # ERROR: Can't reference postgres query from mysql
+    """
+
+    def __init__(
+        self,
+        source_query: str,
+        query_source: str,
+        target_query: str,
+        target_source: str,
+    ):
+        """Initialize CrossSourceReferenceError.
+
+        Args:
+            source_query: Name of the query making the reference
+            query_source: Source of the query making the reference
+            target_query: Name of the referenced query
+            target_source: Source of the referenced query
+        """
+        self.source_query = source_query
+        self.query_source = query_source
+        self.target_query = target_query
+        self.target_source = target_source
+        super().__init__(
+            f"Query '{source_query}' (source: {query_source}) cannot reference "
+            f"query '{target_query}' (source: {target_source}). "
+            f"Cross-database references are not supported."
+        )
+
+
+@dataclass
+class DependencyGraph:
+    """Directed graph representing query dependencies.
+
+    Nodes are query names. Edges point from a query to queries that depend on it.
+    For example, if query B references {{ queries.A }}, there's an edge A -> B.
+
+    Attributes:
+        edges: Dict mapping query name to set of queries that depend on it
+        reverse_edges: Dict mapping query name to set of queries it depends on
+        all_nodes: Set of all query names in the graph
+    """
+
+    edges: dict[str, set[str]] = field(default_factory=dict)
+    reverse_edges: dict[str, set[str]] = field(default_factory=dict)
+    all_nodes: set[str] = field(default_factory=set)
+
+    def add_node(self, name: str) -> None:
+        """Add a node to the graph.
+
+        Args:
+            name: Query name to add
+        """
+        self.all_nodes.add(name)
+        if name not in self.edges:
+            self.edges[name] = set()
+        if name not in self.reverse_edges:
+            self.reverse_edges[name] = set()
+
+    def add_edge(self, from_node: str, to_node: str) -> None:
+        """Add a directed edge from from_node to to_node.
+
+        Indicates that to_node depends on from_node.
+
+        Args:
+            from_node: The referenced query (dependency)
+            to_node: The query that references from_node (dependent)
+        """
+        self.add_node(from_node)
+        self.add_node(to_node)
+        self.edges[from_node].add(to_node)
+        self.reverse_edges[to_node].add(from_node)
+
+    def get_dependents(self, node: str) -> set[str]:
+        """Get all queries that depend on this node.
+
+        Args:
+            node: Query name
+
+        Returns:
+            Set of query names that reference this query
+        """
+        return self.edges.get(node, set())
+
+    def get_dependencies(self, node: str) -> set[str]:
+        """Get all queries this node depends on.
+
+        Args:
+            node: Query name
+
+        Returns:
+            Set of query names this query references
+        """
+        return self.reverse_edges.get(node, set())
+
+    def out_degree(self, node: str) -> int:
+        """Get the number of queries that depend on this node.
+
+        Args:
+            node: Query name
+
+        Returns:
+            Number of dependent queries (out-edges)
+        """
+        return len(self.get_dependents(node))
+
+    def get_temp_table_candidates(self) -> set[str]:
+        """Get queries that should be materialized as temp tables.
+
+        A query is a temp table candidate if it is referenced by at least
+        one other query (out_degree > 0).
+
+        Returns:
+            Set of query names that should become temp tables
+        """
+        return {node for node in self.all_nodes if self.out_degree(node) > 0}
+
+
+@dataclass
+class BatchStatement:
+    """A single statement in a batch execution plan.
+
+    Attributes:
+        name: Identifier for this statement (query name or _create_X for temp tables)
+        sql: The SQL to execute
+        is_temp_table_create: True if this creates a temp table
+        query_name: The logical query name this statement is for (may differ from name)
+    """
+
+    name: str
+    sql: str
+    is_temp_table_create: bool = False
+    query_name: str | None = None
+
+    def __post_init__(self) -> None:
+        """Set query_name to name if not provided."""
+        if self.query_name is None:
+            self.query_name = self.name
+
+
+# Pattern to match {{ queries.X }} references
+# Handles variations like {{queries.X}}, {{ queries.X }}, {{  queries.X  }}
+_QUERY_REF_PATTERN = re.compile(r"\{\{\s*queries\.(\w+)\s*\}\}")
+
+# Prefix for temp table names created by batch execution
+# This prefix helps identify dataface-created temp tables and avoids conflicts
+TEMP_TABLE_PREFIX = "_df_temp_"
+
+
+def extract_query_refs(sql: str) -> set[str]:
+    """Extract query references from SQL.
+
+    Finds all {{ queries.X }} patterns in the SQL string and returns
+    the set of referenced query names.
+
+    Args:
+        sql: SQL string potentially containing {{ queries.X }} references
+
+    Returns:
+        Set of query names referenced in the SQL
+
+    Example:
+        >>> extract_query_refs("SELECT * FROM {{ queries.orders }} WHERE amount > 100")
+        {'orders'}
+        >>> extract_query_refs("SELECT a.*, b.* FROM {{ queries.a }} a JOIN {{ queries.b }} b")
+        {'a', 'b'}
+    """
+    if not sql or "queries." not in sql:
+        return set()
+
+    return set(_QUERY_REF_PATTERN.findall(sql))
+
+
+def build_dependency_graph(
+    queries: dict[str, AnyQuery],
+) -> DependencyGraph:
+    """Build a dependency graph from a set of queries.
+
+    Scans each query's SQL for {{ queries.X }} references and builds
+    a directed graph where edges point from referenced queries to
+    the queries that reference them.
+
+    Args:
+        queries: Dict mapping query names to AnyQuery objects
+
+    Returns:
+        DependencyGraph with nodes and edges
+
+    Example:
+        Given queries:
+          orders: SELECT * FROM orders
+          high_value: SELECT * FROM {{ queries.orders }} WHERE amount > 1000
+
+        Returns graph with:
+          nodes: {orders, high_value}
+          edges: {orders: {high_value}}  (orders -> high_value)
+    """
+    graph = DependencyGraph()
+
+    for name, query in queries.items():
+        graph.add_node(name)
+
+        # Only SQL queries can have query references
+        if not is_sql_query(query):
+            continue
+
+        sql = query.sql
+        refs = extract_query_refs(sql)
+
+        for ref in refs:
+            if ref in queries:  # Only add edges for known queries
+                graph.add_edge(ref, name)
+
+    return graph
+
+
+def _find_cycle(graph: DependencyGraph) -> list[str]:
+    """Find a cycle in the dependency graph using DFS.
+
+    Args:
+        graph: The dependency graph to search
+
+    Returns:
+        List of query names forming the cycle, or empty list if no cycle
+    """
+    visited: set[str] = set()
+    rec_stack: set[str] = set()
+    path: list[str] = []
+
+    def dfs(node: str) -> list[str] | None:
+        visited.add(node)
+        rec_stack.add(node)
+        path.append(node)
+
+        for dep in graph.get_dependencies(node):
+            if dep not in visited:
+                cycle = dfs(dep)
+                if cycle:
+                    return cycle
+            elif dep in rec_stack:
+                # Found cycle - extract the cycle path
+                cycle_start = path.index(dep)
+                return path[cycle_start:] + [dep]
+
+        path.pop()
+        rec_stack.remove(node)
+        return None
+
+    for node in graph.all_nodes:
+        if node not in visited:
+            cycle = dfs(node)
+            if cycle:
+                return cycle
+
+    return []
+
+
+def topological_sort(graph: DependencyGraph) -> list[str]:
+    """Topologically sort nodes in the dependency graph.
+
+    Returns nodes in an order where dependencies come before dependents.
+    Uses Kahn's algorithm for topological sorting.
+
+    Args:
+        graph: The dependency graph to sort
+
+    Returns:
+        List of query names in topological order (dependencies first)
+
+    Raises:
+        ValueError: If the graph contains a cycle (includes cycle path in message)
+
+    Example:
+        If orders -> high_value (high_value depends on orders),
+        returns ['orders', 'high_value']
+    """
+    # Compute in-degree (number of dependencies) for each node
+    in_degree: dict[str, int] = dict.fromkeys(graph.all_nodes, 0)
+    for node in graph.all_nodes:
+        for dependent in graph.get_dependents(node):
+            in_degree[dependent] = in_degree.get(dependent, 0) + 1
+
+    # Start with nodes that have no dependencies
+    queue: list[str] = [node for node in graph.all_nodes if in_degree[node] == 0]
+    result: list[str] = []
+
+    while queue:
+        # Sort for deterministic ordering
+        queue.sort()
+        node = queue.pop(0)
+        result.append(node)
+
+        # Reduce in-degree for dependents
+        for dependent in sorted(graph.get_dependents(node)):
+            in_degree[dependent] -= 1
+            if in_degree[dependent] == 0:
+                queue.append(dependent)
+
+    if len(result) != len(graph.all_nodes):
+        # Find and report the actual cycle for better debugging
+        cycle = _find_cycle(graph)
+        if cycle:
+            cycle_str = " -> ".join(cycle)
+            raise ValueError(
+                f"Circular dependency detected in query references: {cycle_str}"
+            )
+        raise ValueError("Circular dependency detected in query references")
+
+    return result
+
+
+def collect_layout_chart_query_names(face: Face) -> set[str]:
+    """Walk the layout tree and return query names directly used by charts.
+
+    Shared layout-walk utility used by both ``collect_face_queries``
+    (which expands transitive deps) and the parallel pre-execution path
+    in ``renderer.py`` (which does not, since ``{{ queries.X }}`` is
+    Jinja-inlined).
+    """
+    names: set[str] = set()
+
+    def _walk(items: list[Any]) -> None:
+        for item in items:
+            if item.chart and item.chart.query_name:
+                names.add(item.chart.query_name)
+            if item.face and item.face.layout:
+                _walk(item.face.layout.items)
+
+    _walk(face.layout.items)
+
+    return names
+
+
+def collect_face_queries(
+    face: Face,
+    chart_queries_only: bool = True,
+) -> dict[str, AnyQuery]:
+    """Collect all queries needed for a dataface.
+
+    Collects queries directly used by charts, plus any queries they
+    reference via {{ queries.X }}.
+
+    Args:
+        face: Face to collect queries from
+        chart_queries_only: If True, only include queries used by charts.
+            If False, include all queries defined in the face.
+
+    Returns:
+        Dict mapping query names to AnyQuery objects
+
+    Example:
+        Given face with:
+          queries: {orders, high_value, recent, unused}
+          charts: [chart using high_value, chart using recent]
+
+        Returns {orders, high_value, recent} (orders is a dependency)
+    """
+    queries: dict[str, AnyQuery] = {}
+
+    if not chart_queries_only:
+        # Include all queries from the face
+        return dict(face.queries)
+
+    # Collect queries used by charts
+    chart_query_names = collect_layout_chart_query_names(face)
+
+    # Also check charts dict directly
+    for chart in face.charts.values():
+        if chart.query_name:
+            chart_query_names.add(chart.query_name)
+
+    # Collect queries and their dependencies
+    def add_query_and_deps(name: str) -> None:
+        """Recursively add a query and its dependencies."""
+        if name in queries or name not in face.queries:
+            return
+
+        query = face.queries[name]
+        queries[name] = query
+
+        # Find dependencies
+        if is_sql_query(query):
+            refs = extract_query_refs(query.sql)
+            for ref in refs:
+                add_query_and_deps(ref)
+
+    for query_name in chart_query_names:
+        add_query_and_deps(query_name)
+
+    return queries
+
+
+def get_query_source(query: AnyQuery) -> str:
+    """Get the database source for a query.
+
+    Args:
+        query: AnyQuery object
+
+    Returns:
+        Source string (defaults to 'default' if not specified)
+    """
+    source = getattr(query, "source", None)
+    if isinstance(source, str):
+        return source
+    return "default"
+
+
+def group_by_source(
+    queries: dict[str, AnyQuery],
+) -> dict[str, dict[str, AnyQuery]]:
+    """Group queries by their database source.
+
+    Args:
+        queries: Dict mapping query names to queries
+
+    Returns:
+        Dict mapping source to dict of queries with that source
+
+    Example:
+        Given queries with sources:
+          orders: postgres
+          users: postgres
+          events: mysql
+
+        Returns:
+          {'postgres': {'orders': ..., 'users': ...}, 'mysql': {'events': ...}}
+    """
+    groups: dict[str, dict[str, AnyQuery]] = {}
+
+    for name, query in queries.items():
+        source = get_query_source(query)
+        if source not in groups:
+            groups[source] = {}
+        groups[source][name] = query
+
+    return groups
+
+
+def validate_cross_source_references(
+    queries: dict[str, AnyQuery],
+    graph: DependencyGraph,
+) -> None:
+    """Validate that no queries reference queries on different sources.
+
+    Args:
+        queries: Dict of queries to validate
+        graph: Dependency graph for the queries
+
+    Raises:
+        CrossSourceReferenceError: If a cross-source reference is detected
+    """
+    for name, query in queries.items():
+        if not is_sql_query(query):
+            continue
+
+        query_source = get_query_source(query)
+        deps = graph.get_dependencies(name)
+
+        for dep_name in deps:
+            if dep_name not in queries:
+                continue
+
+            dep_query = queries[dep_name]
+            dep_source = get_query_source(dep_query)
+
+            if query_source != dep_source:
+                raise CrossSourceReferenceError(
+                    source_query=name,
+                    query_source=query_source,
+                    target_query=dep_name,
+                    target_source=dep_source,
+                )
+
+
+class TempTableNotSupportedError(Exception):
+    """Error raised when the dialect doesn't support temporary tables.
+
+    Some databases (e.g., certain cloud analytics platforms) may not support
+    temp tables, in which case batch optimization cannot be used.
+    """
+
+    def __init__(self, dialect_name: str):
+        self.dialect_name = dialect_name
+        super().__init__(
+            f"Dialect '{dialect_name}' does not support temporary tables. "
+            f"Batch execution with temp table optimization is not available."
+        )
+
+
+def generate_batch_sql(
+    queries: dict[str, AnyQuery],
+    graph: DependencyGraph,
+    dialect: SQLDialect,
+    direct_chart_queries: set[str] | None = None,
+) -> list[BatchStatement]:
+    """Generate batch SQL statements with temp table optimization.
+
+    For queries that are referenced by other queries:
+    1. Create a temp table with the base query
+    2. Rewrite dependent queries to SELECT from the temp table
+
+    Args:
+        queries: Dict of queries to include in the batch
+        graph: Dependency graph for the queries
+        dialect: SQL dialect for temp table syntax
+        direct_chart_queries: Set of query names directly used by charts.
+            If a temp table query is also used directly, we add a SELECT *
+            statement for it.
+
+    Returns:
+        List of BatchStatement objects in execution order
+
+    Raises:
+        TempTableNotSupportedError: If dialect doesn't support temp tables
+
+    Example:
+        Given:
+          orders: SELECT * FROM orders
+          high_value: SELECT * FROM {{ queries.orders }} WHERE amount > 1000
+
+        Returns:
+          [
+            BatchStatement("_create_orders", "CREATE TEMP TABLE _df_temp_orders AS SELECT * FROM orders", True),
+            BatchStatement("orders", "SELECT * FROM _df_temp_orders"),  # If orders is directly used
+            BatchStatement("high_value", "SELECT * FROM _df_temp_orders WHERE amount > 1000"),
+          ]
+    """
+    if direct_chart_queries is None:
+        direct_chart_queries = set()
+
+    statements: list[BatchStatement] = []
+    temp_candidates = graph.get_temp_table_candidates()
+    sorted_queries = topological_sort(graph)
+
+    # Track which temp tables have been created
+    created_temp_tables: set[str] = set()
+
+    # Check if dialect supports temp tables (only needed if we have candidates)
+    if temp_candidates and not dialect.supports_temp_tables:
+        raise TempTableNotSupportedError(dialect.name)
+
+    for query_name in sorted_queries:
+        if query_name not in queries:
+            continue
+
+        query = queries[query_name]
+
+        # Only SQL queries can be batched with temp tables
+        if not is_sql_query(query):
+            # For non-SQL queries, just add as-is (they'll be executed separately)
+            continue
+
+        sql = query.sql
+        setup = query.setup_sql.strip() + "\n" if query.setup_sql else ""
+
+        # Rewrite {{ queries.X }} refs to temp table names
+        # Note: We use a precompiled pattern (_QUERY_REF_PATTERN) for extraction,
+        # but need a dynamic pattern here to match specific query names
+        for ref in extract_query_refs(sql):
+            if ref in created_temp_tables:
+                temp_ref = dialect.temp_table_ref(f"{TEMP_TABLE_PREFIX}{ref}")
+                # Replace the Jinja template with the temp table reference
+                # Use re.escape() to handle any regex metacharacters in query names
+                escaped_ref = re.escape(ref)
+                pattern = rf"\{{\{{\s*queries\.{escaped_ref}\s*\}}\}}"
+                sql = re.sub(pattern, temp_ref, sql)
+
+        if query_name in temp_candidates:
+            # This query is referenced by other queries, so we create a temp table.
+            # This ensures the query runs only ONCE, and dependents read from the
+            # temp table instead of re-executing the query.
+            temp_name = f"{TEMP_TABLE_PREFIX}{query_name}"
+            temp_sql = setup + dialect.temp_table(temp_name, sql)
+            statements.append(
+                BatchStatement(
+                    name=f"_create_{query_name}",
+                    sql=temp_sql,
+                    is_temp_table_create=True,
+                    query_name=query_name,
+                )
+            )
+            created_temp_tables.add(query_name)
+
+            # Special case: Query is BOTH a dependency AND directly used by a chart.
+            # Example: "orders" is referenced by "high_value" AND has its own chart.
+            # We need TWO statements:
+            #   1. CREATE TEMP TABLE (above) - for dependents to reference
+            #   2. SELECT * FROM temp table (below) - to get results for the chart
+            if query_name in direct_chart_queries:
+                temp_ref = dialect.temp_table_ref(temp_name)
+                statements.append(
+                    BatchStatement(
+                        name=query_name,
+                        sql=f"SELECT * FROM {temp_ref}",
+                        is_temp_table_create=False,
+                        query_name=query_name,
+                    )
+                )
+        else:
+            # Regular query (not a temp table candidate)
+            statements.append(
+                BatchStatement(
+                    name=query_name,
+                    sql=setup + sql,
+                    is_temp_table_create=False,
+                    query_name=query_name,
+                )
+            )
+
+    return statements
+
+
+def create_batch_execution_plan(
+    face: Face,
+) -> dict[str, list[BatchStatement]]:
+    """Create a complete batch execution plan for a dataface.
+
+    This is the main entry point for batch execution planning.
+    It collects queries, groups by source, validates cross-source
+    references, and generates batch SQL for each source.
+
+    Args:
+        face: Face to create execution plan for
+
+    Returns:
+        Dict mapping source to list of BatchStatements
+
+    Raises:
+        CrossSourceReferenceError: If a query references another
+            query on a different database source
+
+    Example:
+        >>> plan = create_batch_execution_plan(face)
+        >>> for source, statements in plan.items():
+        ...     print(f"Source: {source}")
+        ...     for stmt in statements:
+        ...         print(f"  {stmt.name}: {stmt.sql[:50]}...")
+    """
+    # Collect all queries used by the dataface
+    all_queries = collect_face_queries(face, chart_queries_only=True)
+
+    if not all_queries:
+        return {}
+
+    # Build the full dependency graph first (for validation)
+    full_graph = build_dependency_graph(all_queries)
+
+    # Validate no cross-source references
+    validate_cross_source_references(all_queries, full_graph)
+
+    # Collect query names directly used by charts
+    direct_chart_queries: set[str] = set()
+    for chart in face.charts.values():
+        if chart.query_name:
+            direct_chart_queries.add(chart.query_name)
+
+    # Group queries by source
+    source_groups = group_by_source(all_queries)
+
+    # Generate batch SQL for each source
+    result: dict[str, list[BatchStatement]] = {}
+
+    for source, source_queries in source_groups.items():
+        # Build graph for this source's queries
+        graph = build_dependency_graph(source_queries)
+
+        # Get dialect for this source
+        dialect = get_dialect(source)
+
+        # Generate batch statements
+        statements = generate_batch_sql(
+            queries=source_queries,
+            graph=graph,
+            dialect=dialect,
+            direct_chart_queries=direct_chart_queries,
+        )
+
+        if statements:
+            result[source] = statements
+
+    return result