Initial commit

Author: anuunchin

dlt/common/libs/sqlglot.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, Set, Any, Iterable, Literal
+from typing import Optional, Union, Set, Any, Iterable, Literal, cast
 
 from dlt.common.utils import without_none
 from dlt.common.exceptions import TerminalValueError
@@ -632,3 +632,14 @@ def _literal(v: Any) -> sge.Expression:
         return sge.Tuple(expressions=[_literal(v) for v in value])
     else:
         return _literal(value)
+
+
+def wrap_identifiers_in_columns(expression: sge.Query) -> sge.Query:
+    """Wrap bare Identifier nodes inside Alias with Column nodes"""
+
+    def transform(node: sge.Expression) -> sge.Expression:
+        if isinstance(node, sge.Alias) and isinstance(node.this, sge.Identifier):
+            return sge.Alias(this=sge.Column(this=node.this.copy()), alias=node.alias)
+        return node
+
+    return cast(sge.Query, expression.transform(transform))

dlt/dataset/lineage.py

@@ -2,7 +2,6 @@
 from typing import Optional, Tuple, cast
 
 import sqlglot.expressions as sge
-
 from sqlglot.errors import OptimizeError
 from sqlglot.schema import Schema as SQLGlotSchema, ensure_schema
 from sqlglot.optimizer.annotate_types import annotate_types
@@ -16,6 +15,7 @@
     set_metadata,
     get_metadata,
     TSqlGlotDialect,
+    wrap_identifiers_in_columns,
 )
 from dlt.common.schema.typing import (
     TTableSchemaColumns,
@@ -111,6 +111,8 @@ def compute_columns_schema(
     else:
         select_expression = expression
 
+    select_expression = wrap_identifiers_in_columns(select_expression)
+
     # prevent normalization
     select_expression.meta["case_sensitive"] = True
 
@@ -132,7 +134,7 @@ def compute_columns_schema(
             f"Failed to resolve SQL query against the schema received: {e}"
         ) from e
 
-    expression = annotate_types(expression, schema=sqlglot_schema)
+    select_expression = annotate_types(select_expression, schema=sqlglot_schema)
 
     # NOTE: this has to be fixed
     if allow_anonymous_columns is False:
@@ -173,4 +175,4 @@ def compute_columns_schema(
         if propagated_name and col.output_name != propagated_name:
             dlt_table_schema[col.output_name]["x-original-name"] = propagated_name  # type: ignore[typeddict-unknown-key]
 
-    return dlt_table_schema, expression
+    return dlt_table_schema, select_expression

tests/hub/test_transformations.py

@@ -1,4 +1,5 @@
 import dlt
+from dlt.sources.rest_api import rest_api_resources
 
 
 def test_transformation_decorator() -> None:
@@ -12,3 +13,71 @@ def get_even_rows(dataset: dlt.Dataset):
     # get instance without license
     transformation = get_even_rows(dlt.dataset("duckdb", "mock_dataset"))
     assert transformation.name == "get_even_rows"
+
+
+def test_missing_columns_bug() -> None:
+    """Regression test: bare Identifier nodes were not properly type annotated in dlt.dataset.lineage.compute_columns_schema,
+    causing success_count and success_rate to have UNKNOWN typetype and be excluded as incomplete columns by dlt.
+    """
+    import dlthub.data_quality as dq
+
+    @dlt.source
+    def jaffleshop():
+        jaffle_rest_resources = rest_api_resources(
+            {
+                "client": {
+                    "base_url": "https://jaffle-shop.dlthub.com/api/v1",
+                    "paginator": {"type": "header_link"},
+                },
+                "resources": [
+                    "customers",
+                    "products",
+                    "orders",
+                ],
+                "resource_defaults": {
+                    "endpoint": {
+                        "params": {
+                            "start_date": "2017-01-01",
+                            "end_date": "2017-01-15",
+                        },
+                    },
+                },
+            }
+        )
+
+        return jaffle_rest_resources
+
+    @dlt.hub.transformation
+    def jaffle_checks(dataset: dlt.Dataset) -> dlt.Relation:
+        checks = {"orders": [dq.checks.is_unique("id"), dq.checks.case("subtotal > 0")]}
+        return dq.prepare_checks(dataset, checks=checks)  # type: ignore
+
+    pipeline = dlt.pipeline("test_missing_columns", destination="duckdb")
+    pipeline.run([jaffleshop()])
+    pipeline.run(jaffle_checks(pipeline.dataset()))
+
+    expected_column_names = [
+        "table_name",
+        "check_qualified_name",
+        "row_count",
+        "success_count",  # was missing due to unqualified UNION columns
+        "success_rate",  # was missing due to unqualified UNION columns
+    ]
+
+    # direct query execution returns raw select output (no dlt columns)
+    query = dq.prepare_checks(
+        pipeline.dataset(),
+        checks={
+            "orders": [dq.checks.is_unique("id"), dq.checks.case("subtotal > 0")],  # type: ignore
+        },
+    )
+    assert query.arrow().column_names == expected_column_names
+
+    # materialized table includes _dlt_load_id added by pipeline
+    with pipeline.sql_client() as client:
+        with client.execute_query(
+            f"SELECT * FROM {pipeline.pipeline_name}.{pipeline.dataset_name}.jaffle_checks"
+        ) as cursor:
+            df = cursor.df()
+            columns = list(df.columns)
+    assert columns == expected_column_names + ["_dlt_load_id"]