feat(connection): add TCP keepalive configuration options

Author: Luis Garza

redshift_connector/__init__.py

@@ -103,6 +103,18 @@
 
 __author__ = "Mathieu Fenniak"
 
+def validate_keepalive_values(idle, interval, count):
+    if idle is not None:
+        if idle <= 0:
+            raise ValueError("tcp_keepalive_idle must be positive")
+
+    if interval is not None:
+        if interval <= 0:
+            raise ValueError("tcp_keepalive_interval must be positive")
+
+    if count is not None:
+        if count <= 0:
+            raise ValueError("tcp_keepalive_count must be positive")
 
 def connect(
     user: typing.Optional[str] = None,
@@ -117,6 +129,9 @@ def connect(
     timeout: typing.Optional[int] = None,
     max_prepared_statements: typing.Optional[int] = None,
     tcp_keepalive: typing.Optional[bool] = None,
+    tcp_keepalive_idle: typing.Optional[int] = None,
+    tcp_keepalive_interval: typing.Optional[int] = None,
+    tcp_keepalive_count: typing.Optional[int] = None,
     application_name: typing.Optional[str] = None,
     replication: typing.Optional[str] = None,
     idp_host: typing.Optional[str] = None,
@@ -193,7 +208,13 @@ def connect(
         The number of seconds before the connection to the server will timeout. By default there is no timeout.
     max_prepared_statements : Optional[int]
     tcp_keepalive : Optional[bool]
-        Is `TCP keepalive <https://en.wikipedia.org/wiki/Keepalive#TCP_keepalive>`_ used. The default value is ``True``.
+        A boolean specifying whether the driver uses TCP keepalives to prevent connections from timing out. Defaults to ``True``.
+    tcp_keepalive_idle : Optional[int]
+        Time (in seconds) before sending keepalive probes. Defaults to None (system default).
+    tcp_keepalive_interval : Optional[int]
+        Time (in seconds) between keepalive probes. Defaults to None (system default).
+    tcp_keepalive_count : Optional[int]
+        Number of failed probes before connection is considered dead. Defaults to None (system default).
     application_name : Optional[str]
         Sets the application name. The default value is None.
     replication : Optional[str]
@@ -343,6 +364,9 @@ def connect(
     info.put("ssl_insecure", ssl_insecure)
     info.put("sslmode", sslmode)
     info.put("tcp_keepalive", tcp_keepalive)
+    info.put("tcp_keepalive_idle", tcp_keepalive_idle)
+    info.put("tcp_keepalive_interval", tcp_keepalive_interval)
+    info.put("tcp_keepalive_count", tcp_keepalive_count)
     info.put("timeout", timeout)
     info.put("token", token)
     info.put("token_type", token_type)
@@ -385,6 +409,16 @@ def connect(
     if not redshift_native_auth:
         IamHelper.set_iam_properties(info)
 
+    if info.tcp_keepalive:
+        try:
+            validate_keepalive_values(
+                info.tcp_keepalive_idle,
+                info.tcp_keepalive_interval,
+                info.tcp_keepalive_count
+            )
+        except ValueError as e:
+            raise InterfaceError(str(e))
+
     _logger.debug(make_divider_block())
     _logger.debug("Connection arguments following validation and IAM auth (if applicable)")
     _logger.debug(make_divider_block())
@@ -404,6 +438,9 @@ def connect(
         timeout=info.timeout,
         max_prepared_statements=info.max_prepared_statements,
         tcp_keepalive=info.tcp_keepalive,
+        tcp_keepalive_idle=info.tcp_keepalive_idle,
+        tcp_keepalive_interval=info.tcp_keepalive_interval,
+        tcp_keepalive_count=info.tcp_keepalive_count,
         application_name=info.application_name,
         replication=info.replication,
         client_protocol_version=info.client_protocol_version,

redshift_connector/core.py

@@ -363,7 +363,6 @@ 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 Connection:
     # DBAPI Extension: supply exceptions as attributes on the connection
     Warning = property(lambda self: self._getError(Warning))
@@ -424,6 +423,9 @@ def __init__(
         timeout: typing.Optional[int] = None,
         max_prepared_statements: int = DEFAULT_MAX_PREPARED_STATEMENTS,
         tcp_keepalive: typing.Optional[bool] = True,
+        tcp_keepalive_idle: typing.Optional[int] = None,
+        tcp_keepalive_interval: typing.Optional[int] = None,
+        tcp_keepalive_count: typing.Optional[int] = None,
         application_name: typing.Optional[str] = None,
         replication: typing.Optional[str] = None,
         client_protocol_version: int = DEFAULT_PROTOCOL_VERSION,
@@ -462,7 +464,13 @@ def __init__(
             The number of seconds before the connection to the server will timeout. By default there is no timeout.
         max_prepared_statements : int
         tcp_keepalive : Optional[bool]
-            Is `TCP keepalive <https://en.wikipedia.org/wiki/Keepalive#TCP_keepalive>`_ used. The default value is ``True``.
+            A boolean specifying whether the driver uses TCP keepalives to prevent connections from timing out. Defaults to ``True``.
+        tcp_keepalive_idle : Optional[int]
+            Time (in seconds) before sending TCP keepalive probes on an idle connection.  Defaults to None (system default).
+        tcp_keepalive_interval : Optional[int]
+            Time (in seconds) between individual TCP keepalive probes.  Defaults to None (system default).
+        tcp_keepalive_count : Optional[int]
+            Maximum number of TCP keepalive probes to send before dropping the connection if no response is received.  Defaults to None (system default).
         application_name : Optional[str]
             Sets the application name. The default value is None.
         replication : Optional[str]
@@ -686,6 +694,32 @@ def get_calling_module() -> str:
                 _logger.debug("enabling tcp keepalive on socket")
                 self._usock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
 
+                # Set TCP keepalive parameters if supported by platform and values are defined
+                if tcp_keepalive_idle is not None:
+                    # Mac OS X uses TCP_KEEPALIVE instead of TCP_KEEPIDLE
+                    if hasattr(socket, 'TCP_KEEPIDLE'):
+                        self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, tcp_keepalive_idle)
+                        _logger.debug(f"Set TCP_KEEPIDLE to {tcp_keepalive_idle}")
+                    elif hasattr(socket, 'TCP_KEEPALIVE'):  # macOS/BSD
+                        self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, tcp_keepalive_idle)
+                        _logger.debug(f"Set TCP_KEEPALIVE to {tcp_keepalive_idle}")
+                    else:
+                        _logger.warning("Neither TCP_KEEPIDLE nor TCP_KEEPALIVE supported on this platform")
+
+                if tcp_keepalive_interval is not None:
+                    if hasattr(socket, 'TCP_KEEPINTVL'):
+                        self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, tcp_keepalive_interval)
+                        _logger.debug(f"Set TCP_KEEPINTVL to {tcp_keepalive_interval}")
+                    else:
+                        _logger.warning("TCP_KEEPINTVL not supported on this platform")
+
+                if tcp_keepalive_count is not None:
+                    if hasattr(socket, 'TCP_KEEPCNT'):
+                        self._usock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, tcp_keepalive_count)
+                        _logger.debug(f"Set TCP_KEEPCNT to {tcp_keepalive_count}")
+                    else:
+                        _logger.warning("TCP_KEEPCNT not supported on this platform")
+
         except socket.timeout as timeout_error:
             self._usock.close()
             raise OperationalError("connection time out", timeout_error)

redshift_connector/redshift_property.py

@@ -112,6 +112,12 @@ def __init__(self: "RedshiftProperty", **kwargs):
             self.sslmode: str = "verify-ca"
             # Use this property to enable or disable TCP keepalives.
             self.tcp_keepalive: bool = True
+            # Time (in seconds) before sending keepalive probes
+            self.tcp_keepalive_idle: typing.Optional[int] = None
+            # Time (in seconds) between keepalive probes
+            self.tcp_keepalive_interval: typing.Optional[int] = None
+            # Number of failed probes before connection is considered dead
+            self.tcp_keepalive_count: typing.Optional[int] = None
             # This is the time in seconds before the connection to the server will time out.
             self.timeout: typing.Optional[int] = None
             self.token: typing.Optional[str] = None

redshift_connector/utils/logging_utils.py

@@ -71,6 +71,9 @@ def mask_secure_info_in_props(info: "RedshiftProperty") -> "RedshiftProperty":
         "ssl_insecure",
         "sslmode",
         "tcp_keepalive",
+        "tcp_keepalive_idle",
+        "tcp_keepalive_interval",
+        "tcp_keepalive_count",
         "token_type",
         "timeout",
         "unix_sock",

test/integration/test_connection.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import random
+import socket
 import string
 import sys
 import typing
@@ -334,3 +335,72 @@ def test_socket_timeout(db_kwargs) -> None:
 
     with pytest.raises(redshift_connector.InterfaceError):
         redshift_connector.connect(**db_kwargs)
+
+def test_tcp_keepalive(db_kwargs) -> None:
+    """Test TCP keepalive configuration is properly applied"""
+    db_kwargs["tcp_keepalive"] = True
+    db_kwargs["tcp_keepalive_idle"] = 30
+    db_kwargs["tcp_keepalive_interval"] = 10
+    db_kwargs["tcp_keepalive_count"] = 3
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        # Verify the socket options after connection is established
+        sock = conn._usock
+
+        # Get the socket options
+        keepalive = sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE)
+        assert keepalive > 0, "TCP keepalive should be enabled"
+
+        if hasattr(socket, 'TCP_KEEPIDLE'):
+            idle = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)
+            assert idle == 30, "TCP keepalive idle time should be 30"
+        elif hasattr(socket, 'TCP_KEEPALIVE'):
+            idle = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)
+            assert idle == 30, "TCP keepalive idle time should be 30"
+
+        if hasattr(socket, 'TCP_KEEPINTVL'):
+            interval = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)
+            assert interval == 10, "TCP keepalive interval should be 10"
+
+        if hasattr(socket, 'TCP_KEEPCNT'):
+            count = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)
+            assert count == 3, "TCP keepalive count should be 3"
+
+def test_tcp_keepalive_defaults(db_kwargs) -> None:
+    """Test TCP keepalive with default values"""
+    with redshift_connector.connect(**db_kwargs) as conn:
+        # Verify the socket options after connection is established
+        sock = conn._usock
+
+        # Verify SO_KEEPALIVE is enabled
+        keepalive = sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE)
+        assert keepalive > 0, "TCP keepalive should be enabled"
+
+        # Verify other parameters were not explicitly set by checking they match system defaults
+        # We can do this by creating another connection without keepalive and comparing values
+        with redshift_connector.connect(**db_kwargs) as conn2:
+            db_kwargs["tcp_keepalive"] = False
+            sock2 = conn2._usock
+
+            if hasattr(socket, 'TCP_KEEPIDLE'):
+                val1 = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)
+                val2 = sock2.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)
+                assert val1 == val2, "TCP_KEEPIDLE should match system default. Got {} but expected {}".format(val1,
+                                                                                                               val2)
+            elif hasattr(socket, 'TCP_KEEPALIVE'):
+                val1 = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)
+                val2 = sock2.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)
+                assert val1 == val2, "TCP_KEEPALIVE should match system default. Got {} but expected {}".format(val1,
+                                                                                                                val2)
+
+            if hasattr(socket, 'TCP_KEEPINTVL'):
+                val1 = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)
+                val2 = sock2.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)
+                assert val1 == val2, "TCP_KEEPINTVL should match system default. Got {} but expected {}".format(val1,
+                                                                                                                val2)
+
+            if hasattr(socket, 'TCP_KEEPCNT'):
+                val1 = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)
+                val2 = sock2.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)
+                assert val1 == val2, "TCP_KEEPCNT should match system default. Got {} but expected {}".format(val1,
+                                                                                                              val2)

