From 667845e74c9ffebea197b77b8eaeb476f02911f6 Mon Sep 17 00:00:00 2001
From: ulixius9 <mayursingal9@gmail.com>
Date: Wed, 4 Dec 2024 23:19:17 +0530
Subject: [PATCH] MINOR: Tableau Capture SQL in Lineage

---
 .../source/dashboard/dashboard_service.py     |  2 ++
 .../source/dashboard/tableau/metadata.py      | 28 +++++++++++++------
 .../source/dashboard/tableau/models.py        | 10 +++++++
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py
index 96df9f912f78..80a8895288f0 100644
--- a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py
+++ b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py
@@ -479,6 +479,7 @@ def _get_add_lineage_request(
         to_entity: Union[Dashboard, DashboardDataModel],
         from_entity: Union[Table, DashboardDataModel, Dashboard],
         column_lineage: List[ColumnLineage] = None,
+        sql: Optional[str] = None,
     ) -> Optional[Either[AddLineageRequest]]:
         if from_entity and to_entity:
             return Either(
@@ -494,6 +495,7 @@ def _get_add_lineage_request(
                         ),
                         lineageDetails=LineageDetails(
                             source=LineageSource.DashboardLineage,
+                            sqlQuery=sql,
                             columnsLineage=column_lineage,
                         ),
                     )
diff --git a/ingestion/src/metadata/ingestion/source/dashboard/tableau/metadata.py b/ingestion/src/metadata/ingestion/source/dashboard/tableau/metadata.py
index 1d0887ed037c..39626bbfde13 100644
--- a/ingestion/src/metadata/ingestion/source/dashboard/tableau/metadata.py
+++ b/ingestion/src/metadata/ingestion/source/dashboard/tableau/metadata.py
@@ -69,6 +69,7 @@
     ChartUrl,
     DataSource,
     DatasourceField,
+    TableAndQuery,
     TableauDashboard,
     TableauTag,
     UpstreamTable,
@@ -149,7 +150,7 @@ def get_owner_ref(
         return None
 
     @staticmethod
-    def _get_data_models_tags(dataModels: [DataSource]) -> Set[str]:
+    def _get_data_models_tags(dataModels: List[DataSource]) -> Set[str]:
         """
         Get the tags from the data model in the upstreamDatasources
         """
@@ -454,14 +455,18 @@ def _get_table_datamodel_lineage(
             }
             for table in datamodel.upstreamTables or []:
                 om_tables = self._get_database_tables(db_service_entity, table)
-                for om_table in om_tables or []:
+                for om_table_and_query in om_tables or []:
                     column_lineage = self._get_column_lineage(
-                        table, om_table, upstream_data_model_entity, upstream_col_set
+                        table,
+                        om_table_and_query.table,
+                        upstream_data_model_entity,
+                        upstream_col_set,
                     )
                     yield self._get_add_lineage_request(
                         to_entity=upstream_data_model_entity,
-                        from_entity=om_table,
+                        from_entity=om_table_and_query.table,
                         column_lineage=column_lineage,
+                        sql=om_table_and_query.query,
                     )
         except Exception as err:
             yield Either(
@@ -698,7 +703,7 @@ def close(self):
 
     def _get_table_entities_from_api(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         In case we get the table details from the Graphql APIs we process them
         """
@@ -734,7 +739,7 @@ def _get_table_entities_from_api(
                     fqn=table_fqn,
                 )
                 if table_entity:
-                    return [table_entity]
+                    return [TableAndQuery(table=table_entity)]
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Error to get tables for lineage using GraphQL Apis: {exc}")
@@ -742,7 +747,7 @@ def _get_table_entities_from_api(
 
     def _get_table_entities_from_query(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         In case we get the table details from the Graphql APIs we process them
         """
@@ -778,7 +783,12 @@ def _get_table_entities_from_query(
                         database_schema=schema_name,
                         table=table_name,
                     )
-                    tables_list.extend(from_entities)
+                    tables_list.extend(
+                        [
+                            TableAndQuery(table=table, query=custom_sql_table.query)
+                            for table in from_entities
+                        ]
+                    )
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
@@ -787,7 +797,7 @@ def _get_table_entities_from_query(
 
     def _get_database_tables(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         Get the table entities for lineage
         """
diff --git a/ingestion/src/metadata/ingestion/source/dashboard/tableau/models.py b/ingestion/src/metadata/ingestion/source/dashboard/tableau/models.py
index 10837cac04f4..307b74ccb32f 100644
--- a/ingestion/src/metadata/ingestion/source/dashboard/tableau/models.py
+++ b/ingestion/src/metadata/ingestion/source/dashboard/tableau/models.py
@@ -18,6 +18,7 @@
 from pydantic import BaseModel, ConfigDict, Field, validator
 
 from metadata.generated.schema.entity.data.chart import ChartType
+from metadata.generated.schema.entity.data.table import Table
 
 
 class TableauBaseModel(BaseModel):
@@ -172,3 +173,12 @@ class TableauDashboard(TableauBaseModel):
     webpageUrl: Optional[str] = None
     charts: Optional[List[TableauChart]] = None
     dataModels: List[DataSource] = []
+
+
+class TableAndQuery(BaseModel):
+    """
+    Wrapper class for Table entity and associated Query for lineage
+    """
+
+    table: Table
+    query: Optional[str] = None