feat(ingestion/dbt): add support for delta lake table_changes table valued function as DBT source

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

@@ -1,11 +1,16 @@
 import functools
+import re
 from typing import Any, Optional
 
 import sqlglot
 from pydantic import BaseModel
 
 from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
 from datahub.metadata.schema_classes import SchemaFieldDataTypeClass
+from datahub.sql_parsing.sql_parsing_common import DIALECTS_WITH_DELTALAKE_TVF_SUPPORT
+from datahub.sql_parsing.sqlglot_utils import (
+    is_dialect_instance,
+)
 
 
 class _ParserBaseModel(
@@ -78,6 +83,7 @@
         table: sqlglot.exp.Table,
         default_db: Optional[str] = None,
         default_schema: Optional[str] = None,
+        dialect: sqlglot.Dialect = None,
     ) -> "_TableName":
         if isinstance(table.this, sqlglot.exp.Dot):
             # For tables that are more than 3 parts, the extra parts will be in a Dot.
@@ -89,6 +95,12 @@
                 exp = exp.expression
             parts.append(exp.name)
             table_name = ".".join(parts)
+        elif (
+            is_dialect_instance(dialect, DIALECTS_WITH_DELTALAKE_TVF_SUPPORT)
+            and table.this.name == "table_changes"
+            and re.match(r"TABLE_CHANGES\s*\([^)]*\)", table.this)
+        ):
+            table_name = table.this.expressions[0].alias_or_name
         else:
             table_name = table.this.name
         return cls(

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

@@ -31,6 +31,9 @@
     # automatically lowercase unquoted identifiers.
     "snowflake",
 }
+# To support delta lake functions
+DIALECTS_WITH_DELTALAKE_TVF_SUPPORT = {"hive", "spark"}
+
 assert DIALECTS_WITH_DEFAULT_UPPERCASE_COLS.issubset(
     DIALECTS_WITH_CASE_INSENSITIVE_COLS
 )

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

@@ -203,7 +203,7 @@ def _table_level_lineage(
     # Generate table-level lineage.
     modified = (
         {
-            _TableName.from_sqlglot_table(expr.this)
+            _TableName.from_sqlglot_table(table=expr.this, dialect=dialect)
             for expr in statement.find_all(
                 sqlglot.exp.Create,
                 sqlglot.exp.Insert,
@@ -220,7 +220,7 @@ def _table_level_lineage(
             # For statements that include a column list, like
             # CREATE DDL statements and `INSERT INTO table (col1, col2) SELECT ...`
             # the table name is nested inside a Schema object.
-            _TableName.from_sqlglot_table(expr.this.this)
+            _TableName.from_sqlglot_table(table=expr.this.this, dialect=dialect)
             for expr in statement.find_all(
                 sqlglot.exp.Create,
                 sqlglot.exp.Insert,
@@ -231,7 +231,7 @@ def _table_level_lineage(
         | {
             # For drop statements, we only want it if a table/view is being dropped.
             # Other "kinds" will not have table.name populated.
-            _TableName.from_sqlglot_table(expr.this)
+            _TableName.from_sqlglot_table(table=expr.this, dialect=dialect)
             for expr in ([statement] if isinstance(statement, sqlglot.exp.Drop) else [])
             if isinstance(expr.this, sqlglot.exp.Table)
             and expr.this.this
@@ -241,7 +241,7 @@ def _table_level_lineage(
 
     tables = (
         {
-            _TableName.from_sqlglot_table(table)
+            _TableName.from_sqlglot_table(table=table, dialect=dialect)
             for table in statement.find_all(sqlglot.exp.Table)
             if not isinstance(table.parent, sqlglot.exp.Drop)
         }
@@ -518,7 +518,9 @@ def _select_statement_cll(  # noqa: C901
             # )
 
             # Generate SELECT lineage.
-            direct_raw_col_upstreams = _get_direct_raw_col_upstreams(lineage_node)
+            direct_raw_col_upstreams = _get_direct_raw_col_upstreams(
+                lineage_node, dialect
+            )
 
             # column_logic = lineage_node.source
 
@@ -652,7 +654,7 @@ def _column_level_lineage(
 
 
 def _get_direct_raw_col_upstreams(
-    lineage_node: sqlglot.lineage.Node,
+    lineage_node: sqlglot.lineage.Node, dialect: sqlglot.Dialect
 ) -> Set[_ColumnRef]:
     # Using a set here to deduplicate upstreams.
     direct_raw_col_upstreams: Set[_ColumnRef] = set()
@@ -663,7 +665,9 @@ def _get_direct_raw_col_upstreams(
             pass
 
         elif isinstance(node.expression, sqlglot.exp.Table):
-            table_ref = _TableName.from_sqlglot_table(node.expression)
+            table_ref = _TableName.from_sqlglot_table(
+                table=node.expression, dialect=dialect
+            )
 
             if node.name == "*":
                 # This will happen if we couldn't expand the * to actual columns e.g. if