Skip to content

Commit

Permalink
Handler parameter from JdbcOperator to JdbcHook.run
Browse files Browse the repository at this point in the history
  • Loading branch information
kazanzhy committed May 20, 2022
1 parent 4b731f4 commit 4adb576
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 3 deletions.
8 changes: 7 additions & 1 deletion airflow/providers/jdbc/operators/jdbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -24,6 +25,11 @@
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.
Expand Down Expand Up @@ -67,4 +73,4 @@ def __init__(
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)
7 changes: 5 additions & 2 deletions tests/providers/jdbc/operators/test_jdbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -33,5 +33,8 @@ def test_execute(self, mock_jdbc_hook):

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,
)

0 comments on commit 4adb576

Please sign in to comment.