dataface 0.1.2__py3-none-any.whl

dataface/core/compile/normalize_queries.py

@@ -0,0 +1,297 @@
+"""Query normalization and external query resolution."""
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import ValidationError as PydanticValidationError
+
+from dataface.core.compile.errors import CompilationError
+from dataface.core.compile.filter_injection import extract_filter_variable_deps
+from dataface.core.compile.jinja import extract_variable_dependencies
+from dataface.core.compile.models.query.authored import AuthoredQuery
+from dataface.core.compile.models.query.compiled import (
+    AnyQuery,
+    CsvQuery,
+    DbtModelQuery,
+    HttpQuery,
+    MetricFlowQuery,
+    SchemaResolverQuery,
+    SqlQuery,
+    ValuesQuery,
+    is_sql_query,
+)
+from dataface.core.compile.models.refs import infer_query_type_from_keys
+from dataface.core.compile.sql_authoring_lint import has_literal_escaped_newlines
+from dataface.core.errors.codes_compile import DF_COMPILE_SQL_LITERAL_NEWLINES
+from dataface.core.execute.errors import MutatingSqlError, UnparseableSqlError
+from dataface.core.execute.sql_guard import validate_select_only, validate_setup_sql
+
+logger = logging.getLogger(__name__)
+
+
+def normalize_query(
+    name: str,
+    query_def: Any,
+    default_source: str | None = None,
+) -> AnyQuery:
+    """Normalize a query definition to AnyQuery.
+
+    Creates the appropriate query type (SqlQuery, CsvQuery, etc.)
+    based on the query definition.
+
+    Args:
+        name: Query name
+        query_def: AuthoredQuery definition (AuthoredQuery, dict, or bare SQL string shorthand)
+        default_source: Default source to apply if query has no explicit source
+
+    Returns:
+        Query object (SqlQuery, CsvQuery, etc.)
+    """
+    # Non-SQL pre-built query types need no guard — return immediately.
+    if isinstance(
+        query_def,
+        (
+            CsvQuery,
+            MetricFlowQuery,
+            HttpQuery,
+            DbtModelQuery,
+            ValuesQuery,
+            SchemaResolverQuery,
+        ),
+    ):
+        return query_def
+
+    # Pre-built SqlQuery: skip reconstruction but still run guard + dep-extraction below.
+    if isinstance(query_def, SqlQuery):
+        query: AnyQuery = query_def
+    else:
+        # Convert to dict
+        if isinstance(query_def, AuthoredQuery):
+            query_dict = query_def.model_dump(exclude_none=True)
+        elif isinstance(query_def, dict):
+            query_dict = dict(query_def)
+        elif isinstance(query_def, str):
+            query_dict = {"sql": query_def}
+        else:
+            raise CompilationError(
+                f"Query '{name}': Invalid query definition. Expected dict with 'sql' and 'source' fields for SQL queries."
+            )
+
+        # Ensure type field
+        if "type" not in query_dict:
+            query_dict["type"] = infer_query_type_from_keys(query_dict)
+
+        query_type = query_dict.pop("type")
+
+        # Apply default source if query doesn't have one.
+        # ValuesQuery is inline data; SchemaResolverQuery's source is a dbt source name,
+        # not a connection reference; HttpQuery uses url, not source — all three are excluded.
+        has_source = query_dict.get("source") is not None
+        if (
+            not has_source
+            and default_source
+            and query_type not in {"values", "dbt_model", "schema_resolver", "http"}
+        ):
+            query_dict["source"] = default_source
+
+        # Create the appropriate query type
+        if query_type == "sql":
+            # Ensure sql is present
+            if "sql" not in query_dict:
+                raise CompilationError(
+                    f"Query '{name}': SQL queries must have 'sql' field"
+                )
+            # Check for source
+            if "source" not in query_dict:
+                raise CompilationError(
+                    f"Query '{name}': SQL queries must have 'source' field "
+                    "to specify which database connection to use. "
+                    "You can also set a default source at the face level."
+                )
+            query = SqlQuery(**query_dict)
+        elif query_type == "csv":
+            # CSV queries need either file or source with file config
+            if "file" not in query_dict and "source" not in query_dict:
+                raise CompilationError(
+                    f"Query '{name}': CSV queries must have 'file' field or 'source' reference"
+                )
+            query = CsvQuery(**query_dict)
+        elif query_type == "metricflow":
+            query = MetricFlowQuery(**query_dict)
+        elif query_type == "http":
+            if "url" not in query_dict:
+                raise CompilationError(
+                    f"Query '{name}': HTTP queries require a 'url' field"
+                )
+            try:
+                query = HttpQuery(**query_dict)
+            except PydanticValidationError as e:
+                fields = ", ".join(
+                    str(err["loc"][0]) for err in e.errors() if err["loc"]
+                )
+                raise CompilationError(
+                    f"Query '{name}': invalid field(s) for HTTP query: {fields}"
+                ) from e
+        elif query_type == "dbt_model":
+            query = DbtModelQuery(**query_dict)
+        elif query_type == "values":
+            query = ValuesQuery(**query_dict)
+        elif query_type == "schema_resolver":
+            query = SchemaResolverQuery(**query_dict)
+        else:
+            # Default to SQL
+            query = SqlQuery(**query_dict)
+
+    # Validate setup_sql doesn't contain query references (non-nestable by design)
+    if is_sql_query(query) and query.setup_sql and "{{ queries." in query.setup_sql:
+        raise CompilationError(
+            f"Query '{name}': setup_sql must not contain {{{{ queries.X }}}} references. "
+            "setup_sql is for non-nestable setup statements (CREATE TEMP FUNCTION, etc.)."
+        )
+
+    # Compile-time SQL guard + variable-dep extraction for SQL queries.
+    # UnparseableSqlError is silently deferred (compile lacks the dbt manifest for
+    # {{ ref() }} resolution and must not reject faces with complex Jinja macros).
+    # MutatingSqlError is fatal at compile time — wrap in CompilationError with query name.
+    deps: set = set()
+    if is_sql_query(query):
+        dialect = query.source.get("type") if isinstance(query.source, dict) else None
+
+        # Authoring lint: detect literal \n (backslash + n) from single-quoted YAML.
+        # Must run before validate_select_only because literal \n causes parse failure
+        # that would otherwise be silently deferred as UnparseableSqlError.
+        if has_literal_escaped_newlines(query.sql, dialect=dialect):
+            raise CompilationError.from_code(
+                DF_COMPILE_SQL_LITERAL_NEWLINES, query_name=name, field_label="sql"
+            )
+        if query.setup_sql and has_literal_escaped_newlines(
+            query.setup_sql, dialect=dialect
+        ):
+            raise CompilationError.from_code(
+                DF_COMPILE_SQL_LITERAL_NEWLINES,
+                query_name=name,
+                field_label="setup_sql",
+            )
+
+        try:
+            validate_select_only(query.sql, dialect=dialect)
+        except UnparseableSqlError:
+            logger.debug("compile-time SQL validation deferred to runtime: %s", name)
+        except MutatingSqlError as e:
+            raise CompilationError(f"Query '{name}': {e}") from e
+
+        if query.setup_sql:
+            try:
+                validate_setup_sql(query.setup_sql, dialect=dialect)
+            except UnparseableSqlError:
+                logger.debug(
+                    "compile-time setup_sql validation deferred to runtime: %s", name
+                )
+            except MutatingSqlError as e:
+                raise CompilationError(f"Query '{name}' setup_sql: {e}") from e
+
+        if query.sql:
+            deps |= extract_variable_dependencies(query.sql)
+        if query.setup_sql:
+            deps |= extract_variable_dependencies(query.setup_sql)
+    # Declarative filters: variable names are direct values, not Jinja templates
+    if query.filters:
+        deps |= extract_filter_variable_deps(query.filters)
+    query.variable_dependencies = deps
+
+    return query
+
+
+# Cache for loaded external query files
+_external_query_cache: dict[str, dict[str, AnyQuery]] = {}
+
+
+def resolve_external_query(
+    query_ref: str,
+    query_registry: dict[str, AnyQuery],
+    base_path: Path | None = None,
+) -> tuple[str, dict[str, AnyQuery]]:
+    """Resolve an external query reference.
+
+    Parses references like `path/to/file.yml#query_name` and loads
+    queries from the external file.
+
+    Args:
+        query_ref: Query reference with file path (e.g., "_shared.yml#sales")
+        query_registry: Current query registry to merge into
+        base_path: Base path for resolving relative file paths
+
+    Returns:
+        Tuple of (resolved_query_name, updated_query_registry)
+
+    Raises:
+        CompilationError: If file not found or query not in file
+    """
+    # Parse file path and query name
+    file_part, query_name = query_ref.split("#", 1)
+
+    if not file_part or not query_name:
+        raise CompilationError(
+            f"Invalid external query reference: '{query_ref}'. "
+            "Expected format: 'path/to/file.yml#query_name'"
+        )
+
+    # Resolve file path
+    file_path = base_path / file_part if base_path else Path(file_part)
+
+    # Create a unique registry key for this file's queries
+    file_key = str(file_path.resolve()) if file_path.exists() else str(file_path)
+
+    # Check cache first
+    if file_key not in _external_query_cache:
+        # Load and parse the external file
+        if not file_path.exists():
+            raise CompilationError(
+                f"External query file not found: '{file_path}' "
+                f"(referenced as '{query_ref}')"
+            )
+
+        try:
+            with open(file_path) as f:
+                external_data = yaml.safe_load(f)
+        except yaml.YAMLError as e:
+            raise CompilationError(
+                f"Failed to parse external query file '{file_path}': {e}"
+            ) from e
+
+        if not isinstance(external_data, dict):
+            raise CompilationError(
+                f"External query file '{file_path}' must be a YAML dictionary"
+            )
+
+        # Extract and normalize queries from the file
+        external_queries: dict[str, AnyQuery] = {}
+        queries_section = external_data.get("queries", {})
+
+        for ext_query_name, ext_query_def in queries_section.items():
+            compiled_query = normalize_query(ext_query_name, ext_query_def)
+            external_queries[ext_query_name] = compiled_query
+
+        _external_query_cache[file_key] = external_queries
+
+    # Get cached queries for this file
+    file_queries = _external_query_cache[file_key]
+
+    if query_name not in file_queries:
+        available = ", ".join(file_queries.keys()) if file_queries else "none"
+        raise CompilationError(
+            f"Query '{query_name}' not found in '{file_path}'. "
+            f"Available queries: {available}"
+        )
+
+    # Create a unique key for this external query in the registry
+    # Use file path + query name to avoid collisions
+    registry_key = f"{file_part}#{query_name}"
+
+    # Add to registry if not already present
+    if registry_key not in query_registry:
+        query_registry[registry_key] = file_queries[query_name]
+
+    return registry_key, query_registry