Skip to content

Commit d67d4fd

Browse files
kvlongekvlonge
committed
fix: preserve CTE facts with upstream Polyglot
1 parent d0a9296 commit d67d4fd

4 files changed

Lines changed: 76 additions & 47 deletions

File tree

‎src/sqlbuild/compiler/compile/_helpers/analysis/columns.py‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@
1010
from sqlbuild.adapter.contract.models import ExpressionInferenceProfile
1111
from sqlbuild.adapter.contract.types import FunctionNullabilityRule
1212
from sqlbuild.compiler.compile._helpers.analysis.cte_facts import (
13+
_polyglot_cte_passthrough_facts,
1314
_polyglot_cte_passthrough_nullability_from_parsed,
14-
_polyglot_cte_passthrough_type_facts,
1515
_polyglot_cte_passthrough_types_from_parsed,
1616
)
1717
from sqlbuild.compiler.compile.constants import (
@@ -27,6 +27,7 @@
2727
CompiledLineageColumnFact,
2828
CompiledLineageSourceFact,
2929
CompileSqlReference,
30+
CteFactResolvers,
3031
InferredColumn,
3132
PolyglotAnalysisResult,
3233
)
@@ -482,18 +483,24 @@ def _analyze_columns_and_lineage_with_compact_polyglot(
482483
return None
483484
reference_map: dict[str, tuple[CompiledResourceType, str]] = _lineage_reference_map(references)
484485
relation_alias_by_name: dict[str, str | None] = _compact_relation_alias_by_name(analysis)
485-
cte_passthrough_types, direct_cte_outputs = (
486-
_polyglot_cte_passthrough_type_facts(
486+
cte_passthrough_types, cte_passthrough_nullability, direct_cte_outputs = (
487+
_polyglot_cte_passthrough_facts(
487488
polyglot_module=polyglot_module,
488489
cleaned_sql=cleaned_sql,
489490
dialect=dialect,
490491
column_types_by_table=column_types_by_table,
492+
column_nullability_by_table=column_nullability_by_table,
491493
inference_profile=inference_profile,
492494
analysis=analysis,
493-
expression_type_resolver=_polyglot_expression_type,
495+
resolvers=CteFactResolvers(
496+
expression_type=_polyglot_expression_type,
497+
nullability=_infer_polyglot_nullability,
498+
shallow_nullability=_infer_polyglot_shallow_nullability,
499+
alias_nullability=_polyglot_alias_nullability_from_select,
500+
),
494501
)
495502
if recover_cte_facts
496-
else ({}, frozenset())
503+
else ({}, {}, frozenset())
497504
)
498505
columns: list[InferredColumn] = []
499506
lineage_columns: list[CompiledLineageColumnFact] = []
@@ -519,9 +526,13 @@ def _analyze_columns_and_lineage_with_compact_polyglot(
519526
projection=projection, inference_profile=inference_profile
520527
)
521528
),
522-
nullability=_compact_projection_nullability(
523-
projection=projection,
524-
infer_nullability=infer_nullability,
529+
nullability=(
530+
cte_passthrough_nullability.get(output_column)
531+
if output_column in direct_cte_outputs
532+
else None
533+
)
534+
or _compact_projection_nullability(
535+
projection=projection, infer_nullability=infer_nullability
525536
),
526537
)
527538
)

‎src/sqlbuild/compiler/compile/_helpers/analysis/cte_facts.py‎

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from sqlbuild.adapter.type_system.main.normalize_type import normalize_type
1010
from sqlbuild.adapter.type_system.main.types_equal import types_equal
1111
from sqlbuild.compiler.compile.constants import SQL_WILDCARD_TOKEN
12+
from sqlbuild.compiler.compile.models import CteFactResolvers
1213
from sqlbuild.compiler.lineage.types import InferredNullability
1314
from sqlbuild.compiler.sql_analysis.constants import (
1415
POLYGLOT_BINARY_OPERAND_COUNT as _POLYGLOT_BINARY_OPERAND_COUNT,
@@ -97,36 +98,57 @@ def __call__(
9798
) -> dict[str, InferredNullability]: ...
9899

99100

100-
def _polyglot_cte_passthrough_type_facts(
101+
def _polyglot_cte_passthrough_facts(
101102
*,
102103
polyglot_module: Any,
103104
cleaned_sql: str,
104105
dialect: str | None,
105106
column_types_by_table: dict[str, dict[str, str]],
107+
column_nullability_by_table: dict[str, dict[str, InferredNullability]],
106108
inference_profile: ExpressionInferenceProfile,
107109
analysis: dict[str, Any],
108-
expression_type_resolver: _ExpressionTypeResolver,
109-
) -> tuple[dict[str, str], frozenset[str]]:
110-
"""Recover types that compact analysis incorrectly skips across direct CTE reads."""
110+
resolvers: CteFactResolvers,
111+
) -> tuple[dict[str, str], dict[str, InferredNullability], frozenset[str]]:
112+
"""Recover facts that compact analysis skips across direct CTE reads."""
111113

112-
if not column_types_by_table or not analysis.get("cteFacts"):
113-
return {}, frozenset()
114+
if not analysis.get("cteFacts") or not (column_types_by_table or column_nullability_by_table):
115+
return {}, {}, frozenset()
114116
try:
115117
parsed: Any = polyglot_module.parse_one(cleaned_sql, dialect=dialect or "generic")
116118
except polyglot_module.PolyglotError:
117-
return {}, frozenset()
119+
return {}, {}, frozenset()
118120
ctes: tuple[tuple[str, Any, bool], ...] = _polyglot_top_level_ctes(parsed)
119121
if not ctes:
120-
return {}, frozenset()
121-
inferred_types: dict[str, str] = _polyglot_cte_passthrough_types_from_parsed(
122-
parsed=parsed,
123-
column_types_by_table=column_types_by_table,
124-
inference_profile=inference_profile,
125-
expression_type_resolver=expression_type_resolver,
122+
return {}, {}, frozenset()
123+
inferred_types: dict[str, str] = (
124+
_polyglot_cte_passthrough_types_from_parsed(
125+
parsed=parsed,
126+
column_types_by_table=column_types_by_table,
127+
inference_profile=inference_profile,
128+
expression_type_resolver=resolvers.expression_type,
129+
)
130+
if column_types_by_table
131+
else {}
126132
)
127-
return inferred_types, _polyglot_direct_cte_output_names(
128-
parsed=parsed,
129-
cte_names=frozenset(name for name, _, _ in ctes),
133+
inferred_nullability: dict[str, InferredNullability] = (
134+
_polyglot_cte_passthrough_nullability_from_parsed(
135+
parsed=parsed,
136+
column_nullability_by_table=column_nullability_by_table,
137+
inference_profile=inference_profile,
138+
nullability_resolver=resolvers.nullability,
139+
shallow_nullability_resolver=resolvers.shallow_nullability,
140+
alias_nullability_resolver=resolvers.alias_nullability,
141+
)
142+
if column_nullability_by_table
143+
else {}
144+
)
145+
return (
146+
inferred_types,
147+
inferred_nullability,
148+
_polyglot_direct_cte_output_names(
149+
parsed=parsed,
150+
cte_names=frozenset(name for name, _, _ in ctes),
151+
),
130152
)
131153

132154

‎src/sqlbuild/compiler/compile/models.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from collections.abc import Callable, Mapping
66
from dataclasses import dataclass, field
77
from pathlib import Path
8+
from typing import Any
89

910
from sqlbuild.compiler.auditing.models import MeasurementContract, MeasurementThresholds
1011
from sqlbuild.compiler.auditing.types import AuditEvaluationMode, AuditSeverity
@@ -143,6 +144,16 @@ class InferredColumn:
143144
nullability: InferredNullability = InferredNullability.UNKNOWN
144145

145146

147+
@dataclass(frozen=True)
148+
class CteFactResolvers:
149+
"""Expression resolvers used for conservative CTE fact recovery."""
150+
151+
expression_type: Callable[..., Any]
152+
nullability: Callable[..., Any]
153+
shallow_nullability: Callable[..., Any]
154+
alias_nullability: Callable[..., Any]
155+
156+
146157
@dataclass(frozen=True)
147158
class CompileModelConfig:
148159
"""Pre-semantic effective config layers attached to a model input."""

‎tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py‎

Lines changed: 8 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -900,19 +900,8 @@ def raise_parse_error(*args: object, **kwargs: object) -> object:
900900
InferredColumn(name="status", type="TEXT"),
901901
InferredColumn(name="nullable_status", type="TEXT"),
902902
),
903-
expected_lineage_columns=(
904-
CompiledLineageColumnFact(
905-
output_column="amount",
906-
upstream_columns=(),
907-
transform_kind=ColumnTransformKind.CONSTANT,
908-
confidence=ColumnLineageConfidence.HIGH,
909-
),
910-
CompiledLineageColumnFact(
911-
output_column="status",
912-
upstream_columns=(),
913-
transform_kind=ColumnTransformKind.CONSTANT,
914-
confidence=ColumnLineageConfidence.HIGH,
915-
),
903+
expected_lineage_columns=direct_orders_lineage("amount", "status")
904+
+ (
916905
CompiledLineageColumnFact(
917906
output_column="nullable_status",
918907
upstream_columns=(
@@ -986,7 +975,7 @@ def raise_parse_error(*args: object, **kwargs: object) -> object:
986975
expected_has_star=False,
987976
),
988977
PolyglotAnalysisTestCase(
989-
description="declines renamed star type recovery",
978+
description="preserves renamed star lineage without recovering type",
990979
query_sql=(
991980
"WITH typed AS ("
992981
"SELECT CAST(amount AS INTEGER) AS amount, "
@@ -1002,14 +991,7 @@ def raise_parse_error(*args: object, **kwargs: object) -> object:
1002991
column_types_by_table={"orders": {"amount": "VARCHAR"}},
1003992
inference_profile=ExpressionInferenceProfile(sql_analysis_dialect="snowflake"),
1004993
expected_columns=(InferredColumn(name="amount"),),
1005-
expected_lineage_columns=(
1006-
CompiledLineageColumnFact(
1007-
output_column="amount",
1008-
upstream_columns=(),
1009-
transform_kind=ColumnTransformKind.CONSTANT,
1010-
confidence=ColumnLineageConfidence.HIGH,
1011-
),
1012-
),
994+
expected_lineage_columns=direct_orders_lineage("amount"),
1013995
expected_has_star=False,
1014996
),
1015997
],
@@ -1093,7 +1075,10 @@ def test_given_cte_chain_when_analyzing_direct_passthrough_then_type_is_conserva
10931075
function_return_types={"TO_DATE": "DATE"},
10941076
),
10951077
expected_columns=(
1096-
InferredColumn(name="source_name"),
1078+
InferredColumn(
1079+
name="source_name",
1080+
nullability=InferredNullability.NON_NULL,
1081+
),
10971082
InferredColumn(name="active", type="BOOLEAN"),
10981083
InferredColumn(name="selected", type="BOOLEAN"),
10991084
InferredColumn(name="full_name", type="TEXT"),

0 commit comments

Comments
 (0)