test/unit/test_connection.py

@@ -16,6 +16,7 @@
     InterfaceError,
     OperationalError,
     ProgrammingError,
+    validate_keepalive_values,
 )
 from redshift_connector.config import (
     ClientProtocolVersion,
@@ -560,3 +561,24 @@ def test_make_params_maps_pandas_timestamp_to_timestamp(db_kwargs):
     assert res[0][0] == RedshiftOID.TIMESTAMPTZ
     assert res[0][1] == 1
     assert res[0][2] == timestamptz_send_integer
+
+
+@pytest.mark.parametrize(
+    "idle,interval,count,expected_error",
+    [
+        (-1, None, None, "tcp_keepalive_idle must be positive"),
+        (None, -1, None, "tcp_keepalive_interval must be positive"),
+        (None, None, -1, "tcp_keepalive_count must be positive"),
+        (0, None, None, "tcp_keepalive_idle must be positive"),
+        (None, 0, None, "tcp_keepalive_interval must be positive"),
+        (None, None, 0, "tcp_keepalive_count must be positive"),
+        (1, 1, 1, None),  # valid values should not raise error
+    ],
+)
+def test_tcp_keepalive_validation(idle, interval, count, expected_error):
+    if expected_error:
+        with pytest.raises(ValueError, match=expected_error):
+            validate_keepalive_values(idle, interval, count)
+    else:
+        # Should not raise any exception
+        validate_keepalive_values(idle, interval, count)