vdk-plugins: Fix Trino Lineage (#3428)

vdk-plugins: Fix Trino Lineage First collect lineage for Trino if configured, then execute the statement, otherwise 'explain create as select' fails if the table exists. --------- Signed-off-by: Tonka Zheleva <[email protected]> Signed-off-by: Tonka Zheleva <[email protected]> Co-authored-by: Goran Kokinovski <[email protected]>
vmware · Nov 8, 2024 · 9a01569 · 9a01569
1 parent 295092d
commit 9a01569
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/trino_connection.py b/projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/trino_connection.py
@@ -159,11 +159,12 @@ def db_connection_recover_operation(self, recovery_cursor: RecoveryCursor) -> No
             raise recovery_cursor.get_exception()
 
     def execute_query(self, query):
-        res = self.execute_query_with_retries(query)
+        # first evaluate lineage because current behavior of 'explain create as select' fails if the table exists
         if self._lineage_logger:
             lineage_data = self._get_lineage_data(query)
-            if lineage_data:
-                self._lineage_logger.send(lineage_data)
+        res = self.execute_query_with_retries(query)
+        if self._lineage_logger and lineage_data:
+            self._lineage_logger.send(lineage_data)
         #  TODO: collect lineage for failed query
         return res