fix: reduce large-project planning work

Author: kvlonge

src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py

@@ -95,7 +95,7 @@
 from sqlbuild.sql_values.models import SqlValue
 
 _EXACT_COLUMN_INSPECTION_LIMIT: int = 32
-_BULK_COLUMN_RELATION_CHUNK_SIZE: int = 200
+_BULK_COLUMN_RELATION_CHUNK_SIZE: int = 2_000
 
 
 class SnowflakeAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter):

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

@@ -31,6 +31,7 @@
     FixtureRelationMetadata,
     RelationFixtureCompletion,
     RelationFixtureDiagnostic,
+    RelationFixturePlanningContext,
 )
 from sqlbuild.compiler.planner.types import ContractPolicy, FixtureGroups, FixtureKey
 from sqlbuild.compiler.references.constants import REF_PATTERN, SEED_PATTERN, SOURCE_PATTERN
@@ -47,10 +48,14 @@ def build_relation_fixture_completion(
     adapter: BaseAdapter,
     ordered_model_names: tuple[str, ...],
     fixture_groups: FixtureGroups,
+    planning_context: RelationFixturePlanningContext | None = None,
 ) -> RelationFixtureCompletion:
     """Complete required omitted columns with typed nulls unless explicitly non-nullable."""
 
-    model_map: dict[str, CompiledModel] = {model.name: model for model in project.models}
+    context: RelationFixturePlanningContext = planning_context or build_relation_fixture_context(
+        project=project
+    )
+    model_map: dict[str, CompiledModel] = context.models_by_name
     fixture_sql_by_key: dict[FixtureKey, str] = {}
     inferred_by_fixture: dict[FixtureKey, tuple[InferredColumn, ...]] = {}
     for resource_type, fixtures in fixture_groups:
