From 945cd35148d4b0ddb2eb697efe84b1101ec43678 Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Wed, 15 May 2024 14:46:57 +0530
Subject: [PATCH] MINOR: Fix Oracle SP Lineage for begin...end SP call (#16240)

---
 .../source/database/oracle/queries.py         |  5 ++-
 .../src/metadata/utils/stored_procedures.py   |  2 +-
 .../unit/utils/test_stored_procedures.py      | 42 +++++++++++++++++++
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/source/database/oracle/queries.py b/ingestion/src/metadata/ingestion/source/database/oracle/queries.py
index 2cddcc4bbfb..6464bf0a3e7 100644
--- a/ingestion/src/metadata/ingestion/source/database/oracle/queries.py
+++ b/ingestion/src/metadata/ingestion/source/database/oracle/queries.py
@@ -97,7 +97,7 @@ WITH SP_HISTORY AS (SELECT
     TO_TIMESTAMP(LAST_LOAD_TIME, 'YYYY-MM-DD HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000, 'SECOND') AS end_time,
     PARSING_SCHEMA_NAME as user_name
   FROM gv$sql
-  WHERE UPPER(sql_text) LIKE 'CALL%%'
+  WHERE UPPER(sql_text) LIKE '%%CALL%%' or UPPER(sql_text) LIKE '%%BEGIN%%'
   AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'YYYY-MM-DD HH24:MI:SS') >= TO_TIMESTAMP('{start_date}', 'YYYY-MM-DD HH24:MI:SS')
  ),
  Q_HISTORY AS (SELECT
@@ -114,7 +114,7 @@ WITH SP_HISTORY AS (SELECT
       PARSING_SCHEMA_NAME AS SCHEMA_NAME,
       NULL AS DATABASE_NAME
     FROM gv$sql
-    WHERE UPPER(sql_text) NOT LIKE '%CALL%'
+    WHERE UPPER(sql_text) NOT LIKE '%%CALL%%' AND UPPER(sql_text) NOT LIKE '%%BEGIN%%'
       AND SQL_FULLTEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
       AND SQL_FULLTEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'
       AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'YYYY-MM-DD HH24:MI:SS') 
@@ -135,6 +135,7 @@ JOIN Q_HISTORY Q
   ON Q.start_time between SP.start_time and SP.end_time
   AND Q.end_time between SP.start_time and SP.end_time
   AND Q.user_name = SP.user_name
+  AND Q.QUERY_TYPE <> 'SELECT'
 ORDER BY PROCEDURE_START_TIME DESC
 """
 )
diff --git a/ingestion/src/metadata/utils/stored_procedures.py b/ingestion/src/metadata/utils/stored_procedures.py
index 6cd764d5c7f..a8f02eddd85 100644
--- a/ingestion/src/metadata/utils/stored_procedures.py
+++ b/ingestion/src/metadata/utils/stored_procedures.py
@@ -19,7 +19,7 @@ from metadata.utils.logger import utils_logger
 
 logger = utils_logger()
 
-NAME_PATTERN = r"(?<=call)(.*)(?=\()"
+NAME_PATTERN = r"(?<=call)(.*)(?=\()|(?<=begin)(.*)(?=\()|(?<=begin)(.*)(?=;\s*end)"
 
 
 def get_procedure_name_from_call(
diff --git a/ingestion/tests/unit/utils/test_stored_procedures.py b/ingestion/tests/unit/utils/test_stored_procedures.py
index 65d59e1799c..a94a0ec6d24 100644
--- a/ingestion/tests/unit/utils/test_stored_procedures.py
+++ b/ingestion/tests/unit/utils/test_stored_procedures.py
@@ -49,6 +49,48 @@ class StoredProceduresTests(TestCase):
             "procedure_name",
         )
 
+        self.assertEqual(
+            get_procedure_name_from_call(
+                query_text="BEGIN DB.SCHEMA.PROCEDURE_NAME; END;",
+            ),
+            "procedure_name",
+        )
+
+        self.assertEqual(
+            get_procedure_name_from_call(
+                query_text="BEGIN schema.procedure_name; END;",
+            ),
+            "procedure_name",
+        )
+
+        self.assertEqual(
+            get_procedure_name_from_call(
+                query_text="BEGIN procedure_name; END;",
+            ),
+            "procedure_name",
+        )
+
+        self.assertEqual(
+            get_procedure_name_from_call(
+                query_text="BEGIN DB.SCHEMA.PROCEDURE_NAME(...); END;",
+            ),
+            "procedure_name",
+        )
+
+        self.assertEqual(
+            get_procedure_name_from_call(
+                query_text="BEGIN schema.procedure_name(...); END;",
+            ),
+            "procedure_name",
+        )
+
+        self.assertEqual(
+            get_procedure_name_from_call(
+                query_text="BEGIN procedure_name(...); END;",
+            ),
+            "procedure_name",
+        )
+
         self.assertIsNone(
             get_procedure_name_from_call(
                 query_text="something very random",