From 727cb5d8cda40f1ac4b8220626df6117c51be1f6 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Thu, 27 Oct 2022 14:57:53 +0300
Subject: [PATCH 01/14] XCOM push ORA error code in OracleStoredProcedure

Motivation:
    https://github.com/apache/airflow/discussions/23787

ORA exit codes is controlled by the DB developer and could be a part of a dag logic.
So in order to get access to these codes in dag, I propose an easy enhancement - push ORA exit code to XCOM when DatabaseError occurs.

One file changed - providers/oracle/operators/oracle.py
---
 airflow/providers/oracle/operators/oracle.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index fac193b4f8324..f5b1336b6662a 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -26,7 +26,9 @@
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
-
+    
+import oracledb
+from airflow.exceptions import AirflowException
 
 class OracleOperator(SQLExecuteQueryOperator):
     """
@@ -93,4 +95,10 @@ def __init__(
     def execute(self, context: Context):
         self.log.info("Executing: %s", self.procedure)
         hook = OracleHook(oracle_conn_id=self.oracle_conn_id)
-        return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
+        ti=context.get('task_instance')
+        try:
+            return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
+        except oracledb.DatabaseError as e:
+            code, mesg = e.args[0].message[:-1].split(': ', 1)
+            ti.xcom_push(key='ORA', value=str(code.split('-')[1]))
+            raise AirflowException(f"Task failed with exit code: {code}")

From 741299f84e2e7e2fd46abd64d253d08143697e14 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Fri, 28 Oct 2022 20:14:11 +0300
Subject: [PATCH 02/14] raise original instead of raise AirflowException

https://github.com/apache/airflow/pull/27319/files/daf4ddf632a6e9d2fcf58982a697c4d1e031364b#r1007791853
---
 airflow/providers/oracle/operators/oracle.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index f5b1336b6662a..87388d641ec95 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -28,7 +28,6 @@
     from airflow.utils.context import Context
     
 import oracledb
-from airflow.exceptions import AirflowException
 
 class OracleOperator(SQLExecuteQueryOperator):
     """
@@ -101,4 +100,4 @@ def execute(self, context: Context):
         except oracledb.DatabaseError as e:
             code, mesg = e.args[0].message[:-1].split(': ', 1)
             ti.xcom_push(key='ORA', value=str(code.split('-')[1]))
-            raise AirflowException(f"Task failed with exit code: {code}")
+            raise

From b57ea19d4a710e20e08b7e6e40bd1936fef77c13 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Tue, 1 Nov 2022 21:54:05 +0300
Subject: [PATCH 03/14] Update oracle.py

---
 airflow/providers/oracle/operators/oracle.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 87388d641ec95..7ade2882cfe51 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -99,5 +99,6 @@ def execute(self, context: Context):
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
             code, mesg = e.args[0].message[:-1].split(': ', 1)
-            ti.xcom_push(key='ORA', value=str(code.split('-')[1]))
+            if self.do_xcom_push is True:
+                ti.xcom_push(key='ORA', value=str(code.split('-')[1]))
             raise

From f53a28d1526de64b95785a812caa753261b16811 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Wed, 2 Nov 2022 12:41:55 +0300
Subject: [PATCH 04/14] added double quotes instaed of single quotes