@@ -64,21 +69,23 @@ def build_relation_fixture_completion(
             if inferred is not None:
                 inferred_by_fixture[key] = inferred
 
-    relations: dict[FixtureKey, FixtureRelationMetadata] = _fixture_relation_metadata(
-        project=project,
-        model_map=model_map,
+    relations: dict[FixtureKey, FixtureRelationMetadata] = context.relations
+    fixture_keys_requiring_analysis: frozenset[FixtureKey] = _fixture_keys_requiring_analysis(
+        inferred_by_fixture=inferred_by_fixture,
+        authoritative_columns=context.authoritative_columns,
     )
     required: dict[FixtureKey, dict[str, set[str]]]
     star_fixture_keys: frozenset[FixtureKey]
     required, star_fixture_keys = _required_fixture_columns(
         ordered_model_names=ordered_model_names,
         model_map=model_map,
-        fixture_keys=frozenset(fixture_sql_by_key),
+        fixture_keys=fixture_keys_requiring_analysis,
         relations=relations,
+        authoritative_columns=context.authoritative_columns,
         adapter=adapter,
     )
     expected_types: dict[FixtureKey, dict[str, str]] = _expected_fixture_types(
-        relations=relations,
+        base_expected_types=context.expected_types,
         ordered_model_names=ordered_model_names,
         model_map=model_map,
     )
@@ -177,23 +184,61 @@ def build_relation_fixture_completion(
     )
 
 
+def _fixture_keys_requiring_analysis(
+    *,
+    inferred_by_fixture: dict[FixtureKey, tuple[InferredColumn, ...]],
+    authoritative_columns: dict[FixtureKey, frozenset[str]],
+) -> frozenset[FixtureKey]:
+    required: set[FixtureKey] = set()
+    for key, inferred in inferred_by_fixture.items():
+        expected_names: frozenset[str] | None = authoritative_columns.get(key)
+        supplied_names: frozenset[str] = frozenset(column.name.casefold() for column in inferred)
+        if expected_names is None or not expected_names.issubset(supplied_names):
+            required.add(key)
+    return frozenset(required)
+
+
+def build_relation_fixture_context(*, project: CompiledProject) -> RelationFixturePlanningContext:
+    """Build immutable project-wide fixture metadata once per planning invocation."""
+
+    models_by_name: dict[str, CompiledModel] = {model.name: model for model in project.models}
+    relations: dict[FixtureKey, FixtureRelationMetadata] = _fixture_relation_metadata(
+        project=project,
+        model_map=models_by_name,
+    )
+    authoritative_columns: dict[FixtureKey, frozenset[str]] = _authoritative_fixture_columns(
+        relations=relations
+    )
+    return RelationFixturePlanningContext(
+        models_by_name=models_by_name,
+        relations=relations,
+        authoritative_columns=authoritative_columns,
+        expected_types=_base_expected_fixture_types(relations=relations),
+    )
+
+
+def _authoritative_fixture_columns(
+    *, relations: dict[FixtureKey, FixtureRelationMetadata]
+) -> dict[FixtureKey, frozenset[str]]:
+    authoritative: dict[FixtureKey, frozenset[str]] = {}
+    for key, relation in relations.items():
+        if relation.authoritative_names:
+            authoritative[key] = frozenset(column.name.casefold() for column in relation.columns)
+    return authoritative
+
+
 def _required_fixture_columns(
     *,
     ordered_model_names: tuple[str, ...],
     model_map: dict[str, CompiledModel],
     fixture_keys: frozenset[FixtureKey],
     relations: dict[FixtureKey, FixtureRelationMetadata],
+    authoritative_columns: dict[FixtureKey, frozenset[str]],
     adapter: BaseAdapter,
 ) -> tuple[dict[FixtureKey, dict[str, set[str]]], frozenset[FixtureKey]]:
     required: dict[FixtureKey, dict[str, set[str]]] = {}
     star_fixture_keys: set[FixtureKey] = set()
-    authoritative_columns: dict[FixtureKey, frozenset[str]] = {}
-    for key, relation in relations.items():
-        if not relation.authoritative_names:
-            continue
-        authoritative_columns[key] = frozenset(
-            column.name.casefold() for column in relation.columns
-        )
+    dialect: str | None = adapter.expression_inference_profile().sql_analysis_dialect
     for model_name in ordered_model_names:
         model: CompiledModel | None = model_map.get(model_name)
         if model is None:
@@ -203,7 +248,7 @@ def _required_fixture_columns(
             references=model.references,
             fast_lineage=model.fast_lineage_columns or (),
             fixture_keys=fixture_keys,
-            dialect=adapter.expression_inference_profile().sql_analysis_dialect,
+            dialect=dialect,
         ):
             key: FixtureKey = (
                 CompiledResourceType(upstream.resource_type),
@@ -451,22 +496,18 @@ def _model_fixture_metadata(*, model: CompiledModel) -> FixtureRelationMetadata:
 
 def _expected_fixture_types(
     *,
-    relations: dict[FixtureKey, FixtureRelationMetadata],
+    base_expected_types: dict[FixtureKey, dict[str, str]],
     ordered_model_names: tuple[str, ...],
     model_map: dict[str, CompiledModel],
 ) -> dict[FixtureKey, dict[str, str]]:
-    expected_types: dict[FixtureKey, dict[str, str]] = {}
-    for key, relation in relations.items():
-        column_types: dict[str, str] = {}
-        for column in relation.columns:
-            if column.type is not None:
-                column_types[column.name.casefold()] = column.type
-        expected_types[key] = column_types
+    expected_types: dict[FixtureKey, dict[str, str]] = dict(base_expected_types)
     for model_name in ordered_model_names:
         model: CompiledModel | None = model_map.get(model_name)
         if model is None:
             continue
-        model_types: dict[str, str] = expected_types[(CompiledResourceType.MODEL, model_name)]
+        model_key: FixtureKey = (CompiledResourceType.MODEL, model_name)
+        model_types: dict[str, str] = dict(expected_types[model_key])
+        expected_types[model_key] = model_types
         for lineage_column in model.fast_lineage_columns or ():
             if lineage_column.transform_kind != ColumnTransformKind.DIRECT:
                 continue
@@ -484,6 +525,19 @@ def _expected_fixture_types(
     return expected_types
 
 
+def _base_expected_fixture_types(
+    *, relations: dict[FixtureKey, FixtureRelationMetadata]
+) -> dict[FixtureKey, dict[str, str]]:
+    expected_types: dict[FixtureKey, dict[str, str]] = {}
+    for key, relation in relations.items():
+        column_types: dict[str, str] = {}
+        for column in relation.columns:
+            if column.type is not None:
+                column_types[column.name.casefold()] = column.type
+        expected_types[key] = column_types
+    return expected_types
+
+
 def _unknown_fixture_column_diagnostics(
     *,
     inferred_by_fixture: dict[FixtureKey, tuple[InferredColumn, ...]],

src/sqlbuild/compiler/planner/_helpers/output/plan_output.py

@@ -20,6 +20,9 @@
     PythonHookEntry,
 )
 from sqlbuild.compiler.fingerprints.models import Fingerprint
+from sqlbuild.compiler.planner._helpers.fixtures.completion import (
+    build_relation_fixture_context,
+)
 from sqlbuild.compiler.planner._helpers.graph.loader_dag import upstream_loader_dependency_names
 from sqlbuild.compiler.planner._helpers.graph.source_load_nodes import build_source_load_entries
 from sqlbuild.compiler.planner._helpers.output.audit_entry import plan_audit
@@ -49,6 +52,7 @@
     PlanOutputExtras,
     PlanProviderUsage,
     PlanWarning,
+    RelationFixturePlanningContext,
     SeedPlanEntry,
     SourceLoadPlanEntry,
     SqlTestPlanEntry,
@@ -441,12 +445,15 @@ def build_selected_test_entries(
     entries: list[SqlTestPlanEntry] = []
     warnings: list[PlanWarning] = []
     fixture_diagnostics: list[str] = []
+    fixture_planning_context: RelationFixturePlanningContext | None = None
     sql_test: CompiledSqlTest
     for sql_test in project.sql_tests:
         if not scope_overlaps(scope_deps=sql_test.scope_deps, selected_keys=selected_keys):
             continue
         if case_name is not None and sql_test.case_name != case_name:
             continue
+        if project.settings.sql_analysis and fixture_planning_context is None:
+            fixture_planning_context = build_relation_fixture_context(project=project)
         test_entry: SqlTestPlanEntry
         test_warnings: tuple[PlanWarning, ...]
         try:
@@ -456,6 +463,7 @@ def build_selected_test_entries(
                 adapter=adapter,
                 sql_analysis_enabled=project.settings.sql_analysis,
                 validate_fixtures=True,
+                fixture_planning_context=fixture_planning_context,
             )
         except SqlTestFixtureValidationError as error:
             fixture_diagnostics.extend(error.diagnostics)

src/sqlbuild/compiler/planner/_helpers/sql_tests/assembly.py

@@ -46,6 +46,7 @@
 from sqlbuild.compiler.planner.models import (
     ChainStep,
     PlanWarning,
+    RelationFixturePlanningContext,
     SqlAnalysisResolvedTestSql,
     SqlTestAssertionStep,
     SqlTestPlanEntry,
@@ -133,6 +134,7 @@ def plan_test(
     adapter: BaseAdapter,
     sql_analysis_enabled: bool = False,
     validate_fixtures: bool = False,
+    fixture_planning_context: RelationFixturePlanningContext | None = None,
 ) -> tuple[SqlTestPlanEntry, tuple[PlanWarning, ...]]:
     """Build a test plan entry with chained resolution."""
 
@@ -193,6 +195,7 @@ def plan_test(
             mock_sources=mock_sources,
             mock_seeds=mock_seeds,
             expected_outputs=expected_map,
+            planning_context=fixture_planning_context,
         )
 
     warnings: list[PlanWarning] = []

src/sqlbuild/compiler/planner/_helpers/sql_tests/fixture_validation.py

@@ -19,6 +19,7 @@
 from sqlbuild.compiler.planner.models import (
     RelationFixtureCompletion,
     RelationFixtureDiagnostic,
+    RelationFixturePlanningContext,
 )
 
 _COLLECTION_TYPES: frozenset[str] = frozenset(
@@ -36,6 +37,7 @@ def build_validated_test_fixtures(
     mock_sources: dict[str, str],
     mock_seeds: dict[str, str],
     expected_outputs: dict[str, str],
+    planning_context: RelationFixturePlanningContext | None = None,
 ) -> tuple[dict[str, str], dict[str, str], dict[str, str]]:
     """Report all statically provable missing columns and collection-type mismatches."""
 
@@ -49,6 +51,7 @@ def build_validated_test_fixtures(
         adapter=adapter,
         ordered_model_names=ordered_model_names,
         fixture_groups=fixture_groups,
+        planning_context=planning_context,
     )
     errors: list[str] = [
         _located_diagnostic(test=test, diagnostic=diagnostic)

src/sqlbuild/compiler/planner/models.py

@@ -122,6 +122,16 @@ class FixtureRelationMetadata:
     authoritative_names: bool
 
 
+@dataclass(frozen=True)
+class RelationFixturePlanningContext:
+    """Project-wide fixture metadata shared across test and scenario planning."""
+
+    models_by_name: dict[str, CompiledModel]
+    relations: dict[FixtureKey, FixtureRelationMetadata]
+    authoritative_columns: dict[FixtureKey, frozenset[str]]
+    expected_types: dict[FixtureKey, dict[str, str]]
+
+
 @dataclass(frozen=True)
 class GraphNodeKey:
     """Neutral graph key matching fingerprint identity fields."""

tests/integration/src/sqlbuild/compiler/planner/_helpers/_test_types.py

@@ -7,6 +7,12 @@
 from sqlbuild.spec.contracts.models import SourceEntry
 
 
+@dataclass(frozen=True)
+class FixturePlanningIntegrationTestCase:
+    description: str
+    expected_sql_fragment: str
+
+
 @dataclass(frozen=True)
 class GatherWatermarkTypeTestCase:
     description: str

tests/integration/src/sqlbuild/compiler/planner/_helpers/test_fixture_planning.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+
+from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter
+from sqlbuild.compiler.compile.models import CompiledProject
+from sqlbuild.compiler.discovery.main.discover import discover_project_inputs
+from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs
+from sqlbuild.compiler.pipeline.main.project import compile_project
+from sqlbuild.compiler.planner._helpers.output.plan_output import build_selected_test_entries
+from sqlbuild.compiler.planner.models import SqlTestPlanEntry
+from tests.integration.src.sqlbuild.compiler.planner._helpers._test_types import (
+    FixturePlanningIntegrationTestCase,
+)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        FixturePlanningIntegrationTestCase(
+            description="partial source fixture",
+            expected_sql_fragment='CAST(NULL AS VARCHAR) AS "status"',
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_partial_fixture_when_compiling_and_planning_then_completes_required_columns(
+    test_case: FixturePlanningIntegrationTestCase,
+    tmp_path: Path,
+    write_repo_files: Callable[[Path, dict[str, str]], None],
+) -> None:
+    write_repo_files(
+        tmp_path,
+        {
+            "sqlbuild_project.toml": (
+                'name = "fixture_cache_demo"\n'
+                'adapter = "duckdb"\n\n'
+                "[connection]\n"
+                'database = ":memory:"\n'
+            ),
+            "models/orders.sql": (
+                'MODEL ();\n\nSELECT order_id, status FROM __source("raw_orders")\n'
+            ),
+            "sources/raw_orders.yml": (
+                "sources:\n"
+                "  - name: raw_orders\n"
+                "    expression: SELECT 1 AS order_id, 'open' AS status\n"
+                "    columns:\n"
+                "      - name: order_id\n        type: INTEGER\n        nullable: false\n"
+                "      - name: status\n        type: VARCHAR\n        nullable: true\n"
+            ),
+            "tests/unit/test_orders.sql": (
+                "TEST();\n\n"
+                "WITH\n"
+                "__source__raw_orders AS (SELECT 1 AS order_id),\n"
+                "__expected__orders AS (\n"
+                "  SELECT 1 AS order_id, CAST(NULL AS VARCHAR) AS status\n"
+                ")\n"
+                "SELECT 1\n"
+            ),
+        },
+    )
+    discovered_inputs: DiscoveredProjectInputs = discover_project_inputs(project_dir=tmp_path)
+    adapter: DuckDbAdapter = DuckDbAdapter()
+
+    project: CompiledProject = compile_project(
+        discovered_inputs=discovered_inputs,
+        adapter=adapter,
+    )
+    entries, warnings = build_selected_test_entries(
+        project=project,
+        adapter=adapter,
+        selected_keys=frozenset((project.models[0].key,)),
+    )
+
+    assert warnings == []
+    assert len(entries) == 1
+    entry: SqlTestPlanEntry = entries[0]
+    assert test_case.expected_sql_fragment in entry.chain[0].resolved_sql
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vv"])