This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b205a67afe Handler parameter from `JdbcOperator` to `JdbcHook.run`
(#23817)
b205a67afe is described below
commit b205a67afe6b23d2863bdf5f657dd235d9a901d8
Author: Dmytro Kazanzhy <[email protected]>
AuthorDate: Sun May 22 02:59:22 2022 +0300
Handler parameter from `JdbcOperator` to `JdbcHook.run` (#23817)
---
airflow/providers/jdbc/operators/jdbc.py | 8 +++++++-
tests/providers/jdbc/operators/test_jdbc.py | 7 +++++--
2 files changed, 12 insertions(+), 3 deletions(-)
diff --git a/airflow/providers/jdbc/operators/jdbc.py
b/airflow/providers/jdbc/operators/jdbc.py
index 7fefc74a93..2c023d9afe 100644
--- a/airflow/providers/jdbc/operators/jdbc.py
+++ b/airflow/providers/jdbc/operators/jdbc.py
@@ -15,6 +15,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
from typing import TYPE_CHECKING, Iterable, List, Mapping, Optional, Sequence,
Union
from airflow.models import BaseOperator
@@ -24,6 +25,11 @@ if TYPE_CHECKING:
from airflow.utils.context import Context
+def fetch_all_handler(cursor):
+ """Handler for DbApiHook.run() to return results"""
+ return cursor.fetchall()
+
+
class JdbcOperator(BaseOperator):
"""
Executes sql code in a database using jdbc driver.
@@ -67,4 +73,4 @@ class JdbcOperator(BaseOperator):
def execute(self, context: 'Context') -> None:
self.log.info('Executing: %s', self.sql)
hook = JdbcHook(jdbc_conn_id=self.jdbc_conn_id)
- hook.run(self.sql, self.autocommit, parameters=self.parameters)
+ return hook.run(self.sql, self.autocommit, parameters=self.parameters,
handler=fetch_all_handler)
diff --git a/tests/providers/jdbc/operators/test_jdbc.py
b/tests/providers/jdbc/operators/test_jdbc.py
index e16ac446e6..812d60fd45 100644
--- a/tests/providers/jdbc/operators/test_jdbc.py
+++ b/tests/providers/jdbc/operators/test_jdbc.py
@@ -19,7 +19,7 @@
import unittest
from unittest.mock import patch
-from airflow.providers.jdbc.operators.jdbc import JdbcOperator
+from airflow.providers.jdbc.operators.jdbc import JdbcOperator,
fetch_all_handler
class TestJdbcOperator(unittest.TestCase):
@@ -33,5 +33,8 @@ class TestJdbcOperator(unittest.TestCase):
mock_jdbc_hook.assert_called_once_with(jdbc_conn_id=jdbc_operator.jdbc_conn_id)
mock_jdbc_hook.return_value.run.assert_called_once_with(
- jdbc_operator.sql, jdbc_operator.autocommit,
parameters=jdbc_operator.parameters
+ jdbc_operator.sql,
+ jdbc_operator.autocommit,
+ parameters=jdbc_operator.parameters,
+ handler=fetch_all_handler,
)
