From 0514002cd1327136e80383ff7326c48cca0a95b2 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 5 May 2025 10:15:30 -0700 Subject: [PATCH] add support for tcp keepalive options --- redshift_connector/core.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/redshift_connector/core.py b/redshift_connector/core.py index 3c12530..393c0f4 100644 --- a/redshift_connector/core.py +++ b/redshift_connector/core.py @@ -1,6 +1,7 @@ import logging import os import socket +import sys import typing from collections import deque from copy import deepcopy @@ -363,6 +364,18 @@ def create_message(code: bytes, data: bytes = b"") -> bytes: arr_trans: typing.Mapping[int, typing.Optional[str]] = dict(zip(map(ord, "[] 'u"), ["{", "}", None, None, None])) +class TCPKeepAlive: + after_idle_sec: int + """ + TCP Keepalive settings + """ + + def __init__(self, after_idle_sec: int, interval_sec: int, max_fails: int): + self.after_idle_sec = after_idle_sec + self.interval_sec = interval_sec + self.max_fails = max_fails + + class Connection: # DBAPI Extension: supply exceptions as attributes on the connection Warning = property(lambda self: self._getError(Warning)) @@ -422,7 +435,7 @@ def __init__( sslmode: str = "verify-ca", timeout: typing.Optional[int] = None, max_prepared_statements: int = 1000, - tcp_keepalive: typing.Optional[bool] = True, + tcp_keepalive: typing.Optional[typing.Union[bool, TCPKeepAlive]] = True, application_name: typing.Optional[str] = None, replication: typing.Optional[str] = None, client_protocol_version: int = DEFAULT_PROTOCOL_VERSION, @@ -682,8 +695,26 @@ def get_calling_module() -> str: self._sock = self._usock.makefile(mode="rwb") if tcp_keepalive: + # logic copied from: https://github.com/sfinktah/keepalive/blob/master/keepalive/keepalive.py _logger.debug("enabling tcp keepalive on socket") self._usock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + # set a sane default value for tcp_keepalive if not supplied + if not isinstance(tcp_keepalive, TCPKeepAlive): + tcp_keepalive = TCPKeepAlive(after_idle_sec=200, interval_sec=200, max_fails=5) + if sys.platform == 'linux': + self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, + tcp_keepalive.after_idle_sec) + self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, + tcp_keepalive.interval_sec) + self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, + tcp_keepalive.max_fails) + elif sys.platform == 'darwin': + # scraped from /usr/include, not exported by python's socket module + _TCP_KEEPALIVE = 0x10 + self._usock.setsockopt(socket.IPPROTO_TCP, _TCP_KEEPALIVE, tcp_keepalive.interval_sec) + elif sys.platform == 'win32': + self._usock.ioctl(socket.SIO_KEEPALIVE_VALS, (1, tcp_keepalive.after_idle_sec * 1000, + tcp_keepalive.interval_sec * 1000)) except socket.timeout as timeout_error: self._usock.close()