feat(tableau): ability to force extraction of table/column level lina…

…ge from SQL queries (#9838)

metadata-ingestion/src/datahub/ingestion/source/tableau_common.py

@@ -2,7 +2,7 @@
 import logging
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from pydantic.fields import Field
 from tableauserverclient import Server
@@ -762,8 +762,19 @@ def make_upstream_class(
 
 
 def make_fine_grained_lineage_class(
-    parsed_result: Optional[SqlParsingResult], dataset_urn: str
+    parsed_result: Optional[SqlParsingResult],
+    dataset_urn: str,
+    out_columns: List[Dict[Any, Any]],
 ) -> List[FineGrainedLineage]:
+    # 1) fine grained lineage links are case sensitive
+    # 2) parsed out columns are always lower cased
+    # 3) corresponding Custom SQL output columns can be in any case lower/upper/mix
+    #
+    # we need a map between 2 and 3 that will be used during building column level linage links (see below)
+    out_columns_map = {
+        col.get(c.NAME, "").lower(): col.get(c.NAME, "") for col in out_columns
+    }
+
     fine_grained_lineages: List[FineGrainedLineage] = []
 
     if parsed_result is None:
@@ -775,7 +786,15 @@ def make_fine_grained_lineage_class(
 
     for cll_info in cll:
         downstream = (
-            [builder.make_schema_field_urn(dataset_urn, cll_info.downstream.column)]
+            [
+                builder.make_schema_field_urn(
+                    dataset_urn,
+                    out_columns_map.get(
+                        cll_info.downstream.column.lower(),
+                        cll_info.downstream.column,
+                    ),
+                )
+            ]
             if cll_info.downstream is not None
             and cll_info.downstream.column is not None
             else []

metadata-ingestion/src/datahub/sql_parsing/sql_parsing_result_utils.py

@@ -0,0 +1,23 @@
+from typing import Dict, Set
+
+from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult, Urn
+
+
+def transform_parsing_result_to_in_tables_schemas(
+    parsing_result: SqlParsingResult,
+) -> Dict[Urn, Set[str]]:
+    table_urn_to_schema_map: Dict[str, Set[str]] = (
+        {it: set() for it in parsing_result.in_tables}
+        if parsing_result.in_tables
+        else {}
+    )
+
+    if parsing_result.column_lineage:
+        for cli in parsing_result.column_lineage:
+            for upstream in cli.upstreams:
+                if upstream.table in table_urn_to_schema_map:
+                    table_urn_to_schema_map[upstream.table].add(upstream.column)
+                else:
+                    table_urn_to_schema_map[upstream.table] = {upstream.column}
+
+    return table_urn_to_schema_map

metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py

@@ -529,6 +529,9 @@ def _schema_aware_fuzzy_column_resolve(
 
                     # Parse the column name out of the node name.
                     # Sqlglot calls .sql(), so we have to do the inverse.
+                    if node.name == "*":
+                        continue
+
                     normalized_col = sqlglot.parse_one(node.name).this.name
                     if node.subfield:
                         normalized_col = f"{normalized_col}.{node.subfield}"
@@ -834,6 +837,7 @@ def _sqlglot_lineage_inner(
     # Fetch schema info for the relevant tables.
     table_name_urn_mapping: Dict[_TableName, str] = {}
     table_name_schema_mapping: Dict[_TableName, SchemaInfo] = {}
+
     for table in tables | modified:
         # For select statements, qualification will be a no-op. For other statements, this
         # is where the qualification actually happens.
@@ -1016,8 +1020,9 @@ def create_lineage_sql_parsed_result(
     env: str,
     default_schema: Optional[str] = None,
     graph: Optional[DataHubGraph] = None,
+    schema_aware: bool = True,
 ) -> SqlParsingResult:
-    if graph:
+    if graph and schema_aware:
         needs_close = False
         schema_resolver = graph._make_schema_resolver(
             platform=platform,

metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py

@@ -16,6 +16,11 @@ def _get_dialect_str(platform: str) -> str:
         return "tsql"
     elif platform == "athena":
         return "trino"
+    # TODO: define SalesForce SOQL dialect
+    # Temporary workaround is to treat SOQL as databricks dialect
+    # At least it allows to parse simple SQL queries and built linage for them
+    elif platform == "salesforce":
+        return "databricks"
     elif platform in {"mysql", "mariadb"}:
         # In sqlglot v20+, MySQL is now case-sensitive by default, which is the
         # default behavior on Linux. However, MySQL's default case sensitivity
@@ -31,6 +36,7 @@ def _get_dialect_str(platform: str) -> str:
 def get_dialect(platform: DialectOrStr) -> sqlglot.Dialect:
     if isinstance(platform, sqlglot.Dialect):
         return platform
+
     return sqlglot.Dialect.get_or_raise(_get_dialect_str(platform))
 
 

metadata-ingestion/tests/integration/tableau/tableau_cll_mces_golden.json

@@ -42870,6 +42870,38 @@
         "lastRunId": "no-run-id-provided"
     }
 },
+{
+    "entityType": "dataset",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:bigquery,demo-custom-323403.bigquery_demo.order_items,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1638860400000,
+        "runId": "tableau-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
+{
+    "entityType": "dataset",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:bigquery,demo-custom-323403.bigquery_demo.sellers,PROD)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1638860400000,
+        "runId": "tableau-test",
+        "lastRunId": "no-run-id-provided"
+    }
+},
 {
     "entityType": "dataset",
     "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:external,sample - superstore%2C %28new%29.xls.orders,PROD)",

metadata-ingestion/tests/integration/tableau/test_tableau_ingest.py

@@ -319,6 +319,8 @@ def test_tableau_cll_ingest(pytestconfig, tmp_path, mock_datahub_graph):
     new_pipeline_config: Dict[Any, Any] = {
         **config_source_default,
         "extract_lineage_from_unsupported_custom_sql_queries": True,
+        "force_extraction_of_lineage_from_custom_sql_queries": False,
+        "sql_parsing_disable_schema_awareness": False,
         "extract_column_level_lineage": True,
     }
 
@@ -834,6 +836,7 @@ def test_tableau_unsupported_csql(mock_datahub_graph):
                     "connectionType": "bigquery",
                 },
             },
+            out_columns=[],
         )
 
         mcp = cast(MetadataChangeProposalClass, next(iter(lineage)).metadata)

metadata-ingestion/tests/unit/sql_parsing/test_sql_parsing_result_utils.py

@@ -0,0 +1,67 @@
+from datahub.sql_parsing.sql_parsing_result_utils import (
+    transform_parsing_result_to_in_tables_schemas,
+)
+from datahub.sql_parsing.sqlglot_lineage import (
+    ColumnLineageInfo,
+    ColumnRef,
+    DownstreamColumnRef,
+    SqlParsingResult,
+)
+
+
+def test_transform_parsing_result_to_in_tables_schemas__empty_parsing_result():
+    parsing_result = SqlParsingResult(in_tables=[], out_tables=[], column_lineage=None)
+
+    in_tables_schema = transform_parsing_result_to_in_tables_schemas(parsing_result)
+    assert not in_tables_schema
+
+
+def test_transform_parsing_result_to_in_tables_schemas__in_tables_only():
+    parsing_result = SqlParsingResult(
+        in_tables=["table_urn1", "table_urn2", "table_urn3"],
+        out_tables=[],
+        column_lineage=None,
+    )
+
+    in_tables_schema = transform_parsing_result_to_in_tables_schemas(parsing_result)
+    assert in_tables_schema == {
+        "table_urn1": set(),
+        "table_urn2": set(),
+        "table_urn3": set(),
+    }
+
+
+def test_transform_parsing_result_to_in_tables_schemas__in_tables_and_column_linage():
+    parsing_result = SqlParsingResult(
+        in_tables=["table_urn1", "table_urn2", "table_urn3"],
+        out_tables=[],
+        column_lineage=[
+            ColumnLineageInfo(
+                downstream=DownstreamColumnRef(column="out_col1"),
+                upstreams=[
+                    ColumnRef(table="table_urn1", column="col11"),
+                ],
+            ),
+            ColumnLineageInfo(
+                downstream=DownstreamColumnRef(column="out_col2"),
+                upstreams=[
+                    ColumnRef(table="table_urn2", column="col21"),
+                    ColumnRef(table="table_urn2", column="col22"),
+                ],
+            ),
+            ColumnLineageInfo(
+                downstream=DownstreamColumnRef(column="out_col3"),
+                upstreams=[
+                    ColumnRef(table="table_urn1", column="col12"),
+                    ColumnRef(table="table_urn2", column="col23"),
+                ],
+            ),
+        ],
+    )
+
+    in_tables_schema = transform_parsing_result_to_in_tables_schemas(parsing_result)
+    assert in_tables_schema == {
+        "table_urn1": {"col11", "col12"},
+        "table_urn2": {"col21", "col22", "col23"},
+        "table_urn3": set(),
+    }

metadata-ingestion/tests/unit/test_tableau_source.py

@@ -0,0 +1,123 @@
+import pytest
+
+from datahub.ingestion.source.tableau import TableauSource
+
+
+def test_tableau_source_unescapes_lt():
+    res = TableauSource._clean_tableau_query_parameters(
+        "select * from t where c1 << 135"
+    )
+
+    assert res == "select * from t where c1 < 135"
+
+
+def test_tableau_source_unescapes_gt():
+    res = TableauSource._clean_tableau_query_parameters(
+        "select * from t where c1 >> 135"
+    )
+
+    assert res == "select * from t where c1 > 135"
+
+
+def test_tableau_source_unescapes_gte():
+    res = TableauSource._clean_tableau_query_parameters(
+        "select * from t where c1 >>= 135"
+    )
+
+    assert res == "select * from t where c1 >= 135"
+
+
+def test_tableau_source_unescapeslgte():
+    res = TableauSource._clean_tableau_query_parameters(
+        "select * from t where c1 <<= 135"
+    )
+
+    assert res == "select * from t where c1 <= 135"
+
+
+def test_tableau_source_doesnt_touch_not_escaped():
+    res = TableauSource._clean_tableau_query_parameters(
+        "select * from t where c1 < 135 and c2 > 15"
+    )
+
+    assert res == "select * from t where c1 < 135 and c2 > 15"
+
+
+TABLEAU_PARAMS = [
+    "<Parameters.MyParam>",
+    "<Parameters.MyParam_1>",
+    "<Parameters.My Param _ 1>",
+    "<Parameters.My Param 1 !@\"',.#$%^:;&*()-_+={}|\\ /<>",
+    "<[Parameters].MyParam>",
+    "<[Parameters].MyParam_1>",
+    "<[Parameters].My Param _ 1>",
+    "<[Parameters].My Param 1 !@\"',.#$%^:;&*()-_+={}|\\ /<>",
+    "<Parameters.[MyParam]>",
+    "<Parameters.[MyParam_1]>",
+    "<Parameters.[My Param _ 1]>",
+    "<Parameters.[My Param 1 !@\"',.#$%^:;&*()-_+={}|\\ /<]>",
+    "<[Parameters].[MyParam]>",
+    "<[Parameters].[MyParam_1]>",
+    "<[Parameters].[My Param _ 1]>",
+    "<[Parameters].[My Param 1 !@\"',.#$%^:;&*()-_+={}|\\ /<]>",
+    "<Parameters.[My Param 1 !@\"',.#$%^:;&*()-_+={}|\\ /<>]>",
+    "<[Parameters].[My Param 1 !@\"',.#$%^:;&*()-_+={}|\\ /<>]>",
+]
+
+
+@pytest.mark.parametrize("p", TABLEAU_PARAMS)
+def test_tableau_source_cleanups_tableau_parameters_in_equi_predicates(p):
+    assert (
+        TableauSource._clean_tableau_query_parameters(
+            f"select * from t where c1 = {p} and c2 = {p} and c3 = 7"
+        )
+        == "select * from t where c1 = 1 and c2 = 1 and c3 = 7"
+    )
+
+
+@pytest.mark.parametrize("p", TABLEAU_PARAMS)
+def test_tableau_source_cleanups_tableau_parameters_in_lt_gt_predicates(p):
+    assert (
+        TableauSource._clean_tableau_query_parameters(
+            f"select * from t where c1 << {p} and c2<<{p} and c3 >> {p} and c4>>{p} or {p} >> c1 and {p}>>c2 and {p} << c3 and {p}<<c4"
+        )
+        == "select * from t where c1 < 1 and c2<1 and c3 > 1 and c4>1 or 1 > c1 and 1>c2 and 1 < c3 and 1<c4"
+    )
+
+
+@pytest.mark.parametrize("p", TABLEAU_PARAMS)
+def test_tableau_source_cleanups_tableau_parameters_in_lte_gte_predicates(p):
+    assert (
+        TableauSource._clean_tableau_query_parameters(
+            f"select * from t where c1 <<= {p} and c2<<={p} and c3 >>= {p} and c4>>={p} or {p} >>= c1 and {p}>>=c2 and {p} <<= c3 and {p}<<=c4"
+        )
+        == "select * from t where c1 <= 1 and c2<=1 and c3 >= 1 and c4>=1 or 1 >= c1 and 1>=c2 and 1 <= c3 and 1<=c4"
+    )
+
+
+@pytest.mark.parametrize("p", TABLEAU_PARAMS)
+def test_tableau_source_cleanups_tableau_parameters_in_join_predicate(p):
+    assert (
+        TableauSource._clean_tableau_query_parameters(
+            f"select * from t1 inner join t2 on t1.id = t2.id and t2.c21 = {p} and t1.c11 = 123 + {p}"
+        )
+        == "select * from t1 inner join t2 on t1.id = t2.id and t2.c21 = 1 and t1.c11 = 123 + 1"
+    )
+
+
+@pytest.mark.parametrize("p", TABLEAU_PARAMS)
+def test_tableau_source_cleanups_tableau_parameters_in_complex_expressions(p):
+    assert (
+        TableauSource._clean_tableau_query_parameters(
+            f"select myudf1(c1, {p}, c2) / myudf2({p}) > ({p} + 3 * {p} * c5) * {p} - c4"
+        )
+        == "select myudf1(c1, 1, c2) / myudf2(1) > (1 + 3 * 1 * c5) * 1 - c4"
+    )
+
+
+@pytest.mark.parametrize("p", TABLEAU_PARAMS)
+def test_tableau_source_cleanups_tableau_parameters_in_udfs(p):
+    assert (
+        TableauSource._clean_tableau_query_parameters(f"select myudf({p}) from t")
+        == "select myudf(1) from t"
+    )