fix: harden SQL analysis for complex projects

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
   "duckdb>=1.3.2",
   "filelock>=3.20.3",
   "orjson>=3.10.0",
-  "polyglot-sql==0.11.0",
+  "polyglot-sql==0.12.0",
   "pydantic-settings>=2,<3",
   "PyYAML>=6.0.2",
   "rich>=15.0.0",

src/sqlbuild/compiler/compile/_helpers/analysis/dynamic_pivot.py

@@ -140,6 +140,7 @@ def _analyze_passthrough(
             DynamicColumnFamilyProof(name=family.name, inferred_type=family.type)
             for family in families
         ),
+        input_relations=(name,),
     )
 
 
@@ -174,6 +175,13 @@ def _analyze_pivot(
     )
     if source_columns is None:
         return _failure("dynamic pivot input does not have an authoritative, explicit schema")
+    input_relation: str | None = _external_relation_name(
+        node=source_node,
+        ctes=ctes,
+        seen=frozenset(),
+    )
+    if input_relation is None:
+        return _failure("dynamic pivot input relation could not be resolved")
     pivot_column: str = pivot_columns[0]
     matched_aggregate_indices: set[int] = set()
     family_proofs: list[DynamicColumnFamilyProof] = []
@@ -238,6 +246,7 @@ def _analyze_pivot(
             for name in fixed_names
         ),
         families=tuple(family_proofs),
+        input_relations=(input_relation,),
         bare_dynamic_pivot=bare_dynamic_pivot,
     )
 
@@ -393,6 +402,46 @@ def _relation_schema(
     )
 
 
