Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Suppress jaydebeapi.Error when setAutoCommit or getAutoCommit is unsupported by JDBC driver #38707

Merged
merged 37 commits into from
Apr 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
37 commits
Select commit Hold shift + click to select a range
ff557a2
refactor: Suppress jaydebeapi.Error when setAutoCommit or getAutoComm…
davidblain-infrabel Apr 3, 2024
3e99705
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 3, 2024
1822f90
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 4, 2024
79b3c75
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 6, 2024
3a8bf82
refactor: Added suppress_and_warn context manager which allows specif…
davidblain-infrabel Apr 9, 2024
321ad02
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 9, 2024
ebbeffa
refactor: Print the stacktrace beside the exception when an exception…
davidblain-infrabel Apr 9, 2024
81fc3f6
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 9, 2024
1b37930
refactor: Re-organized imports of context
davidblain-infrabel Apr 9, 2024
0d8f4e2
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 9, 2024
da6a7d7
refactor: Added return type to suppress_and_warn contextmanager
davidblain-infrabel Apr 9, 2024
545c5d0
refactor: Use type instead of Type for Python 3.8 compat
davidblain-infrabel Apr 9, 2024
76e8e1f
refactor: Reorganized typing imports on context
davidblain-infrabel Apr 9, 2024
77c65c4
refactor: Added license on top of TestContext file
davidblain-infrabel Apr 9, 2024
b063da8
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 9, 2024
e02462a
refactor: Removed return type of suppress_and_warn
davidblain-infrabel Apr 9, 2024
72ac9d1
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 10, 2024
2b08af9
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 10, 2024
19e0df2
refactor: Moved suppress_and_warn method from Airflow context to comm…
davidblain-infrabel Apr 11, 2024
d0c61d0
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
569ef2c
refactor: Changed version of apache-airflow-providers-common-sql depe…
davidblain-infrabel Apr 11, 2024
2d2db8c
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
54b49a2
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
8eb0ee6
Revert "refactor: Changed version of apache-airflow-providers-common-…
davidblain-infrabel Apr 11, 2024
dcc02b6
refactor: Moved suppress_and_warn method from common sql provider to …
davidblain-infrabel Apr 11, 2024
c3807a7
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
63982cb
refactor: Also updated the apache-airflow dependency version from 2.6…
davidblain-infrabel Apr 11, 2024
36ec020
Revert "refactor: Also updated the apache-airflow dependency version …
davidblain-infrabel Apr 11, 2024
39b30da
fix: Move tests of suppress_and_warn from TestDbApiHook to TestJdbcHook
davidblain-infrabel Apr 11, 2024
522af2e
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
c212fc8
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
f8bfe3f
refactor: Updated common sql.pyi like main
davidblain-infrabel Apr 11, 2024
154052d
refactor: Return False in get_autocommit if exception is suppressed
davidblain-infrabel Apr 11, 2024
c531271
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
fdc3e2a
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
93618ef
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
77df3aa
Merge branch 'main' into feature/jdbc-autocommit
dabla Apr 11, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions airflow/providers/jdbc/hooks/jdbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
# under the License.
from __future__ import annotations

import traceback
import warnings
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any

import jaydebeapi
Expand All @@ -27,6 +30,15 @@
from airflow.models.connection import Connection


@contextmanager
def suppress_and_warn(*exceptions: type[BaseException]):
"""Context manager that suppresses the given exceptions and logs a warning message."""
try:
yield
except exceptions as e:
warnings.warn(f"Exception suppressed: {e}\n{traceback.format_exc()}", category=UserWarning)


class JdbcHook(DbApiHook):
"""General hook for JDBC access.

Expand Down Expand Up @@ -152,7 +164,8 @@ def set_autocommit(self, conn: jaydebeapi.Connection, autocommit: bool) -> None:
:param conn: The connection.
:param autocommit: The connection's autocommit setting.
"""
conn.jconn.setAutoCommit(autocommit)
with suppress_and_warn(jaydebeapi.Error):
conn.jconn.setAutoCommit(autocommit)
dabla marked this conversation as resolved.
Show resolved Hide resolved

def get_autocommit(self, conn: jaydebeapi.Connection) -> bool:
"""Get autocommit setting for the provided connection.
Expand All @@ -162,4 +175,6 @@ def get_autocommit(self, conn: jaydebeapi.Connection) -> bool:
to True on the connection. False if it is either not set, set to
False, or the connection does not support auto-commit.
"""
return conn.jconn.getAutoCommit()
with suppress_and_warn(jaydebeapi.Error):
return conn.jconn.getAutoCommit()
return False
28 changes: 27 additions & 1 deletion tests/providers/jdbc/hooks/test_jdbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,14 @@
from unittest import mock
from unittest.mock import Mock, patch

import jaydebeapi
import pytest

from airflow.exceptions import DeserializingResultError
from airflow.models import Connection
from airflow.providers.jdbc.hooks.jdbc import JdbcHook
from airflow.providers.jdbc.hooks.jdbc import JdbcHook, suppress_and_warn
from airflow.utils import db
from airflow.utils.context import AirflowContextDeprecationWarning

pytestmark = pytest.mark.db_test

Expand Down Expand Up @@ -82,13 +85,27 @@ def test_jdbc_conn_set_autocommit(self, _):
jdbc_hook.set_autocommit(jdbc_conn, False)
jdbc_conn.jconn.setAutoCommit.assert_called_once_with(False)

@patch("airflow.providers.jdbc.hooks.jdbc.jaydebeapi.connect")
def test_jdbc_conn_set_autocommit_when_not_supported(self, _):
jdbc_hook = JdbcHook()
jdbc_conn = jdbc_hook.get_conn()
jdbc_conn.jconn.setAutoCommit.side_effect = jaydebeapi.Error()
jdbc_hook.set_autocommit(jdbc_conn, False)

@patch("airflow.providers.jdbc.hooks.jdbc.jaydebeapi.connect")
def test_jdbc_conn_get_autocommit(self, _):
jdbc_hook = JdbcHook()
jdbc_conn = jdbc_hook.get_conn()
jdbc_hook.get_autocommit(jdbc_conn)
jdbc_conn.jconn.getAutoCommit.assert_called_once_with()

@patch("airflow.providers.jdbc.hooks.jdbc.jaydebeapi.connect")
def test_jdbc_conn_get_autocommit_when_not_supported_then_return_false(self, _):
jdbc_hook = JdbcHook()
jdbc_conn = jdbc_hook.get_conn()
jdbc_conn.jconn.getAutoCommit.side_effect = jaydebeapi.Error()
assert jdbc_hook.get_autocommit(jdbc_conn) is False

def test_driver_hook_params(self):
hook = get_hook(hook_params=dict(driver_path="Blah driver path", driver_class="Blah driver class"))
assert hook.driver_path == "Blah driver path"
Expand Down Expand Up @@ -161,3 +178,12 @@ def test_driver_extra_raises_warning_and_returns_default_driver_by_default(self,
"have supplied 'driver_class' via connection extra but it will not be used"
) in caplog.text
assert driver_class == "Blah driver class"

def test_suppress_and_warn_when_raised_exception_is_suppressed(self):
with suppress_and_warn(AirflowContextDeprecationWarning):
raise AirflowContextDeprecationWarning()

def test_suppress_and_warn_when_raised_exception_is_not_suppressed(self):
with pytest.raises(AirflowContextDeprecationWarning):
with suppress_and_warn(DeserializingResultError):
raise AirflowContextDeprecationWarning()