socket_options

Author: dgalas2

singlestoredb/connection.py

@@ -1320,6 +1320,7 @@ def connect(
     ssl_ca: Optional[str] = None, ssl_disabled: Optional[bool] = None,
     ssl_cipher: Optional[str] = None, ssl_verify_cert: Optional[bool] = None,
     tls_sni_servername: Optional[str] = None,
+    socket_options: Optional[Dict[int, Dict[int, Any]]] = {},
     ssl_verify_identity: Optional[bool] = None,
     conv: Optional[Dict[int, Callable[..., Any]]] = None,
     credential_type: Optional[str] = None,

singlestoredb/mysql/connection.py

@@ -230,6 +230,10 @@ class Connection(BaseConnection):
         Set to true to check the server's identity.
     tls_sni_servername: str, optional
         Set server host name for TLS connection
+    socket_options: Dict[int, Dict[int, any]], optional
+        A dictionary of socket options to set on the connection.
+        The keys are the socket level constants (e.g., socket.SOL_SOCKET),
+        and the values are dictionaries mapping option names to values.
     read_default_group : str, optional
         Group to read from in the configuration file.
     autocommit : bool, optional
@@ -341,6 +345,7 @@ def __init__(  # noqa: C901
         ssl_verify_cert=None,
         ssl_verify_identity=None,
         tls_sni_servername=None,
+        socket_options={},
         parse_json=True,
         invalid_values=None,
         pure_python=None,
@@ -477,7 +482,7 @@ def _config(key, arg):
         self.collation = collation
         self.use_unicode = use_unicode
         self.encoding_errors = encoding_errors
-
+        self.socket_options = socket_options
         self.encoding = charset_by_name(self.charset).encoding
 
         client_flag |= CLIENT.CAPABILITIES
@@ -1110,9 +1115,10 @@ def connect(self, sock=None):
 
                     # setting TCP keepalive for mysql
                     # 60s idle, 30s interval, 5 times before close
-                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60)
-                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 30)
-                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5)
+                    for level, options in self.socket_options.items():
+                        for opt, value in options.items():
+                            sock.setsockopt(level, opt, value)
+
                 sock.settimeout(None)
 
             self._sock = sock