dataface 0.1.2__py3-none-any.whl

dataface/agent_api/query.py

@@ -0,0 +1,274 @@
+"""Typed query verbs for the agent API."""
+
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from dataface.agent_api._paths import resolve_scoped_path
+from dataface.core.compile import compile_file
+from dataface.core.compile.models.query.compiled import SqlQuery
+from dataface.core.execute.adapters import AdapterRegistry
+from dataface.core.validate import normalize_data_for_json
+
+MAX_QUERY_LIMIT = 1000
+
+
+class FaceQueryLookupResult(BaseModel):
+    """Result of extracting SQL from a named face query for offline use (validate/describe)."""
+
+    success: bool
+    sql: str = ""
+    source: str | None = None
+    errors: list[str] = []
+    available_queries: list[str] = []
+
+
+def lookup_face_query_sql(
+    name: str,
+    path: Path,
+    project_dir: Path | None = None,
+) -> FaceQueryLookupResult:
+    """Compile a face file and extract the SQL + source for one named SQL query.
+
+    Used by validate and describe paths that need the raw SQL text without
+    executing it against a warehouse.
+    """
+    try:
+        file_path = resolve_scoped_path(path, project_dir)
+    except ValueError as e:
+        return FaceQueryLookupResult(success=False, errors=[str(e)])
+
+    if not file_path.exists():
+        return FaceQueryLookupResult(
+            success=False, errors=[f"File not found: {file_path}"]
+        )
+
+    compile_result = compile_file(file_path)
+    if not compile_result.success:
+        return FaceQueryLookupResult(
+            success=False, errors=[e.message for e in compile_result.errors]
+        )
+
+    face = compile_result.face
+    assert face is not None  # compile_result.success guarantees this
+
+    if name not in face.queries:
+        return FaceQueryLookupResult(
+            success=False,
+            errors=[f"Unknown query: '{name}'"],
+            available_queries=sorted(face.queries.keys()),
+        )
+
+    query = face.queries[name]
+    if not isinstance(query, SqlQuery):
+        return FaceQueryLookupResult(
+            success=False,
+            errors=[
+                f"Query '{name}' is not a SQL query (type: {type(query).__name__})."
+            ],
+        )
+
+    source: str | None = query.source if isinstance(query.source, str) else None
+    return FaceQueryLookupResult(success=True, sql=query.sql, source=source)
+
+
+class ExecuteQueryArgs(BaseModel):
+    """Execute a SQL query against the project's data sources and return results. Use {{ variable_name }} for parameterized values — these work identically in dashboard YAML, so queries you test here will be cached when reused in dashboards. Great for exploring data shape before writing chart configs."""
+
+    sql: str = Field(
+        ...,
+        description="SQL query to execute. Use {{ variable_name }} for parameterized values — these work identically in dashboard YAML, so queries you test here will be cached when reused in dashboards.",
+    )
+    variables: dict[str, Any] | None = Field(
+        None, description='Variable values to substitute (e.g. {"region": "US"})'
+    )
+    source: str | None = Field(None, description="Data source name to execute against")
+    limit: int | None = Field(None, description="Maximum rows to return (default 50)")
+
+
+class ExecuteQueryResult(BaseModel):
+    success: bool
+    columns: list[str]
+    data: list[dict[str, Any]]
+    error: str | None
+    errors: list[str]
+    row_count: int
+    truncated: bool
+
+
+def execute_query(
+    sql: str,
+    variables: dict[str, Any] | None = None,
+    source: str | None = None,
+    limit: int = 50,
+    *,
+    adapter_registry: AdapterRegistry,
+) -> ExecuteQueryResult:
+    """Execute a SQL query and return the results.
+
+    Runs SQL directly against the configured data sources. Use {{ variable_name }}
+    syntax for parameterized values — these work identically in dashboard YAML,
+    so queries tested here will be cached when reused in dashboards.
+    """
+    limit = min(limit, MAX_QUERY_LIMIT)
+    fetch_limit = limit + 1
+
+    try:
+        query = SqlQuery(
+            sql=sql,
+            source=source,
+            limit=fetch_limit,
+        )
+
+        result = adapter_registry.execute(query, variables=variables)
+
+        if result.error:
+            return ExecuteQueryResult(
+                success=False,
+                data=[],
+                columns=[],
+                error=result.error,
+                errors=[result.error],
+                row_count=0,
+                truncated=False,
+            )
+
+        data = normalize_data_for_json(result.data)
+        columns = list(data[0].keys()) if data else []
+        truncated = len(data) > limit
+        data = data[:limit]
+
+        return ExecuteQueryResult(
+            success=True,
+            data=data,
+            columns=columns,
+            error=None,
+            errors=[],
+            row_count=len(data),
+            truncated=truncated,
+        )
+
+    except Exception as e:  # noqa: BLE001
+        return ExecuteQueryResult(
+            success=False,
+            data=[],
+            columns=[],
+            error=str(e),
+            errors=[str(e)],
+            row_count=0,
+            truncated=False,
+        )
+
+
+class QueryFaceArgs(BaseModel):
+    """Run one named query from a face YAML file and return its columns and sample rows. Unlike execute_query (which takes raw SQL), query_face runs a query by name from the face the user actually authored — including Jinja variable substitution and support for non-SQL query types (csv, dbt_model, metricflow). Use this to inspect what a named query produces when debugging chart errors like 'unknown column foo'."""
+
+    name: str = Field(
+        ..., description="Name of the query inside the face's `queries:` block"
+    )
+    path: Path = Field(..., description="Path to the face YAML file")
+    project_dir: Path | None = Field(
+        None, description="Project root for resolving relative paths"
+    )
+    vars: dict[str, Any] | None = Field(
+        None, description='Variable overrides (e.g. {"country": "FR"})'
+    )
+    limit: int = Field(
+        20, ge=1, le=1000, description="Max rows to return (default 20, max 1000)"
+    )
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class QueryFaceResult(BaseModel):
+    success: bool
+    name: str
+    path: Path
+    query_type: str | None = None
+    sql: str | None = None
+    columns: list[str] = []
+    data: list[dict[str, Any]] = []
+    row_count: int = 0
+    truncated: bool = False
+    errors: list[str] = []
+    available_queries: list[str] = []
+
+
+def _fail(
+    name: str, path: Path, errors: list[str], available: list[str] | None = None
+) -> QueryFaceResult:
+    return QueryFaceResult(
+        success=False,
+        name=name,
+        path=path,
+        errors=errors,
+        available_queries=available or [],
+    )
+
+
+def query_face(
+    name: str,
+    path: Path,
+    project_dir: Path | None = None,
+    vars: dict[str, Any] | None = None,
+    limit: int = 20,
+    *,
+    adapter_registry: AdapterRegistry,
+) -> QueryFaceResult:
+    """Run one named query from a compiled face and return its sample rows."""
+    limit = min(limit, MAX_QUERY_LIMIT)
+
+    try:
+        file_path = resolve_scoped_path(path, project_dir)
+    except ValueError as e:
+        return _fail(name, path, [str(e)])
+
+    if not file_path.exists():
+        from dataface.agent_api._paths import no_project_hint
+
+        hint = no_project_hint(project_dir)
+        return _fail(name, file_path, [f"File not found: {file_path}{hint}"])
+
+    compile_result = compile_file(file_path)
+    if not compile_result.success:
+        return _fail(name, file_path, [e.message for e in compile_result.errors])
+
+    face = compile_result.face
+    assert face is not None  # compile_result.success guarantees this
+    if name not in face.queries:
+        return _fail(
+            name,
+            file_path,
+            [f"Unknown query: '{name}'"],
+            sorted(face.queries.keys()),
+        )
+
+    merged_vars = {**face.variable_defaults, **(vars or {})}
+    query = face.queries[name].model_copy(update={"limit": limit + 1})
+
+    try:
+        exec_result = adapter_registry.execute(
+            query, variables=merged_vars or None, face=face
+        )
+    except Exception as e:  # noqa: BLE001
+        return _fail(name, file_path, [str(e)])
+
+    if exec_result.error:
+        return _fail(name, file_path, [exec_result.error])
+
+    data = normalize_data_for_json(exec_result.data)
+    truncated = len(data) > limit
+    data = data[:limit]
+
+    return QueryFaceResult(
+        success=True,
+        name=name,
+        path=file_path,
+        query_type=query.query_type,
+        sql=query.sql if isinstance(query, SqlQuery) else None,
+        columns=list(data[0].keys()) if data else [],
+        data=data,
+        row_count=len(data),
+        truncated=truncated,
+    )