control-service: Extending Trino lineage

Add lineage for "CREATE AS SELECT" --------- Signed-off-by: Tonka Zheleva <[email protected]>
vmware · Oct 10, 2024 · 984fa49 · 984fa49
1 parent 3d97f64
commit 984fa49
Showing 1 changed file with 7 additions and 1 deletion.
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
@@ -196,7 +196,13 @@ def _get_lineage_data(self, query):
                     "ALTER operation not a RENAME TABLE operation. No lineage will be collected."
                 )
 
-        elif statement.get_type() == "SELECT" or statement.get_type() == "INSERT":
+        elif (
+            statement.get_type() == "SELECT"
+            or statement.get_type() == "INSERT"
+            or (
+                statement.get_type() == "CREATE" and "select" in statement.value.lower()
+            )
+        ):
             if lineage_utils.is_heartbeat_query(query):
                 return None
             log.debug("Collecting lineage for SELECT/INSERT query ...")