diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/columns.py b/src/sqlbuild/compiler/compile/_helpers/analysis/columns.py index 24f3b84265..0f79b0fa5b 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/columns.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/columns.py @@ -13,6 +13,9 @@ _polyglot_cte_passthrough_facts, _polyglot_cte_passthrough_nullability_from_parsed, _polyglot_cte_passthrough_types_from_parsed, + _polyglot_expression_is_non_null_after_filter, + _polyglot_filtered_non_null_outputs, + _polyglot_non_null_filter_context, ) from sqlbuild.compiler.compile.constants import ( DECIMAL_SQL_TYPE_NAME, @@ -29,6 +32,7 @@ CompileSqlReference, CteFactResolvers, InferredColumn, + NonNullFilterContext, PolyglotAnalysisResult, ) from sqlbuild.compiler.compile.types import CompiledResourceType @@ -483,7 +487,7 @@ def _analyze_columns_and_lineage_with_compact_polyglot( return None reference_map: dict[str, tuple[CompiledResourceType, str]] = _lineage_reference_map(references) relation_alias_by_name: dict[str, str | None] = _compact_relation_alias_by_name(analysis) - cte_passthrough_types, cte_passthrough_nullability, direct_cte_outputs = ( + cte_passthrough_types, cte_passthrough_nullability, direct_cte_outputs, parsed = ( _polyglot_cte_passthrough_facts( polyglot_module=polyglot_module, cleaned_sql=cleaned_sql, @@ -500,7 +504,15 @@ def _analyze_columns_and_lineage_with_compact_polyglot( ), ) if recover_cte_facts - else ({}, {}, frozenset()) + else ({}, {}, frozenset(), None) + ) + filtered_non_null_outputs: frozenset[str] = _polyglot_filtered_non_null_outputs( + polyglot_module=polyglot_module, + cleaned_sql=cleaned_sql, + dialect=dialect, + analysis=analysis, + column_nullability_by_table=column_nullability_by_table, + parsed=parsed, ) columns: list[InferredColumn] = [] lineage_columns: list[CompiledLineageColumnFact] = [] @@ -527,12 +539,19 @@ def _analyze_columns_and_lineage_with_compact_polyglot( ) ), nullability=( - cte_passthrough_nullability.get(output_column) - if output_column in direct_cte_outputs - else None - ) - or _compact_projection_nullability( - projection=projection, infer_nullability=infer_nullability + InferredNullability.NON_NULL + if output_column in filtered_non_null_outputs + else ( + ( + cte_passthrough_nullability.get(output_column) + if output_column in direct_cte_outputs + else None + ) + or _compact_projection_nullability( + projection=projection, + infer_nullability=infer_nullability, + ) + ) ), ) ) @@ -848,6 +867,10 @@ def _infer_columns_from_polyglot_ast( alias_nullability_resolver=_polyglot_alias_nullability_from_select, ) ) + non_null_filter_context: NonNullFilterContext | None = _polyglot_non_null_filter_context( + select=select, + column_nullability_by_table=column_nullability_by_table, + ) columns: list[InferredColumn] = [] projection: Any for projection in getattr(select, "expressions", ()): @@ -868,14 +891,21 @@ def _infer_columns_from_polyglot_ast( ) nullability: InferredNullability = InferredNullability.UNKNOWN if infer_nullability: - nullability = cte_passthrough_nullability.get( - name, - _infer_polyglot_nullability( + nullability = ( + InferredNullability.NON_NULL + if _polyglot_expression_is_non_null_after_filter( expression=inner, - alias_nullability=alias_nullability, - column_nullability_by_table=column_nullability_by_table, - inference_profile=inference_profile, - ), + context=non_null_filter_context, + ) + else cte_passthrough_nullability.get( + name, + _infer_polyglot_nullability( + expression=inner, + alias_nullability=alias_nullability, + column_nullability_by_table=column_nullability_by_table, + inference_profile=inference_profile, + ), + ) ) columns.append(InferredColumn(name=name, type=col_type, nullability=nullability)) return tuple(columns) @@ -935,6 +965,10 @@ def _analyze_columns_and_lineage_from_polyglot_ast( if recover_cte_facts else {} ) + non_null_filter_context: NonNullFilterContext | None = _polyglot_non_null_filter_context( + select=select, + column_nullability_by_table=column_nullability_by_table, + ) columns: list[InferredColumn] = [] lineage_columns: list[CompiledLineageColumnFact] = [] has_star: bool = False @@ -963,21 +997,28 @@ def _analyze_columns_and_lineage_from_polyglot_ast( ) nullability: InferredNullability = InferredNullability.UNKNOWN if infer_nullability: - nullability = cte_passthrough_nullability.get( - output_column, - ( - _infer_polyglot_nullability( - expression=inner, - alias_nullability=alias_nullability, - column_nullability_by_table=column_nullability_by_table, - inference_profile=inference_profile, - ) - if has_known_nullability - else _infer_polyglot_shallow_nullability( - expression=inner, - inference_profile=inference_profile, - ) - ), + nullability = ( + InferredNullability.NON_NULL + if _polyglot_expression_is_non_null_after_filter( + expression=inner, + context=non_null_filter_context, + ) + else cte_passthrough_nullability.get( + output_column, + ( + _infer_polyglot_nullability( + expression=inner, + alias_nullability=alias_nullability, + column_nullability_by_table=column_nullability_by_table, + inference_profile=inference_profile, + ) + if has_known_nullability + else _infer_polyglot_shallow_nullability( + expression=inner, + inference_profile=inference_profile, + ) + ), + ) ) columns.append(InferredColumn(name=output_column, type=col_type, nullability=nullability)) diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/cte_facts.py b/src/sqlbuild/compiler/compile/_helpers/analysis/cte_facts.py index da0e49dbcb..d9b892311d 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/cte_facts.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/cte_facts.py @@ -9,8 +9,20 @@ from sqlbuild.adapter.type_system.main.normalize_type import normalize_type from sqlbuild.adapter.type_system.main.types_equal import types_equal from sqlbuild.compiler.compile.constants import SQL_WILDCARD_TOKEN -from sqlbuild.compiler.compile.models import CteFactResolvers +from sqlbuild.compiler.compile.models import CteFactResolvers, NonNullFilterContext from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_ANALYSIS_COLUMN_USES as _POLYGLOT_ANALYSIS_COLUMN_USES, +) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_ANALYSIS_CONTEXT as _POLYGLOT_ANALYSIS_CONTEXT, +) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_ANALYSIS_CONTEXT_FILTER as _POLYGLOT_ANALYSIS_CONTEXT_FILTER, +) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_ANALYSIS_EXPRESSION_SQL as _POLYGLOT_ANALYSIS_EXPRESSION_SQL, +) from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_BINARY_OPERAND_COUNT as _POLYGLOT_BINARY_OPERAND_COUNT, ) @@ -18,10 +30,12 @@ POLYGLOT_BOOLEAN_RESULT_KINDS as _POLYGLOT_BOOLEAN_RESULT_KINDS, ) from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_ALIAS as _POLYGLOT_KIND_ALIAS +from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_AND as _POLYGLOT_KIND_AND from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_KIND_ANNOTATED as _POLYGLOT_KIND_ANNOTATED, ) from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_CASE as _POLYGLOT_KIND_CASE +from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_CAST as _POLYGLOT_KIND_CAST from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_KIND_COALESCE as _POLYGLOT_KIND_COALESCE, ) @@ -33,8 +47,12 @@ from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_KIND_IF_FUNC as _POLYGLOT_KIND_IF_FUNC, ) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_KIND_IS_NULL as _POLYGLOT_KIND_IS_NULL, +) from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_LITERAL as _POLYGLOT_KIND_LITERAL from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_NULL as _POLYGLOT_KIND_NULL +from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_PAREN as _POLYGLOT_KIND_PAREN from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_KIND_SELECT as _POLYGLOT_KIND_SELECT from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_KIND_SUBSTRING as _POLYGLOT_KIND_SUBSTRING, @@ -49,19 +67,35 @@ from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_PAYLOAD_COLUMN as _POLYGLOT_PAYLOAD_COLUMN, ) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_PAYLOAD_LEFT as _POLYGLOT_PAYLOAD_LEFT, +) from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_PAYLOAD_LITERAL_TYPE as _POLYGLOT_PAYLOAD_LITERAL_TYPE, ) from sqlbuild.compiler.sql_analysis.constants import POLYGLOT_PAYLOAD_NAME as _POLYGLOT_PAYLOAD_NAME +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_PAYLOAD_NOT as _POLYGLOT_PAYLOAD_NOT, +) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_PAYLOAD_RIGHT as _POLYGLOT_PAYLOAD_RIGHT, +) from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_PAYLOAD_TABLE as _POLYGLOT_PAYLOAD_TABLE, ) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_PAYLOAD_THIS as _POLYGLOT_PAYLOAD_THIS, +) +from sqlbuild.compiler.sql_analysis.constants import ( + POLYGLOT_PAYLOAD_WHERE_CLAUSE as _POLYGLOT_PAYLOAD_WHERE_CLAUSE, +) from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_SET_OPERATION_KINDS as _POLYGLOT_SET_OPERATION_KINDS, ) from sqlbuild.compiler.sql_analysis.constants import ( POLYGLOT_TYPE_PASSTHROUGH_KINDS as _POLYGLOT_TYPE_PASSTHROUGH_KINDS, ) +from sqlbuild.compiler.sql_analysis.constants import SQL_NULL_KEYWORD as _SQL_NULL_KEYWORD _NULL_SET_OPERATION_TYPE: str = "__SQLBUILD_NULL_SET_OPERATION_TYPE__" @@ -108,18 +142,18 @@ def _polyglot_cte_passthrough_facts( inference_profile: ExpressionInferenceProfile, analysis: dict[str, Any], resolvers: CteFactResolvers, -) -> tuple[dict[str, str], dict[str, InferredNullability], frozenset[str]]: +) -> tuple[dict[str, str], dict[str, InferredNullability], frozenset[str], Any | None]: """Recover facts that compact analysis skips across direct CTE reads.""" if not analysis.get("cteFacts") or not (column_types_by_table or column_nullability_by_table): - return {}, {}, frozenset() + return {}, {}, frozenset(), None try: parsed: Any = polyglot_module.parse_one(cleaned_sql, dialect=dialect or "generic") except polyglot_module.PolyglotError: - return {}, {}, frozenset() + return {}, {}, frozenset(), None ctes: tuple[tuple[str, Any, bool], ...] = _polyglot_top_level_ctes(parsed) if not ctes: - return {}, {}, frozenset() + return {}, {}, frozenset(), parsed inferred_types: dict[str, str] = ( _polyglot_cte_passthrough_types_from_parsed( parsed=parsed, @@ -149,6 +183,7 @@ def _polyglot_cte_passthrough_facts( parsed=parsed, cte_names=frozenset(name for name, _, _ in ctes), ), + parsed, ) @@ -303,6 +338,10 @@ def _polyglot_direct_select_output_nullability( if has_known_nullability else {} ) + non_null_filter_context: NonNullFilterContext | None = _polyglot_non_null_filter_context( + select=select, + column_nullability_by_table=scoped_nullability, + ) inferred: dict[str, InferredNullability] = {} projection: Any for raw_projection in getattr(select, "expressions", ()): @@ -318,16 +357,23 @@ def _polyglot_direct_select_output_nullability( else projection ) nullability: InferredNullability = ( - nullability_resolver( + InferredNullability.NON_NULL + if _polyglot_expression_is_non_null_after_filter( expression=expression, - alias_nullability=alias_nullability, - column_nullability_by_table=scoped_nullability, - inference_profile=inference_profile, + context=non_null_filter_context, ) - if has_known_nullability - else shallow_nullability_resolver( - expression=expression, - inference_profile=inference_profile, + else ( + nullability_resolver( + expression=expression, + alias_nullability=alias_nullability, + column_nullability_by_table=scoped_nullability, + inference_profile=inference_profile, + ) + if has_known_nullability + else shallow_nullability_resolver( + expression=expression, + inference_profile=inference_profile, + ) ) ) if nullability != InferredNullability.UNKNOWN: @@ -335,6 +381,187 @@ def _polyglot_direct_select_output_nullability( return inferred +def _polyglot_filtered_non_null_outputs( + *, + polyglot_module: Any, + cleaned_sql: str, + dialect: str | None, + analysis: dict[str, Any], + column_nullability_by_table: dict[str, dict[str, InferredNullability]], + parsed: Any | None = None, +) -> frozenset[str]: + column_uses: object = analysis.get(_POLYGLOT_ANALYSIS_COLUMN_USES) + if not isinstance(column_uses, list) or not any( + isinstance(value, dict) + and value.get(_POLYGLOT_ANALYSIS_CONTEXT) == _POLYGLOT_ANALYSIS_CONTEXT_FILTER + and _SQL_NULL_KEYWORD in str(value.get(_POLYGLOT_ANALYSIS_EXPRESSION_SQL) or "").upper() + for value in column_uses + ): + return frozenset() + if parsed is None: + try: + parsed = polyglot_module.parse_one( + cleaned_sql, + dialect=dialect or "generic", + ) + except polyglot_module.PolyglotError: + return frozenset() + select: Any | None = parsed + if str(getattr(select, "kind", "")) != _POLYGLOT_KIND_SELECT: + select = parsed.find(_POLYGLOT_KIND_SELECT) + if select is None or str(getattr(select, "kind", "")) != _POLYGLOT_KIND_SELECT: + return frozenset() + outputs: set[str] = set() + context: NonNullFilterContext | None = _polyglot_non_null_filter_context( + select=select, + column_nullability_by_table=column_nullability_by_table, + ) + for raw_projection in getattr(select, "expressions", ()): + projection: Any = _unwrap_polyglot_annotations(raw_projection) + output_name: str = str(getattr(projection, "output_name", "") or "") + if not output_name or output_name == SQL_WILDCARD_TOKEN: + continue + expression: Any = ( + projection.this + if str(getattr(projection, "kind", "")) == _POLYGLOT_KIND_ALIAS + else projection + ) + if _polyglot_expression_is_non_null_after_filter( + expression=expression, + context=context, + ): + outputs.add(output_name) + return frozenset(outputs) + + +def _polyglot_expression_is_non_null_after_filter( + *, + expression: Any, + context: NonNullFilterContext | None, +) -> bool: + if context is None: + return False + expression_reference: tuple[str, str] | None = _polyglot_direct_column_reference(expression) + if expression_reference is None: + return False + resolved_expression: tuple[str, str] | None = _polyglot_resolved_relation_column( + reference=expression_reference, + relations=context.relations, + ) + return resolved_expression is not None and resolved_expression in context.columns + + +def _polyglot_non_null_filter_context( + *, + select: Any, + column_nullability_by_table: dict[str, dict[str, InferredNullability]], +) -> NonNullFilterContext | None: + select_args: object = getattr(select, "args", None) + if not isinstance(select_args, dict): + return None + select_args_dict: dict[str, object] = cast(dict[str, object], select_args) + where_payload: object = select_args_dict.get(_POLYGLOT_PAYLOAD_WHERE_CLAUSE) + if not isinstance(where_payload, dict): + return None + predicate: object = cast(dict[str, object], where_payload).get(_POLYGLOT_PAYLOAD_THIS) + non_null_references: tuple[tuple[str, str], ...] = _polyglot_non_null_conjuncts(predicate) + if not non_null_references: + return None + relations: list[tuple[str, dict[str, InferredNullability]]] = [] + for table in _polyglot_direct_select_tables(select): + table_name: str = str(getattr(table, "name", "") or "") + alias: str = str(getattr(table, "alias_or_name", "") or "") + columns: dict[str, InferredNullability] | None = _case_insensitive_mapping_get( + mapping=column_nullability_by_table, + key=table_name, + ) + if columns is not None: + relations.append((alias or table_name, columns)) + relation_tuple: tuple[tuple[str, dict[str, InferredNullability]], ...] = tuple(relations) + columns: frozenset[tuple[str, str]] = frozenset( + resolved + for reference in non_null_references + if ( + resolved := _polyglot_resolved_relation_column( + reference=reference, + relations=relation_tuple, + ) + ) + ) + return NonNullFilterContext(relations=relation_tuple, columns=columns) if columns else None + + +def _polyglot_direct_column_reference(expression: Any) -> tuple[str, str] | None: + expression = _unwrap_polyglot_annotations(expression) + if str(getattr(expression, "kind", "")) == _POLYGLOT_KIND_CAST: + expression = _unwrap_polyglot_annotations(getattr(expression, "this", None)) + if str(getattr(expression, "kind", "")) != _POLYGLOT_KIND_COLUMN: + return None + column_name: str = str(getattr(expression, "name", "") or "") + if not column_name: + return None + expression_args: object = getattr(expression, "args", None) + table_name: str = "" + if isinstance(expression_args, dict): + table_name = _polyglot_name_payload_value( + cast(dict[str, object], expression_args).get(_POLYGLOT_PAYLOAD_TABLE) + ) + return table_name, column_name + + +def _polyglot_resolved_relation_column( + *, + reference: tuple[str, str], + relations: tuple[tuple[str, dict[str, InferredNullability]], ...], +) -> tuple[str, str] | None: + table_name, column_name = reference + matching: list[tuple[str, str]] = [] + for alias, columns in relations: + if table_name and alias.casefold() != table_name.casefold(): + continue + column_key: str | None = _case_insensitive_mapping_key( + mapping=columns, + key=column_name, + ) + if column_key is not None: + matching.append((alias.casefold(), column_key.casefold())) + return matching[0] if len(matching) == 1 else None + + +def _polyglot_non_null_conjuncts(payload: object) -> tuple[tuple[str, str], ...]: + if not isinstance(payload, dict): + return () + payload_dict: dict[str, object] = cast(dict[str, object], payload) + paren_payload: object = payload_dict.get(_POLYGLOT_KIND_PAREN) + if isinstance(paren_payload, dict): + paren_payload_dict: dict[str, object] = cast(dict[str, object], paren_payload) + return _polyglot_non_null_conjuncts(paren_payload_dict.get(_POLYGLOT_PAYLOAD_THIS)) + and_payload: object = payload_dict.get(_POLYGLOT_KIND_AND) + if isinstance(and_payload, dict): + and_payload_dict: dict[str, object] = cast(dict[str, object], and_payload) + return ( + *_polyglot_non_null_conjuncts(and_payload_dict.get(_POLYGLOT_PAYLOAD_LEFT)), + *_polyglot_non_null_conjuncts(and_payload_dict.get(_POLYGLOT_PAYLOAD_RIGHT)), + ) + is_null_payload: object = payload_dict.get(_POLYGLOT_KIND_IS_NULL) + is_null_payload_dict: dict[str, object] = ( + cast(dict[str, object], is_null_payload) if isinstance(is_null_payload, dict) else {} + ) + if not is_null_payload_dict or is_null_payload_dict.get(_POLYGLOT_PAYLOAD_NOT) is not True: + return () + expression_payload: object = is_null_payload_dict.get(_POLYGLOT_PAYLOAD_THIS) + if not isinstance(expression_payload, dict): + return () + expression_payload_dict: dict[str, object] = cast(dict[str, object], expression_payload) + column_payload: object = expression_payload_dict.get(_POLYGLOT_PAYLOAD_COLUMN) + if not isinstance(column_payload, dict): + return () + column_payload_dict: dict[str, object] = cast(dict[str, object], column_payload) + column_name: str = _polyglot_name_payload_value(column_payload_dict.get(_POLYGLOT_PAYLOAD_NAME)) + table_name: str = _polyglot_name_payload_value(column_payload_dict.get(_POLYGLOT_PAYLOAD_TABLE)) + return ((table_name, column_name),) if column_name else () + + def _polyglot_star_output_types( *, projection: Any, types_by_alias: dict[str, dict[str, str]] ) -> dict[str, str]: @@ -907,6 +1134,16 @@ def _case_insensitive_mapping_get[T](*, mapping: dict[str, T], key: str) -> T | return matches[0] if len(matches) == 1 else None +def _case_insensitive_mapping_key[T](*, mapping: dict[str, T], key: str) -> str | None: + if key in mapping: + return key + normalized_key: str = key.casefold() + matches: list[str] = [ + candidate for candidate in mapping if candidate.casefold() == normalized_key + ] + return matches[0] if len(matches) == 1 else None + + def _has_known_nullability( column_nullability_by_table: dict[str, dict[str, InferredNullability]], ) -> bool: diff --git a/src/sqlbuild/compiler/compile/models.py b/src/sqlbuild/compiler/compile/models.py index 37eed5cf47..eb63487385 100644 --- a/src/sqlbuild/compiler/compile/models.py +++ b/src/sqlbuild/compiler/compile/models.py @@ -154,6 +154,14 @@ class CteFactResolvers: alias_nullability: Callable[..., Any] +@dataclass(frozen=True) +class NonNullFilterContext: + """Resolved relations and columns constrained by non-null filters.""" + + relations: tuple[tuple[str, dict[str, InferredNullability]], ...] + columns: frozenset[tuple[str, str]] + + @dataclass(frozen=True) class CompileModelConfig: """Pre-semantic effective config layers attached to a model input.""" diff --git a/src/sqlbuild/compiler/sql_analysis/constants.py b/src/sqlbuild/compiler/sql_analysis/constants.py index 0dc4fc4a4d..8dbe2e8bd3 100644 --- a/src/sqlbuild/compiler/sql_analysis/constants.py +++ b/src/sqlbuild/compiler/sql_analysis/constants.py @@ -12,10 +12,12 @@ SQL_ESCAPE_CHARACTER: str = "\\" SQL_DOLLAR_QUOTE_DELIMITER: str = "$$" SQL_IDENTIFIER_PREFIX: str = "_" +SQL_NULL_KEYWORD: str = "NULL" SNOWFLAKE_DIALECT_NAME: str = "snowflake" BINDING_UNKNOWN_TABLE_INTERNAL_CODE: str = "B000" POLYGLOT_KIND_ALIAS: str = "alias" +POLYGLOT_KIND_AND: str = "and" POLYGLOT_KIND_ANNOTATED: str = "annotated" POLYGLOT_KIND_ARRAY_AGG: str = "array_agg" POLYGLOT_KIND_AVG: str = "avg" @@ -30,10 +32,12 @@ POLYGLOT_KIND_FUNCTION: str = "function" POLYGLOT_KIND_IF_FUNC: str = "if_func" POLYGLOT_KIND_INTERSECT: str = "intersect" +POLYGLOT_KIND_IS_NULL: str = "is_null" POLYGLOT_KIND_LITERAL: str = "literal" POLYGLOT_KIND_MAX: str = "max" POLYGLOT_KIND_MIN: str = "min" POLYGLOT_KIND_NULL: str = "null" +POLYGLOT_KIND_PAREN: str = "paren" POLYGLOT_KIND_SELECT: str = "select" POLYGLOT_KIND_STRING_AGG: str = "string_agg" POLYGLOT_KIND_SUBSTRING: str = "substring" @@ -93,16 +97,20 @@ POLYGLOT_PAYLOAD_DATA_TYPE: str = "data_type" POLYGLOT_VARCHAR_DATA_TYPES: frozenset[str] = frozenset({"var_char", "varchar"}) POLYGLOT_PAYLOAD_FROM: str = "from" +POLYGLOT_PAYLOAD_LEFT: str = "left" POLYGLOT_PAYLOAD_EXPRESSIONS: str = "expressions" POLYGLOT_PAYLOAD_JOINS: str = "joins" POLYGLOT_PAYLOAD_KIND: str = "kind" POLYGLOT_PAYLOAD_LITERAL_TYPE: str = "literal_type" POLYGLOT_PAYLOAD_NAME: str = "name" +POLYGLOT_PAYLOAD_NOT: str = "not" POLYGLOT_PAYLOAD_PRECISION: str = "precision" POLYGLOT_PAYLOAD_SCALE: str = "scale" POLYGLOT_PAYLOAD_SELECT: str = "select" +POLYGLOT_PAYLOAD_RIGHT: str = "right" POLYGLOT_PAYLOAD_TABLE: str = "table" POLYGLOT_PAYLOAD_THIS: str = "this" +POLYGLOT_PAYLOAD_WHERE_CLAUSE: str = "where_clause" POLYGLOT_PAYLOAD_TIMEZONE: str = "timezone" POLYGLOT_PAYLOAD_TO: str = "to" TIMESTAMP_WITH_TIME_ZONE_SQL_TYPE_NAME: str = "TIMESTAMPTZ" @@ -115,8 +123,12 @@ POLYGLOT_ANALYSIS_BASE_TABLES: str = "baseTables" POLYGLOT_ANALYSIS_BODY_SQL: str = "bodySql" POLYGLOT_ANALYSIS_CAST_TYPE: str = "castType" +POLYGLOT_ANALYSIS_COLUMN_USES: str = "columnUses" +POLYGLOT_ANALYSIS_CONTEXT: str = "context" +POLYGLOT_ANALYSIS_CONTEXT_FILTER: str = "filter" POLYGLOT_ANALYSIS_CTES: str = "ctes" POLYGLOT_ANALYSIS_CTE_FACTS: str = "cteFacts" +POLYGLOT_ANALYSIS_EXPRESSION_SQL: str = "expressionSql" POLYGLOT_ANALYSIS_IS_STAR: str = "isStar" POLYGLOT_ANALYSIS_NAME: str = "name" POLYGLOT_ANALYSIS_NULLABILITY: str = "nullability" diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py index b60a5df3b8..2ef2ffbda6 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py @@ -279,6 +279,91 @@ InferredColumn(name="status", nullability=InferredNullability.UNKNOWN), ), ), + InferColumnsTestCase( + description="refines qualified filtered output nullability", + query_sql=( + 'SELECT o.order_id FROM __ref("orders") AS o WHERE (O.order_id IS NOT NULL)' + ), + column_nullability_by_table={"orders": {"order_id": InferredNullability.NULLABLE}}, + expected_columns=( + InferredColumn( + name="order_id", + nullability=InferredNullability.NON_NULL, + ), + ), + ), + InferColumnsTestCase( + description="does not refine nullability through a disjunction", + query_sql=( + 'SELECT order_id FROM __ref("orders") ' + "WHERE order_id IS NOT NULL OR status = 'ready'" + ), + column_nullability_by_table={ + "orders": { + "order_id": InferredNullability.NULLABLE, + "status": InferredNullability.NON_NULL, + } + }, + expected_columns=( + InferredColumn( + name="order_id", + nullability=InferredNullability.NULLABLE, + ), + ), + ), + InferColumnsTestCase( + description="does not leak a filtered self join alias", + query_sql=( + "WITH selected AS (" + 'SELECT b.order_id FROM __ref("orders") AS a ' + 'CROSS JOIN __ref("orders") AS b ' + "WHERE a.order_id IS NOT NULL" + ") SELECT order_id FROM selected" + ), + column_nullability_by_table={"orders": {"order_id": InferredNullability.NULLABLE}}, + expected_columns=( + InferredColumn( + name="order_id", + nullability=InferredNullability.NULLABLE, + ), + ), + ), + InferColumnsTestCase( + description="does not leak an outer filter into a CTE", + query_sql=( + "WITH selected AS (" + 'SELECT order_id FROM __ref("orders")' + ") SELECT selected.order_id FROM selected " + 'CROSS JOIN __ref("orders") AS o ' + "WHERE o.order_id IS NOT NULL" + ), + column_nullability_by_table={"orders": {"order_id": InferredNullability.NULLABLE}}, + expected_columns=( + InferredColumn( + name="order_id", + nullability=InferredNullability.NULLABLE, + ), + ), + ), + InferColumnsTestCase( + description="refines a filtered outer join column", + query_sql=( + 'SELECT c.customer_id FROM __ref("orders") AS o ' + 'LEFT JOIN __ref("customers") AS c ' + "ON o.order_id = c.customer_id " + "WHERE c.customer_id IS NOT NULL" + ), + column_nullability_by_table={ + "orders": {"order_id": InferredNullability.NON_NULL}, + "customers": {"customer_id": InferredNullability.NON_NULL}, + }, + expected_columns=( + InferredColumn( + name="customer_id", + nullability=InferredNullability.NON_NULL, + ), + ), + ), InferColumnsTestCase( description="marks right side of left join nullable", query_sql=( @@ -348,6 +433,25 @@ ), ), ), + InferColumnsTestCase( + description="refines filtered CTE output nullability", + query_sql=( + "WITH mapped AS (" + 'SELECT order_id FROM __ref("orders")' + "), final AS (" + "SELECT CAST(order_id AS VARCHAR) AS order_id FROM mapped " + "WHERE order_id IS NOT NULL" + ") SELECT order_id FROM final" + ), + column_nullability_by_table={"orders": {"order_id": InferredNullability.NULLABLE}}, + expected_columns=( + InferredColumn( + name="order_id", + type="TEXT", + nullability=InferredNullability.NON_NULL, + ), + ), + ), InferColumnsTestCase( description="leaves arbitrary expressions unknown", query_sql='SELECT quantity * price AS total FROM __ref("orders")', @@ -419,6 +523,65 @@ def test_given_query_sql_when_inferring_columns_then_returns_expected( ), expected_has_star=False, ), + PolyglotAnalysisTestCase( + description="refines a compact filtered column across whitespace", + query_sql=('SELECT order_id FROM __ref("orders") WHERE order_id IS\nNOT NULL'), + references=(CompileSqlReference(SqlReferenceKind.REF, "orders"),), + column_nullability_by_table={"orders": {"order_id": InferredNullability.NULLABLE}}, + column_types_by_table={"orders": {"order_id": "INTEGER"}}, + expected_columns=( + InferredColumn( + name="order_id", + type="INT", + nullability=InferredNullability.NON_NULL, + ), + ), + expected_lineage_columns=direct_orders_lineage("order_id"), + expected_has_star=False, + ), + PolyglotAnalysisTestCase( + description="refines a compact filtered outer join across a comment", + query_sql=( + 'SELECT c.customer_id FROM __ref("orders") AS o ' + 'LEFT JOIN __ref("customers") AS c ' + "ON o.order_id = c.customer_id " + "WHERE c.customer_id IS /* required customer */ NOT NULL" + ), + references=( + CompileSqlReference(SqlReferenceKind.REF, "orders"), + CompileSqlReference(SqlReferenceKind.REF, "customers"), + ), + column_nullability_by_table={ + "orders": {"order_id": InferredNullability.NON_NULL}, + "customers": {"customer_id": InferredNullability.NON_NULL}, + }, + column_types_by_table={ + "orders": {"order_id": "INTEGER"}, + "customers": {"customer_id": "INTEGER"}, + }, + expected_columns=( + InferredColumn( + name="customer_id", + type="INT", + nullability=InferredNullability.NON_NULL, + ), + ), + expected_lineage_columns=( + CompiledLineageColumnFact( + output_column="customer_id", + upstream_columns=( + CompiledLineageSourceFact( + resource_type=CompiledResourceType.MODEL, + resource_name="customers", + column_name="customer_id", + ), + ), + transform_kind=ColumnTransformKind.DIRECT, + confidence=ColumnLineageConfidence.HIGH, + ), + ), + expected_has_star=False, + ), PolyglotAnalysisTestCase( description="terminates direct lineage at a managed table function output", query_sql='SELECT order_id FROM __table_fn("customer_orders")(42)', @@ -475,8 +638,11 @@ def test_given_ref_query_when_analyzing_columns_and_lineage_then_returns_compact result: PolyglotAnalysisResult = analyze_columns_and_lineage_with_polyglot( query_sql=test_case.query_sql, references=test_case.references, + column_nullability_by_table=test_case.column_nullability_by_table, + column_types_by_table=test_case.column_types_by_table, inference_profile=test_case.inference_profile, allow_compact_analysis=True, + recover_cte_facts=True, ) assert result.analysis_succeeded