+def _external_relation_name(
+    *,
+    node: object,
+    ctes: _CteMap,
+    seen: frozenset[str],
+) -> str | None:
+    key: str = _node_key(node)
+    payload: object = cast(dict[str, object], node).get(key) if isinstance(node, dict) else None
+    if key in _QUERY_AST_KINDS and isinstance(payload, dict):
+        return _external_relation_name(
+            node=cast(dict[str, object], payload).get("this"),
+            ctes=ctes,
+            seen=seen,
+        )
+    if key == _SELECT_AST_KIND and isinstance(payload, dict):
+        payload_dict: dict[str, object] = cast(dict[str, object], payload)
+        if payload_dict.get("joins"):
+            return None
+        relations: tuple[dict[str, object], ...] = _nested_dict_list(
+            node=payload_dict.get("from"), key="expressions"
+        )
+        if len(relations) != 1:
+            return None
+        return _external_relation_name(node=relations[0], ctes=ctes, seen=seen)
+    relation_name: str | None = _relation_name(node)
+    if relation_name is None:
+        return None
+    normalized_name: str = relation_name.casefold()
+    if normalized_name not in ctes:
+        return relation_name
+    cte: dict[str, object] | None = ctes[normalized_name]
+    if cte is None or normalized_name in seen:
+        return None
+    return _external_relation_name(
+        node=cte,
+        ctes=ctes,
+        seen=seen | {normalized_name},
+    )
+
+
 def _select_schema(
     *,
     node: object,

src/sqlbuild/compiler/compile/models.py

@@ -159,6 +159,7 @@ class DynamicColumnContractProof:
     output_proven: bool
     fixed_columns: tuple[InferredColumn, ...] = field(default_factory=tuple)
     families: tuple[DynamicColumnFamilyProof, ...] = field(default_factory=tuple)
+    input_relations: tuple[str, ...] = field(default_factory=tuple)
     failure_reason: str | None = None
     bare_dynamic_pivot: bool = False
 

src/sqlbuild/compiler/discovery/_helpers/filesystem/python_paths.py

@@ -37,7 +37,7 @@
 _LANGUAGE_PYTHON_ROOTS: frozenset[tuple[str, str]] = frozenset(
     {("functions", "python"), ("hooks", "python")}
 )
-_IGNORED_ROOTS: frozenset[str] = frozenset({"logs", "target", "venv"})
+_IGNORED_ROOTS: frozenset[str] = frozenset({"logs", "target", "tests", "venv"})
 _LEGACY_DIAGNOSTIC_ROOTS: frozenset[str] = frozenset({"event_exporters"})
 _SUPPORTED_ROOT_FILES: frozenset[Path] = frozenset({Path("adapter.py"), Path("definitions.py")})
 _SCOPED_DECLARATION_DIRECTORIES: frozenset[str] = (

src/sqlbuild/compiler/planner/_helpers/fixtures/completion.py

@@ -17,6 +17,7 @@
     CompiledModel,
     CompiledProject,
     CompileSqlReference,
+    DynamicColumnContractProof,
     InferredColumn,
 )
 from sqlbuild.compiler.compile.types import CompiledResourceType
@@ -295,6 +296,12 @@ def _star_fixture_keys(
         for dependency in model.deps
         if (CompiledResourceType(dependency.resource_type), dependency.name) in fixture_keys
     )
+    dynamic_contract: DynamicColumnContractProof | None = model.dynamic_column_contract
+    if dynamic_contract is not None and dynamic_contract.output_proven:
+        input_relations: frozenset[str] = frozenset(
+            name.casefold() for name in dynamic_contract.input_relations
+        )
+        return frozenset(key for key in dependency_keys if key[1].casefold() in input_relations)
     schema_tables: list[dict[str, object]] = []
     for key in dependency_keys:
         relation: FixtureRelationMetadata | None = relations.get(key)

src/sqlbuild/compiler/sql_analysis/_helpers/polyglot.py

@@ -6,11 +6,76 @@
 
 import polyglot_sql
 
+from sqlbuild.compiler.sql_analysis.constants import (
+    POLYGLOT_ANALYZE_QUERY_ENTRY_POINT,
+    POLYGLOT_ANALYZE_QUERY_OPTIONS_ARGUMENT_COUNT,
+    POLYGLOT_ANALYZE_QUERY_OPTIONS_ARGUMENT_INDEX,
+    POLYGLOT_COMPLEXITY_GUARD_OPTION,
+    POLYGLOT_COMPLEXITY_GUARD_PARAMETER,
+    POLYGLOT_MAX_FUNCTION_CALL_DEPTH,
+    POLYGLOT_MAX_FUNCTION_CALL_DEPTH_OPTION,
+)
+
+_COMPLEXITY_GUARD: dict[str, int] = {
+    POLYGLOT_MAX_FUNCTION_CALL_DEPTH_OPTION: POLYGLOT_MAX_FUNCTION_CALL_DEPTH,
+}
+_GUARDED_ENTRY_POINTS: frozenset[str] = frozenset(
+    {
+        POLYGLOT_ANALYZE_QUERY_ENTRY_POINT,
+        "parse",
+        "parse_data_type",
+        "parse_one",
+        "transpile",
+        "validate",
+        "validate_with_schema",
+    }
+)
+
+
+class _GuardedEntryPoint:
+    """One Polyglot entry point with SQLBuild's bounded trusted-SQL budget."""
+
+    def __init__(self, *, name: str, function: Any) -> None:
+        self._name: str = name
+        self._function: Any = function
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if self._has_explicit_options_guard(args=args, kwargs=kwargs):
+            return self._function(*args, **kwargs)
+        guarded_kwargs: dict[str, Any] = {
+            POLYGLOT_COMPLEXITY_GUARD_PARAMETER: dict(_COMPLEXITY_GUARD),
+            **kwargs,
+        }
+        return self._function(*args, **guarded_kwargs)
+
+    def _has_explicit_options_guard(self, *, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool:
+        if self._name != POLYGLOT_ANALYZE_QUERY_ENTRY_POINT:
+            return False
+        options: object = kwargs.get("options")
+        if len(args) >= POLYGLOT_ANALYZE_QUERY_OPTIONS_ARGUMENT_COUNT:
+            options = args[POLYGLOT_ANALYZE_QUERY_OPTIONS_ARGUMENT_INDEX]
+        return isinstance(options, dict) and POLYGLOT_COMPLEXITY_GUARD_OPTION in options
+
+
+class _PolyglotSqlModule:
+    """Proxy Polyglot parsing through SQLBuild's trusted-SQL policy."""
+
+    def __getattr__(self, name: str) -> Any:
+        attribute: Any = getattr(polyglot_sql, name)
+        return (
+            _GuardedEntryPoint(name=name, function=attribute)
+            if name in _GUARDED_ENTRY_POINTS
+            else attribute
+        )
+
+
+_POLYGLOT_SQL_MODULE: _PolyglotSqlModule = _PolyglotSqlModule()
+
 
 def import_polyglot() -> Any:
     """Return SQLBuild's required Polyglot SQL module."""
 
-    return polyglot_sql
+    return _POLYGLOT_SQL_MODULE
 
 
 def import_polyglot_sql() -> Any:

src/sqlbuild/compiler/sql_analysis/constants.py

@@ -14,6 +14,13 @@
 SQL_IDENTIFIER_PREFIX: str = "_"
 SQL_NULL_KEYWORD: str = "NULL"
 SNOWFLAKE_DIALECT_NAME: str = "snowflake"
+POLYGLOT_ANALYZE_QUERY_ENTRY_POINT: str = "analyze_query"
+POLYGLOT_ANALYZE_QUERY_OPTIONS_ARGUMENT_COUNT: int = 2
+POLYGLOT_ANALYZE_QUERY_OPTIONS_ARGUMENT_INDEX: int = 1
+POLYGLOT_COMPLEXITY_GUARD_PARAMETER: str = "complexity_guard"
+POLYGLOT_COMPLEXITY_GUARD_OPTION: str = "complexityGuard"
+POLYGLOT_MAX_FUNCTION_CALL_DEPTH_OPTION: str = "maxFunctionCallDepth"
+POLYGLOT_MAX_FUNCTION_CALL_DEPTH: int = 512
 BINDING_UNKNOWN_TABLE_INTERNAL_CODE: str = "B000"
 
 POLYGLOT_KIND_ALIAS: str = "alias"

tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py

@@ -100,3 +100,10 @@ class CompileSelectionTestCase:
     selection_args: tuple[str, ...]
     expected_stdout_fragments: tuple[str, ...]
     unexpected_stdout_fragments: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class DeepSqlAnalysisCompileTestCase:
+    description: str
+    function_depth: int
+    expected_stdout_fragment: str

tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py

@@ -40,6 +40,13 @@ class CompileBenchmarkMeasurement(NamedTuple):
     summary: dict[str, int]
 
 
+def nested_coalesce_expression(*, function_depth: int) -> str:
+    expression: str = "1"
+    for _ in range(function_depth):
+        expression = f"COALESCE({expression}, 0)"
+    return expression
+
+
 class LayeredProductionCompileBenchmarkResult(NamedTuple):
     cold: CompileBenchmarkMeasurement
     warm: CompileBenchmarkMeasurement

tests/e2e/src/sqlbuild/cli/commands/main/compile/test_compile_complexity_guard.py

@@ -0,0 +1,66 @@
+"""E2E coverage for SQL analysis complexity budgets."""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import (
+    DeepSqlAnalysisCompileTestCase,
+)
+from tests.e2e.src.sqlbuild.cli.commands.main.compile.helpers import (
+    nested_coalesce_expression,
+)
+from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project, run_sqb
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        DeepSqlAnalysisCompileTestCase(
+            description="trusted model exceeds the defensive default",
+            function_depth=65,
+            expected_stdout_fragment="Project compiled  1 model",
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_deep_trusted_model_when_compiling_then_sql_analysis_completes(
+    test_case: DeepSqlAnalysisCompileTestCase,
+    tmp_path: Path,
+) -> None:
+    project_dir: Path = prepare_inline_project(
+        tmp_path=tmp_path,
+        project_name="deep_sql_analysis",
+        repo_files={
+            "sqlbuild_project.toml": (
+                'name = "deep_sql_analysis"\n'
+                'adapter = "duckdb"\n'
+                'default_target = "dev"\n\n'
+                "[connection]\n"
+                'database = ":memory:"\n\n'
+                "[targets.dev]\n"
+                'schema = "dev"\n'
+            ),
+            "models/deep_expression.sql": (
+                "MODEL (columns (resolved_value (type NUMBER)));\n\n"
+                f"SELECT {nested_coalesce_expression(function_depth=test_case.function_depth)} "
+                "AS resolved_value\n"
+            ),
+        },
+    )
+
+    result: subprocess.CompletedProcess[str] = run_sqb(
+        command=("--no-color", "compile"),
+        project_dir=project_dir,
+        working_dir=project_dir,
+    )
+
+    assert result.returncode == 0, result.stdout + result.stderr
+    assert test_case.expected_stdout_fragment in result.stdout
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vv"])

tests/e2e/src/sqlbuild/cli/commands/main/compile/test_python_project_layout.py

@@ -50,5 +50,40 @@ def test_given_python_under_unsupported_root_when_compiling_then_command_fails_w
     assert all(fragment in result.stderr for fragment in test_case.expected_stderr_fragments)
 
 
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        PythonProjectLayoutCompileTestCase(
+            description="custom Rule harness test path",
+            repo_files={
+                "sqlbuild_project.toml": 'name = "rule_harness_layout"\nadapter = "duckdb"\n',
+                "models/orders.sql": "MODEL ();\nSELECT 1 AS order_id\n",
+                "tests/rules/test_order_policy.py": "def test_order_policy(): pass\n",
+            },
+            expected_exit_code=0,
+            expected_stderr_fragments=(),
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_rule_harness_test_when_compiling_then_python_test_path_is_accepted(
+    test_case: PythonProjectLayoutCompileTestCase,
+    tmp_path: Path,
+) -> None:
+    project_dir: Path = prepare_inline_project(
+        tmp_path=tmp_path,
+        project_name="rule_harness_layout",
+        repo_files=test_case.repo_files,
+    )
+
+    result: subprocess.CompletedProcess[str] = run_sqb(
+        command=("--no-color", "compile"),
+        project_dir=project_dir,
+    )
+
+    assert result.returncode == test_case.expected_exit_code, result.stdout + result.stderr
+    assert "Project compiled  1 model" in result.stdout
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-vv"])