dataface 0.1.2__py3-none-any.whl

dataface/core/execute/parallel.py

@@ -0,0 +1,162 @@
+"""Query execution helpers for the two-phase render pipeline.
+
+Phase 1: Execute chart-direct queries that do *not* use ``{{ results.X }}``
+concurrently with ``ThreadPoolExecutor``.
+
+Phase 2: Execute chart-direct ``{{ results.X }}`` queries sequentially after
+phase 1, before the render walk begins.
+
+Only queries directly referenced by charts are pre-executed. Transitive
+``{{ queries.X }}`` dependencies are *not* pre-executed separately because
+the executor inlines them as subqueries — executing the dependency query on
+its own would be wasted work.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from concurrent.futures import Future, ThreadPoolExecutor, as_completed
+from typing import TYPE_CHECKING
+
+from dataface.core.compile.models.face.compiled import Face, VariableValues
+from dataface.core.compile.models.query.compiled import (
+    AnyQuery,
+    is_sql_query,
+)
+from dataface.core.execute.batch import DependencyGraph, topological_sort
+
+if TYPE_CHECKING:
+    from dataface.core.execute.executor import Executor
+
+logger = logging.getLogger(__name__)
+
+_RESULTS_REF_PATTERN = re.compile(r"\{\{\s*results\.(\w+)\s*\}\}")
+
+
+def _has_results_refs(query: AnyQuery) -> bool:
+    """Check if a SQL query references ``{{ results.X }}``."""
+    if not is_sql_query(query):
+        return False
+    return "{{ results." in query.sql or "{{results." in query.sql
+
+
+def _extract_results_refs(query: AnyQuery) -> set[str]:
+    """Extract ``{{ results.X }}`` references from a SQL query."""
+    if not is_sql_query(query):
+        return set()
+    return set(_RESULTS_REF_PATTERN.findall(query.sql))
+
+
+def plan_query_execution_phases(
+    face: Face,
+    query_names: set[str],
+) -> tuple[set[str], list[str]]:
+    """Split chart-direct queries into parallel and sequential execution phases.
+
+    Queries that reference ``{{ results.X }}`` must run after the parallel-safe
+    phase because they depend on previously executed query results in DuckDB.
+    These queries still run before rendering; they are not deferred as a render
+    time fallback.
+    """
+    parallel_names: set[str] = set()
+    sequential_queries: dict[str, AnyQuery] = {}
+
+    for name in query_names:
+        query = face.queries.get(name)
+        if query and _has_results_refs(query):
+            sequential_queries[name] = query
+        else:
+            parallel_names.add(name)
+
+    if not sequential_queries:
+        return parallel_names, []
+
+    graph = DependencyGraph()
+    sequential_names = set(sequential_queries)
+    for name, query in sequential_queries.items():
+        graph.add_node(name)
+        for ref in _extract_results_refs(query) & sequential_names:
+            graph.add_edge(ref, name)
+
+    return parallel_names, topological_sort(graph)
+
+
+def execute_queries_parallel(
+    executor: Executor,
+    face: Face,
+    query_names: set[str],
+    variables: VariableValues | None = None,
+    max_workers: int = 8,
+) -> dict[str, Exception | None]:
+    """Execute *query_names* concurrently via ``ThreadPoolExecutor``.
+
+    All queries are submitted immediately — ``{{ queries.X }}`` references are
+    Jinja-inlined by the executor so no ordering is needed.
+
+    ``{{ results.X }}`` queries are intentionally *not* accepted here. They
+    must be routed through the explicit sequential phase instead of falling
+    through to a render-time execution path.
+
+    Returns ``{query_name: None}`` on success or ``{query_name: exc}``
+    on failure.  Errors in one query don't block others.
+    """
+    if not query_names:
+        return {}
+
+    unsafe_names = sorted(
+        name
+        for name in query_names
+        if (query := face.queries.get(name)) and _has_results_refs(query)
+    )
+    if unsafe_names:
+        joined = ", ".join(unsafe_names)
+        raise ValueError(
+            "execute_queries_parallel received queries with {{ results.X }} "
+            f"references: {joined}"
+        )
+
+    results: dict[str, Exception | None] = {}
+
+    with ThreadPoolExecutor(max_workers=max_workers) as pool:
+        futures: dict[Future, str] = {
+            pool.submit(executor.execute_query, name, variables): name
+            for name in query_names
+        }
+
+        for fut in as_completed(futures):
+            name = futures[fut]
+            try:
+                fut.result()
+                results[name] = None
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("Parallel execution of query '%s' failed: %s", name, exc)
+                results[name] = exc
+                executor._query_errors[name] = exc
+
+    return results
+
+
+def execute_queries_sequential(
+    executor: Executor,
+    query_names: list[str],
+    variables: VariableValues | None = None,
+) -> dict[str, Exception | None]:
+    """Execute *query_names* sequentially before the render walk.
+
+    This phase is for direct chart queries that use ``{{ results.X }}``.
+    They have runtime dependencies on previously executed query results, so
+    they run after the parallel-safe phase but still before rendering.
+    """
+    results: dict[str, Exception | None] = {}
+
+    for name in query_names:
+        try:
+            executor.execute_query(name, variables)
+            results[name] = None
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("Sequential execution of query '%s' failed: %s", name, exc)
+            results[name] = exc
+            executor._query_errors[name] = exc
+
+    return results

dataface/core/execute/setup_sql.py

@@ -0,0 +1,58 @@
+"""Collect and order setup_sql preambles across query dependencies.
+
+When query B depends on query A via {{ queries.A }}, and A declares
+setup_sql, that preamble must execute before B's main body runs.
+This module walks the dependency graph, collects setup_sql blocks
+in topological order, and deduplicates identical blocks.
+"""
+
+import re
+
+from dataface.core.compile.models.query.compiled import (
+    AnyQuery,
+    is_sql_query,
+)
+
+_QUERY_REF_PATTERN = re.compile(r"\{\{\s*queries\.(\w+)\s*\}\}")
+
+
+def collect_setup_sql(
+    query_name: str,
+    queries: dict[str, AnyQuery],
+) -> list[str]:
+    """Collect setup_sql statements for a query and all its dependencies.
+
+    Walks the dependency graph depth-first, returns setup_sql blocks in
+    dependency order (dependencies before dependents). Identical blocks
+    are deduplicated while preserving the earliest occurrence.
+
+    Args:
+        query_name: Name of the query to collect setup_sql for.
+        queries: Full query registry.
+
+    Returns:
+        Ordered list of unique setup_sql strings to execute before the main query.
+    """
+    seen_names: set[str] = set()
+    ordered: list[str] = []
+    seen_sql: set[str] = set()
+
+    def _walk(name: str) -> None:
+        if name in seen_names or name not in queries:
+            return
+        seen_names.add(name)
+
+        query = queries[name]
+        # Walk dependencies first (depth-first, dependencies before self)
+        if is_sql_query(query):
+            for ref in _QUERY_REF_PATTERN.findall(query.sql):
+                _walk(ref)
+
+        # Then collect this query's setup_sql (after deps)
+        if is_sql_query(query) and query.setup_sql:
+            if query.setup_sql not in seen_sql:
+                seen_sql.add(query.setup_sql)
+                ordered.append(query.setup_sql)
+
+    _walk(query_name)
+    return ordered

dataface/core/execute/source_registry.py

@@ -0,0 +1,72 @@
+"""Project-source state for the execute boundary.
+
+Owns the project-source allowlist that the resolver consults: ``_sources.yaml``
+/ ``dataface.yml`` (already cached by ``get_project_sources``) plus any
+synthetic sources registered at runtime (the ``dft inspect`` HTML server is
+the canonical registrar — it synthesizes a ``warehouse`` source when no
+``_sources.yaml`` exists).
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+from dataface.core.compile.compiler import resolve_project_source_paths
+from dataface.core.compile.config import ProjectSourcesConfig, get_project_sources
+
+logger = logging.getLogger(__name__)
+
+
+class SourceRegistry:
+    """Merged view of project-level sources plus synthetic registrations.
+
+    A project source by a given name always wins over a synthetic registration
+    of the same name (``register`` uses set-if-absent semantics on the synthetic
+    map; the project map is layered on top at read time).
+
+    ``project_root=None`` is supported because ``AdapterRegistry`` accepts the
+    same — ad-hoc constructions (tests, agent_api one-offs) yield an empty
+    source map.
+    """
+
+    def __init__(self, project_root: Path | None = None) -> None:
+        self._project_root = project_root
+        self._synthetic: dict[str, dict[str, Any]] = {}
+
+    def get(self, name: str) -> dict[str, Any] | None:
+        """Return the resolved config for ``name``, or ``None`` if unknown."""
+        return self.all().get(name)
+
+    def all(self) -> dict[str, dict[str, Any]]:
+        """Return the merged source map (project sources win over synthetics)."""
+        return {**self._synthetic, **self._load_project()}
+
+    def register(self, name: str, config: dict[str, Any]) -> None:
+        """Register a synthetic source. Idempotent: first registration wins."""
+        self._synthetic.setdefault(name, dict(config))
+
+    def project_sources_config(self) -> ProjectSourcesConfig:
+        """Snapshot for the resolver — merged sources plus project default."""
+        return ProjectSourcesConfig(default=self._project_default(), sources=self.all())
+
+    def _load_project(self) -> dict[str, dict[str, Any]]:
+        if self._project_root is None:
+            return {}
+        try:
+            project_sources = get_project_sources(self._project_root)
+            return resolve_project_source_paths(
+                project_sources.sources, self._project_root
+            )
+        except (OSError, TypeError, ValueError) as e:
+            logger.warning("Failed to load project sources: %s", e)
+            return {}
+
+    def _project_default(self) -> str | None:
+        if self._project_root is None:
+            return None
+        try:
+            return get_project_sources(self._project_root).default
+        except (OSError, TypeError, ValueError):
+            return None

