Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
100 changes: 77 additions & 23 deletions src/sqlbuild/compiler/planner/_helpers/fixtures/completion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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,
)
Expand Down Expand Up @@ -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:
Expand All @@ -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),
Expand Down Expand Up @@ -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
Expand All @@ -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, ...]],
Expand Down
8 changes: 8 additions & 0 deletions src/sqlbuild/compiler/planner/_helpers/output/plan_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -49,6 +52,7 @@
PlanOutputExtras,
PlanProviderUsage,
PlanWarning,
RelationFixturePlanningContext,
SeedPlanEntry,
SourceLoadPlanEntry,
SqlTestPlanEntry,
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down
3 changes: 3 additions & 0 deletions src/sqlbuild/compiler/planner/_helpers/sql_tests/assembly.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from sqlbuild.compiler.planner.models import (
ChainStep,
PlanWarning,
RelationFixturePlanningContext,
SqlAnalysisResolvedTestSql,
SqlTestAssertionStep,
SqlTestPlanEntry,
Expand Down Expand Up @@ -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."""

Expand Down Expand Up @@ -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] = []
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from sqlbuild.compiler.planner.models import (
RelationFixtureCompletion,
RelationFixtureDiagnostic,
RelationFixturePlanningContext,
)

_COLLECTION_TYPES: frozenset[str] = frozenset(
Expand All @@ -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."""

Expand All @@ -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)
Expand Down
10 changes: 10 additions & 0 deletions src/sqlbuild/compiler/planner/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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"])
Loading
Loading