control-service: Extending Trino lineage

Fix lineage for templates.

---------

Signed-off-by: Tonka Zheleva <tonka.zheleva@broadcom.com>

Author: gorankokin

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/templates/load/dimension/scd1/03-handle-quality-checks_and_move_data.py

@@ -1,10 +1,13 @@
-# Copyright 2023-2024 Broadcom
+# Copyright 2023-2025 Broadcom
 # SPDX-License-Identifier: Apache-2.0
 import logging
 import os
 import re
 
 from vdk.api.job_input import IJobInput
+from vdk.api.lineage.model.logger.lineage_logger import ILineageLogger
+from vdk.api.lineage.model.sql.model import LineageData
+from vdk.internal.core.statestore import StoreKey
 from vdk.plugin.trino.templates.data_quality_exception import DataQualityException
 from vdk.plugin.trino.trino_utils import CommonUtilities
 
@@ -77,7 +80,13 @@ def run(job_input: IJobInput):
     # copy the data if there's no quality check configure or if it passes
     if not check or check(staging_table_full_name):
         copy_staging_table_to_target_table(
-            job_input, target_schema, target_table, staging_schema, staging_table
+            job_input,
+            target_schema,
+            target_table,
+            staging_schema,
+            staging_table,
+            source_schema,
+            source_view,
         )
     else:
         target_table_full_name = f"{target_schema}.{target_table}"
@@ -94,6 +103,8 @@ def copy_staging_table_to_target_table(
     target_table,
     source_schema,
     source_table,
+    original_source_schema,
+    original_source_view,
 ):
     # check if table is partitioned
     if is_partitioned_table(job_input, target_schema, target_table):
@@ -181,6 +192,20 @@ def copy_staging_table_to_target_table(
         )
         job_input.execute_query(insert_into_table)
 
+    lineage_data = LineageData(
+        query="template",
+        query_type="template",
+        query_status="OK",
+        input_tables=[original_source_schema + "." + original_source_view],
+        output_table=target_schema + "." + target_table,
+    )
+
+    LINEAGE_LOGGER_KEY = StoreKey[ILineageLogger]("trino-lineage-logger")
+    lineage_logger = job_input._JobInput__templates._core_context.state.get(
+        LINEAGE_LOGGER_KEY
+    )
+    lineage_logger.send(lineage_data)
+
 
 def is_partitioned_table(job_input: IJobInput, target_schema, target_table) -> bool:
     show_create_query = CommonUtilities.get_file_content(

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/templates/load/dimension/scd1_upsert/03-handle-quality-checks_and_move_data.py

@@ -1,10 +1,13 @@
-# Copyright 2023-2024 Broadcom
+# Copyright 2023-2025 Broadcom
 # SPDX-License-Identifier: Apache-2.0
 import logging
 import os
 import re
 
 from vdk.api.job_input import IJobInput
+from vdk.api.lineage.model.logger.lineage_logger import ILineageLogger
+from vdk.api.lineage.model.sql.model import LineageData
+from vdk.internal.core.statestore import StoreKey
 from vdk.plugin.trino.templates.data_quality_exception import DataQualityException
 from vdk.plugin.trino.trino_utils import CommonUtilities
 
@@ -78,7 +81,13 @@ def run(job_input: IJobInput):
     # copy the data if there's no quality check configure or if it passes
     if not check or check(staging_table_full_name):
         copy_staging_table_to_target_table(
-            job_input, target_schema, target_table, staging_schema, staging_table
+            job_input,
+            target_schema,
+            target_table,
+            staging_schema,
+            staging_table,
+            source_schema,
+            source_view,
         )
     else:
         target_table_full_name = f"{target_schema}.{target_table}"
@@ -95,6 +104,8 @@ def copy_staging_table_to_target_table(
     target_table,
     source_schema,
     source_table,
+    original_source_schema,
+    original_source_view,
 ):
     # non-partitioned tables:
     # - Since truncate and delete do not work for non-partitioned tables - get the create statement, drop the table and then re-create it - we do this to preserve and metadata like user comments
@@ -135,6 +146,20 @@ def copy_staging_table_to_target_table(
     )
     job_input.execute_query(insert_into_table)
 
+    lineage_data = LineageData(
+        query="template",
+        query_type="template",
+        query_status="OK",
+        input_tables=[original_source_schema + "." + original_source_view],
+        output_table=target_schema + "." + target_table,
+    )
+
+    LINEAGE_LOGGER_KEY = StoreKey[ILineageLogger]("trino-lineage-logger")
+    lineage_logger = job_input._JobInput__templates._core_context.state.get(
+        LINEAGE_LOGGER_KEY
+    )
+    lineage_logger.send(lineage_data)
+
 
 def remove_external_location(sql_statement):
     # Regular expression pattern to match the external_location clause

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/templates/load/fact/insert/02-handle-quality-checks_and_move_data.py

@@ -1,9 +1,12 @@
-# Copyright 2023-2024 Broadcom
+# Copyright 2023-2025 Broadcom
 # SPDX-License-Identifier: Apache-2.0
 import logging
 import os
 
 from vdk.api.job_input import IJobInput
+from vdk.api.lineage.model.logger.lineage_logger import ILineageLogger
+from vdk.api.lineage.model.sql.model import LineageData
+from vdk.internal.core.statestore import StoreKey
 from vdk.plugin.trino.templates.data_quality_exception import DataQualityException
 from vdk.plugin.trino.trino_utils import CommonUtilities
 
@@ -128,3 +131,17 @@ def run(job_input: IJobInput):
             )
     else:
         job_input.execute_query(insert_query)
+
+        lineage_data = LineageData(
+            query="template",
+            query_type="template",
+            query_status="OK",
+            input_tables=[source_schema + "." + source_view],
+            output_table=target_schema + "." + target_table,
+        )
+
+        LINEAGE_LOGGER_KEY = StoreKey[ILineageLogger]("trino-lineage-logger")
+        lineage_logger = job_input._JobInput__templates._core_context.state.get(
+            LINEAGE_LOGGER_KEY
+        )
+        lineage_logger.send(lineage_data)

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/templates/load/fact/snapshot/02-handle-quality-checks_and_move_data.py

@@ -1,9 +1,12 @@
-# Copyright 2023-2024 Broadcom
+# Copyright 2023-2025 Broadcom
 # SPDX-License-Identifier: Apache-2.0
 import logging
 import os
 
 from vdk.api.job_input import IJobInput
+from vdk.api.lineage.model.logger.lineage_logger import ILineageLogger
+from vdk.api.lineage.model.sql.model import LineageData
+from vdk.internal.core.statestore import StoreKey
 from vdk.plugin.trino.templates.data_quality_exception import DataQualityException
 from vdk.plugin.trino.trino_utils import CommonUtilities
 
@@ -116,6 +119,20 @@ def run(job_input: IJobInput):
                 source_view=staging_table,
             )
             job_input.execute_query(create_and_insert_into_target)
+
+            lineage_data = LineageData(
+                query="template",
+                query_type="template",
+                query_status="OK",
+                input_tables=[source_schema + "." + source_view],
+                output_table=target_schema + "." + target_table,
+            )
+
+            LINEAGE_LOGGER_KEY = StoreKey[ILineageLogger]("trino-lineage-logger")
+            lineage_logger = job_input._JobInput__templates._core_context.state.get(
+                LINEAGE_LOGGER_KEY
+            )
+            lineage_logger.send(lineage_data)
         else:
             log.info(
                 f"Target table {target_schema}.{target_table} remains unchanged "