---
 airflow/providers/oracle/operators/oracle.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 7ade2882cfe51..7fe164a3539b5 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -94,11 +94,11 @@ def __init__(
     def execute(self, context: Context):
         self.log.info("Executing: %s", self.procedure)
         hook = OracleHook(oracle_conn_id=self.oracle_conn_id)
-        ti=context.get('task_instance')
+        ti=context.get("task_instance")
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
-            code, mesg = e.args[0].message[:-1].split(': ', 1)
+            code, mesg = e.args[0].message[:-1].split(": ", 1)
             if self.do_xcom_push is True:
-                ti.xcom_push(key='ORA', value=str(code.split('-')[1]))
+                ti.xcom_push(key="ORA", value=str(code.split("-")[1]))
             raise

From 26e60d174a8d698d7e94b1939f8d0cb2a9c04d54 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Thu, 3 Nov 2022 03:39:29 +0300
Subject: [PATCH 05/14] simplify do_xcom_push check

Co-authored-by: Dov Benyomin Sohacheski <b@kloud.email>
---
 airflow/providers/oracle/operators/oracle.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 7fe164a3539b5..e9151e0366821 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -99,6 +99,6 @@ def execute(self, context: Context):
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
             code, mesg = e.args[0].message[:-1].split(": ", 1)
-            if self.do_xcom_push is True:
+            if self.do_xcom_push:
                 ti.xcom_push(key="ORA", value=str(code.split("-")[1]))
             raise

From fc9050f5aeecb29aa051a236d8ebc61ffbab28bf Mon Sep 17 00:00:00 2001
From: kolfild26 <akomrakov26@gmail.com>
Date: Mon, 7 Nov 2022 00:16:27 +0300
Subject: [PATCH 06/14] re.search istead of split

---
 airflow/providers/oracle/operators/oracle.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index e9151e0366821..faf472fefae1b 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -19,6 +19,8 @@
 
 import warnings
 from typing import TYPE_CHECKING, Sequence
+import oracledb
+import re
 
 from airflow.models import BaseOperator
 from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
@@ -26,8 +28,7 @@
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
-    
-import oracledb
+
 
 class OracleOperator(SQLExecuteQueryOperator):
     """
@@ -70,6 +71,9 @@ class OracleStoredProcedureOperator(BaseOperator):
     :param oracle_conn_id: The :ref:`Oracle connection id <howto/connection:oracle>`
         reference to a specific Oracle database.
     :param parameters: (optional, templated) the parameters provided in the call
+
+    :xcom ORA: Oracle db exit code. Optional (if self.do_xcom_push is True).
+        Numeric exit code that Oracle throws in case of failure.
     """
 
     template_fields: Sequence[str] = (
@@ -98,7 +102,7 @@ def execute(self, context: Context):
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
-            code, mesg = e.args[0].message[:-1].split(": ", 1)
+            code = re.search("^ORA-(\\d+):.+", e).group(1)
             if self.do_xcom_push:
-                ti.xcom_push(key="ORA", value=str(code.split("-")[1]))
+                ti.xcom_push(key="ORA", value=code)
             raise

From 6c5f3b509cb2c1432682e46f808fc45d306523af Mon Sep 17 00:00:00 2001
From: kolfild26 <akomrakov26@gmail.com>
Date: Mon, 7 Nov 2022 01:42:29 +0300
Subject: [PATCH 07/14] handing ORA- not found in the begining of returned
 message

---
 airflow/providers/oracle/operators/oracle.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index faf472fefae1b..4147cf013db56 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -102,7 +102,10 @@ def execute(self, context: Context):
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
-            code = re.search("^ORA-(\\d+):.+", e).group(1)
-            if self.do_xcom_push:
-                ti.xcom_push(key="ORA", value=code)
+            try:
+                code = re.search("^ORA-(\\d+):.+", str(e)).group(1)
+                if self.do_xcom_push:
+                    ti.xcom_push(key="ORA", value=code)
+            except AttributeError:
+                pass
             raise

From b5428b90dec031821331830fdd78a2890943708f Mon Sep 17 00:00:00 2001
From: kolfild26 <akomrakov26@gmail.com>
Date: Mon, 7 Nov 2022 02:27:09 +0300
Subject: [PATCH 08/14] fixed mypy Handle Optionals

---
 airflow/providers/oracle/operators/oracle.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 4147cf013db56..4f88914eaec9e 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -102,10 +102,9 @@ def execute(self, context: Context):
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
-            try:
-                code = re.search("^ORA-(\\d+):.+", str(e)).group(1)
+            code_match = re.search("^ORA-(\\d+):.+", str(e))
+            if code_match is not None:
                 if self.do_xcom_push:
-                    ti.xcom_push(key="ORA", value=code)
-            except AttributeError:
-                pass
+                    if ti is not None:
+                        ti.xcom_push(key="ORA", value=code_match.group(1))
             raise

From 1b8dca8a47b4bf2a65559a2269f82cb96795ff7b Mon Sep 17 00:00:00 2001
From: kolfild26 <akomrakov26@gmail.com>
Date: Mon, 7 Nov 2022 13:48:40 +0300
Subject: [PATCH 09/14] simplify xcom_push ablility check

---
 airflow/providers/oracle/operators/oracle.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 4f88914eaec9e..d7d66ad0a4f09 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -17,10 +17,11 @@
 # under the License.
 from __future__ import annotations
 
+import re
 import warnings
 from typing import TYPE_CHECKING, Sequence
+
 import oracledb
-import re
 
 from airflow.models import BaseOperator
 from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
@@ -98,13 +99,11 @@ def __init__(
     def execute(self, context: Context):
         self.log.info("Executing: %s", self.procedure)
         hook = OracleHook(oracle_conn_id=self.oracle_conn_id)
-        ti=context.get("task_instance")
+        ti = context.get("task_instance")
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
             code_match = re.search("^ORA-(\\d+):.+", str(e))
-            if code_match is not None:
-                if self.do_xcom_push:
-                    if ti is not None:
-                        ti.xcom_push(key="ORA", value=code_match.group(1))
+            if code_match and self.do_xcom_push and ti:
+                ti.xcom_push(key="ORA", value=code_match.group(1))
             raise

From a8607a8ffd3f471233361e8ad96891e56b978612 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Mon, 7 Nov 2022 16:05:53 +0300
Subject: [PATCH 10/14] save the regex match unless needed

https://github.com/apache/airflow/pull/27319#discussion_r1015387980

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/providers/oracle/operators/oracle.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index d7d66ad0a4f09..d1c65700c17bf 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -103,7 +103,9 @@ def execute(self, context: Context):
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
+            if not self.do_xcom_push or not ti:
+                raise
             code_match = re.search("^ORA-(\\d+):.+", str(e))
-            if code_match and self.do_xcom_push and ti:
+            if code_match:
                 ti.xcom_push(key="ORA", value=code_match.group(1))
             raise

From 778133de1ea2415b72b719d4dc3b5745a36067c2 Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Mon, 7 Nov 2022 21:15:26 +0800
Subject: [PATCH 11/14] Docstring format

---
 airflow/providers/oracle/operators/oracle.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index d1c65700c17bf..5c78293533774 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -73,8 +73,8 @@ class OracleStoredProcedureOperator(BaseOperator):
         reference to a specific Oracle database.
     :param parameters: (optional, templated) the parameters provided in the call
 
-    :xcom ORA: Oracle db exit code. Optional (if self.do_xcom_push is True).
-        Numeric exit code that Oracle throws in case of failure.
+    If *do_xcom_push* is *True*, the numeric exit code emitted by
+    the database is pused to XCom under key ``ORA`` in case of failure.
     """
 
     template_fields: Sequence[str] = (

From cc6807ff7f22b8a9c25d6417cd0777d7d8f6aa32 Mon Sep 17 00:00:00 2001
From: Andrey Komrakov <akomrakov26@gmail.com>
Date: Tue, 15 Nov 2022 01:17:13 +0300
Subject: [PATCH 12/14] spelling mistake fixed 'pused' to 'pushed'

Co-authored-by: Dov Benyomin Sohacheski <b@kloud.email>
---
 airflow/providers/oracle/operators/oracle.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 5c78293533774..62ffa73d66747 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -74,7 +74,7 @@ class OracleStoredProcedureOperator(BaseOperator):
     :param parameters: (optional, templated) the parameters provided in the call
 
     If *do_xcom_push* is *True*, the numeric exit code emitted by
-    the database is pused to XCom under key ``ORA`` in case of failure.
+    the database is pushed to XCom under key ``ORA`` in case of failure.
     """
 
     template_fields: Sequence[str] = (

From 47ad70a1de26c1f016312f4dd2bb89f9603891e3 Mon Sep 17 00:00:00 2001
From: kolfild26 <akomrakov26@gmail.com>
Date: Mon, 12 Dec 2022 01:15:44 +0300
Subject: [PATCH 13/14] test_push_oracle_exit_to_xcom added

---
 .../providers/oracle/operators/test_oracle.py | 64 ++++++++++++++++++-
 1 file changed, 62 insertions(+), 2 deletions(-)

diff --git a/tests/providers/oracle/operators/test_oracle.py b/tests/providers/oracle/operators/test_oracle.py
index 2e5a8e10e63dc..7aba59a968fcf 100644
--- a/tests/providers/oracle/operators/test_oracle.py
+++ b/tests/providers/oracle/operators/test_oracle.py
@@ -16,13 +16,53 @@
 # under the License.
 from __future__ import annotations
 
+from random import randrange
 from unittest import mock
 
+import oracledb
+import pendulum
 import pytest
 
+from airflow.models import DAG, DagModel, DagRun, TaskInstance
 from airflow.providers.common.sql.hooks.sql import fetch_all_handler
 from airflow.providers.oracle.hooks.oracle import OracleHook
 from airflow.providers.oracle.operators.oracle import OracleOperator, OracleStoredProcedureOperator
+from airflow.utils.session import create_session
+from airflow.utils.timezone import datetime
+from airflow.utils.types import DagRunType
+
+DEFAULT_DATE = datetime(2017, 1, 1)
+
+
+def create_context(task, persist_to_db=False, map_index=None):
+    if task.has_dag():
+        dag = task.dag
+    else:
+        dag = DAG(dag_id="dag", start_date=pendulum.now())
+        dag.add_task(task)
+    dag_run = DagRun(
+        run_id=DagRun.generate_run_id(DagRunType.MANUAL, DEFAULT_DATE),
+        run_type=DagRunType.MANUAL,
+        dag_id=dag.dag_id,
+    )
+    task_instance = TaskInstance(task=task, run_id=dag_run.run_id)
+    task_instance.dag_run = dag_run
+    if map_index is not None:
+        task_instance.map_index = map_index
+    if persist_to_db:
+        with create_session() as session:
+            session.add(DagModel(dag_id=dag.dag_id))
+            session.add(dag_run)
+            session.add(task_instance)
+            session.commit()
+    return {
+        "dag": dag,
+        "ts": DEFAULT_DATE.isoformat(),
+        "task": task,
+        "ti": task_instance,
+        "task_instance": task_instance,
+        "run_id": "test",
+    }
 
 
 class TestOracleOperator:
@@ -56,11 +96,11 @@ def test_execute(self, mock_get_db_hook):
 
 class TestOracleStoredProcedureOperator:
     @mock.patch.object(OracleHook, "run", autospec=OracleHook.run)
-    def test_execute(self, mock_run):
+    def test_execute(self, mock_run, **context):
         procedure = "test"
         oracle_conn_id = "oracle_default"
         parameters = {"parameter": "value"}
-        context = "test_context"
+        # context = "test_context"
         task_id = "test_task_id"
 
         operator = OracleStoredProcedureOperator(
@@ -78,3 +118,23 @@ def test_execute(self, mock_run):
             parameters=parameters,
             handler=mock.ANY,
         )
+
+    @mock.patch.object(OracleHook, "callproc", autospec=OracleHook.callproc)
+    def test_push_oracle_exit_to_xcom(self, mock_callproc):
+        # Test pulls the value previously pushed to xcom and checks if it's the same
+        procedure = "test_push"
+        oracle_conn_id = "oracle_default"
+        parameters = {"parameter": "value"}
+        task_id = "test_push"
+        ora_exit_code = "%05d" % randrange(10**5)
+        task = OracleStoredProcedureOperator(
+            procedure=procedure, oracle_conn_id=oracle_conn_id, parameters=parameters, task_id=task_id
+        )
+        context = create_context(task, persist_to_db=True)
+        mock_callproc.side_effect = oracledb.DatabaseError(
+            "ORA-" + ora_exit_code + ": This is a five-digit ORA error code"
+        )
+        try:
+            task.execute(context=context)
+        except oracledb.DatabaseError:
+            assert task.xcom_pull(key="ORA", context=context, task_ids=[task_id])[0] == ora_exit_code

From f786bffd2485756f39c4751b309d249b449e5b80 Mon Sep 17 00:00:00 2001
From: kolfild26 <akomrakov26@gmail.com>
Date: Mon, 12 Dec 2022 02:02:48 +0300
Subject: [PATCH 14/14] fix execute method: context.get(task_instance) ->
 context[ti]

---
 airflow/providers/oracle/operators/oracle.py    | 4 ++--
 tests/providers/oracle/operators/test_oracle.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/oracle/operators/oracle.py b/airflow/providers/oracle/operators/oracle.py
index 62ffa73d66747..b43787879f2ac 100644
--- a/airflow/providers/oracle/operators/oracle.py
+++ b/airflow/providers/oracle/operators/oracle.py
@@ -99,12 +99,12 @@ def __init__(
     def execute(self, context: Context):
         self.log.info("Executing: %s", self.procedure)
         hook = OracleHook(oracle_conn_id=self.oracle_conn_id)
-        ti = context.get("task_instance")
         try:
             return hook.callproc(self.procedure, autocommit=True, parameters=self.parameters)
         except oracledb.DatabaseError as e:
-            if not self.do_xcom_push or not ti:
+            if not self.do_xcom_push or not context:
                 raise
+            ti = context["ti"]
             code_match = re.search("^ORA-(\\d+):.+", str(e))
             if code_match:
                 ti.xcom_push(key="ORA", value=code_match.group(1))
diff --git a/tests/providers/oracle/operators/test_oracle.py b/tests/providers/oracle/operators/test_oracle.py
index 7aba59a968fcf..debf920f7925d 100644
--- a/tests/providers/oracle/operators/test_oracle.py
+++ b/tests/providers/oracle/operators/test_oracle.py
@@ -96,11 +96,11 @@ def test_execute(self, mock_get_db_hook):
 
 class TestOracleStoredProcedureOperator:
     @mock.patch.object(OracleHook, "run", autospec=OracleHook.run)
-    def test_execute(self, mock_run, **context):
+    def test_execute(self, mock_run):
         procedure = "test"
         oracle_conn_id = "oracle_default"
         parameters = {"parameter": "value"}
-        # context = "test_context"
+        context = "test_context"
         task_id = "test_task_id"
 
         operator = OracleStoredProcedureOperator(
