Skip to content

Commit

Permalink
fix(connection): Raise OperationalError for socket timeouts (aws#179)
Browse files Browse the repository at this point in the history
* Raise OperationalError for socket timeouts

* add unit test and integrationt test

* rectify unit test

* rectify integration test for socket_timeout

* remove empty config.ini
  • Loading branch information
jiezhen-chen authored and soksamnanglim committed Aug 16, 2023
1 parent c7a9ff3 commit 98c4496
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 0 deletions.
5 changes: 5 additions & 0 deletions redshift_connector/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -656,6 +656,11 @@ def get_calling_module() -> str:
self._sock: typing.Optional[typing.BinaryIO] = self._usock.makefile(mode="rwb")
if tcp_keepalive:
self._usock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

except socket.timeout as timeout_error:
self._usock.close()
raise OperationalError("connection time out", timeout_error)

except socket.error as e:
self._usock.close()
raise InterfaceError("communication error", e)
Expand Down
6 changes: 6 additions & 0 deletions test/integration/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -307,3 +307,9 @@ def test_execute_do_parsing_bind_params_when_exist(mocker, db_kwargs, sql, args)
with redshift_connector.connect(**db_kwargs) as conn:
conn.cursor().execute(sql, args)
assert convert_paramstyle_spy.called

def test_socket_timeout(db_kwargs):
db_kwargs["timeout"] = 0

with pytest.raises(redshift_connector.InterfaceError):
redshift_connector.connect(**db_kwargs)
9 changes: 9 additions & 0 deletions test/unit/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
from collections import deque
from decimal import Decimal
from unittest.mock import patch
import socket
from unittest import mock

import pytest # type: ignore

Expand All @@ -12,6 +14,7 @@
IntegrityError,
InterfaceError,
ProgrammingError,
OperationalError
)
from redshift_connector.config import (
ClientProtocolVersion,
Expand Down Expand Up @@ -328,3 +331,9 @@ def test_client_os_version_is_not_present():

with patch("platform.platform", side_effect=Exception("not for you")):
assert mock_connection.client_os_version == "unknown"

def test_socket_timeout_error():
with mock.patch('socket.socket.connect') as mock_socket:
mock_socket.side_effect = (socket.timeout)
with pytest.raises(OperationalError):
Connection(user='mock_user', password='mock_password', host='localhost', port=8080, database='mocked')

0 comments on commit 98c4496

Please sign in to comment.