dataface/core/execute/source_resolver.py

@@ -0,0 +1,255 @@
+"""Source resolver: execute-boundary resolution of authored `source:` values.
+
+The `SourceResolver` Protocol runs in `AdapterRegistry.execute()` before the
+adapter receives the query. When the authored source resolves to a configured
+profile (face-level, project-level, or an inline dict), the adapter receives a
+typed `SourceConfig`. For source-less queries with no project default
+configured, the resolver returns `None` and the adapter falls through to its
+own default connection (DuckDB ``:memory:`` or the dbt-aware lookup).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from dataface.core.compile.config import ProjectSourcesConfig
+from dataface.core.compile.errors import DatafaceError
+from dataface.core.compile.models.source import (
+    VALID_SOURCE_TYPES,
+    SourceConfig,
+    parse_source_config,
+)
+from dataface.core.errors.codes_execute import (
+    DF_EXECUTE_SOURCE_CROSS_FILE_FORBIDDEN,
+    DF_EXECUTE_SOURCE_INLINE_FORBIDDEN,
+    DF_EXECUTE_SOURCE_INVALID_TYPE,
+    DF_EXECUTE_SOURCE_MISSING_TYPE,
+    DF_EXECUTE_SOURCE_NOT_FOUND,
+    DF_EXECUTE_SOURCE_NOT_FOUND_EMPTY,
+)
+
+
+@dataclass
+class DbtContext:
+    """Opaque marker: a dbt project is in scope for this execute call.
+
+    When present, unknown string source names fall through to the dbt adapter
+    (which resolves profile names via the dbt manifest).
+    When absent, unknown names raise DF-EXECUTE-SOURCE-NOT-FOUND.
+    """
+
+
+class SourceResolver(Protocol):
+    """Single execute-boundary chokepoint for source resolution.
+
+    Implementations map an authored `source:` value (str | dict | None) to a
+    typed `SourceConfig | None` before any database driver is loaded.
+
+    Callers (AdapterRegistry) invoke resolve() once per query, then forward the
+    result to the matched adapter via a `source_config` kwarg.
+    """
+
+    def resolve(
+        self,
+        authored: str | dict[str, Any] | None,
+        face_sources: dict[str, dict[str, Any]],
+        project_sources: ProjectSourcesConfig,
+        dbt_context: DbtContext | None,
+    ) -> SourceConfig | None:
+        """Resolve an authored source value to a typed SourceConfig.
+
+        Args:
+            authored: The raw `source:` value from the compiled query.
+            face_sources: Named sources declared in the face's `sources:` block.
+            project_sources: Project-level sources from `_sources.yaml` / dataface.yml.
+            dbt_context: Present when a dbt project is in scope; allows unknown
+                string names to fall through to the dbt adapter.
+
+        Returns:
+            Typed SourceConfig, or None when no source applies (adapter uses
+            its own default connection) or when dbt_context covers the name.
+
+        Raises:
+            DatafaceError: With a DF-EXECUTE-SOURCE-* code on policy violations.
+        """
+        ...
+
+
+class DefaultSourceResolver:
+    """Behavior-preserving default resolver for CLI, local playground, and inspect surfaces.
+
+    Resolves authored source values using face-level sources first, then project-level
+    sources. Does not apply policy gates — inline dicts and unknown names are accepted
+    as-is. Unknown string names with a dbt context fall through to the dbt adapter.
+
+    Resolution order for a string authored value:
+      1. face_sources lookup (face's `sources:` block)
+      2. project_sources lookup (`_sources.yaml` / dataface.yml)
+      3. dbt fallback: return None when dbt_context is set
+      4. raise DF-EXECUTE-SOURCE-NOT-FOUND otherwise
+
+    For None authored: apply project default when configured, else return None.
+    For dict authored: validate and parse directly.
+
+    Subclasses (e.g. an allowlist-enforcing resolver) override resolve(), apply
+    their own policy gates, and delegate to super().resolve() for the shared
+    lookup algorithm.
+    """
+
+    def resolve(
+        self,
+        authored: str | dict[str, Any] | None,
+        face_sources: dict[str, dict[str, Any]],
+        project_sources: ProjectSourcesConfig,
+        dbt_context: DbtContext | None,
+    ) -> SourceConfig | None:
+        if authored is None:
+            default_name = project_sources.default
+            if not default_name:
+                return None
+            source_dict = project_sources.sources.get(default_name)
+            if source_dict is None:
+                return None
+            return self._parse(source_dict)
+
+        if isinstance(authored, dict):
+            return self._parse(authored)
+
+        source_dict = face_sources.get(authored) or project_sources.sources.get(
+            authored
+        )
+        if source_dict is not None:
+            return self._parse(source_dict)
+
+        if dbt_context is not None:
+            return None
+
+        available = sorted(project_sources.sources.keys())
+        if available:
+            raise DatafaceError.from_code(
+                DF_EXECUTE_SOURCE_NOT_FOUND,
+                source=authored,
+                available=available,
+            )
+        raise DatafaceError.from_code(
+            DF_EXECUTE_SOURCE_NOT_FOUND_EMPTY,
+            source=authored,
+        )
+
+    def _parse(self, source_dict: dict[str, Any]) -> SourceConfig:
+        """Validate and parse a source dict to a typed SourceConfig."""
+        if "type" not in source_dict:
+            raise DatafaceError.from_code(
+                DF_EXECUTE_SOURCE_MISSING_TYPE,
+                offending_value=repr(source_dict),
+            )
+        if source_dict["type"] not in VALID_SOURCE_TYPES:
+            raise DatafaceError.from_code(
+                DF_EXECUTE_SOURCE_INVALID_TYPE,
+                offending_value=source_dict["type"],
+                available=sorted(VALID_SOURCE_TYPES),
+            )
+        return parse_source_config(source_dict)
+
+
+def _summarize_inline_source(source_dict: dict[str, Any]) -> str:
+    """Summarize a rejected inline source without echoing secret values.
+
+    The rendered error message surfaces via `QueryResult.error` and structured
+    logs — echoing the raw dict (which may carry `password`, `secret`, `token`,
+    …) compounds the original mistake. Emit `type` plus the sorted key list so
+    authors get actionable feedback ("you authored an inline postgres source")
+    without the values.
+    """
+    type_value = source_dict.get("type", "<missing type>")
+    keys = sorted(source_dict.keys())
+    return f"type={type_value!r}, keys={keys}"
+
+
+class AllowlistedSourceResolver(DefaultSourceResolver):
+    """Strict source resolver: only names in `project_sources.sources` resolve.
+
+    Three rejection rules fire before any resolution:
+
+    1. Inline source dicts — both per-query `source: {type: ..., host: ...}`
+       and any face-level `face.sources` entry are refused with
+       `DF-EXECUTE-SOURCE-INLINE-FORBIDDEN` (face.sources values are inline
+       source definitions by construction, not name references).
+    2. Cross-file `#` references — any authored string containing `#` is
+       refused with `DF-EXECUTE-SOURCE-CROSS-FILE-FORBIDDEN`.
+    3. Unknown names — a string not in `project_sources.sources` is refused
+       with `DF-EXECUTE-SOURCE-NOT-FOUND`; the payload's `available` field
+       carries the sorted allowlist.
+
+    Accepted inputs delegate to `DefaultSourceResolver.resolve` for the shared
+    lookup. The dbt-context fallback is intentionally disabled: an unknown
+    name raises regardless of `dbt_context`, because dbt-profile names are not
+    part of the configured allowlist by definition.
+
+    This is the closed-allowlist resolver — pick it when the caller does not
+    trust authored YAML to declare arbitrary connections (e.g. hosted /
+    multi-tenant deployments). Loosening any rule changes the trust boundary;
+    treat changes here as security decisions, not stylistic ones.
+    """
+
+    def resolve(
+        self,
+        authored: str | dict[str, Any] | None,
+        face_sources: dict[str, dict[str, Any]],
+        project_sources: ProjectSourcesConfig,
+        dbt_context: DbtContext | None,
+        yaml_path: str = "",
+    ) -> SourceConfig | None:
+        if isinstance(authored, dict):
+            raise DatafaceError.from_code(
+                DF_EXECUTE_SOURCE_INLINE_FORBIDDEN,
+                yaml_path=yaml_path,
+                offending_value=_summarize_inline_source(authored),
+            )
+
+        if isinstance(authored, str) and "#" in authored:
+            raise DatafaceError.from_code(
+                DF_EXECUTE_SOURCE_CROSS_FILE_FORBIDDEN,
+                yaml_path=yaml_path,
+                offending_value=authored,
+            )
+
+        if (
+            isinstance(authored, str)
+            and authored in face_sources
+            and authored not in project_sources.sources
+        ):
+            raise DatafaceError.from_code(
+                DF_EXECUTE_SOURCE_INLINE_FORBIDDEN,
+                yaml_path=yaml_path,
+                offending_value=_summarize_inline_source(face_sources[authored]),
+            )
+
+        if isinstance(authored, str):
+            if authored not in project_sources.sources:
+                self._raise_not_found(authored, project_sources, yaml_path)
+        elif authored is None:
+            default_name = project_sources.default
+            if default_name is not None and default_name not in project_sources.sources:
+                self._raise_not_found(default_name, project_sources, yaml_path)
+
+        return super().resolve(
+            authored=authored,
+            face_sources={},
+            project_sources=project_sources,
+            dbt_context=None,
+        )
+
+    @staticmethod
+    def _raise_not_found(
+        offending: str,
+        project_sources: ProjectSourcesConfig,
+        yaml_path: str,
+    ) -> None:
+        raise DatafaceError.from_code(
+            DF_EXECUTE_SOURCE_NOT_FOUND,
+            source=offending,
+            available=sorted(project_sources.sources.keys()),
+            yaml_path=yaml_path,
+        )