build params

Author: dgalas2

singlestoredb/config.py

@@ -9,6 +9,7 @@
 from .utils.config import check_float  # noqa: F401
 from .utils.config import check_int  # noqa: F401
 from .utils.config import check_optional_bool  # noqa: F401
+from .utils.config import check_socket_options  # noqa: F401
 from .utils.config import check_str  # noqa: F401
 from .utils.config import check_url  # noqa: F401
 from .utils.config import describe_option  # noqa: F401
@@ -263,6 +264,11 @@
     environ='SINGLESTOREDB_FUSION_ENABLED',
 )
 
+register_option(
+    'socket_options', 'dict', check_socket_options, None,
+    'Format for socket options',
+)
+
 #
 # Query results options
 #

singlestoredb/connection.py

@@ -1320,7 +1320,6 @@ 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]]] = None,
     ssl_verify_identity: Optional[bool] = None,
     conv: Optional[Dict[int, Callable[..., Any]]] = None,
     credential_type: Optional[str] = None,
@@ -1341,6 +1340,7 @@ def connect(
     vector_data_format: Optional[str] = None,
     parse_json: Optional[bool] = None,
     interpolate_query_with_empty_args: Optional[bool] = None,
+    socket_options: Optional[Dict[int, Dict[int, Any]]] = None,
 ) -> Connection:
     """
     Return a SingleStoreDB connection.
@@ -1429,6 +1429,11 @@ def connect(
     interpolate_query_with_empty_args : bool, optional
         Should the connector apply parameter interpolation even when the
         parameters are empty? This corresponds to pymysql/mysqlclient's handling
+    socket_options : dict, optional
+        Socket options to set on the underlying socket. The keys should be
+        socket level constants (e.g., socket.SOL_SOCKET) and the values should be
+        dictionaries mapping socket option constants (e.g., socket.SO_KEEPALIVE) to
+        the desired value for that option.
 
     Examples
     --------

singlestoredb/utils/config.py

@@ -646,6 +646,55 @@ def check_str(
     return out
 
 
+def check_socket_options(
+    value: Any,
+) -> Optional[Dict[int, Dict[int, Any]]]:
+    """
+    Validate socket options.
+
+    Parameters
+    ----------
+    value : dict
+        The value to validate.  It must be a dictionary where the keys are
+        socket level constants (e.g., socket.SOL_SOCKET) and the values are
+        dictionaries mapping socket option constants (e.g., socket.SO_KEEPALIVE)
+        to the desired value for that option.
+
+    Returns
+    -------
+    dict
+        The validated socket options
+
+    """
+    if value is None:
+        return None
+
+    if not isinstance(value, Mapping):
+        raise ValueError(
+            'value {} must be of type dict'.format(value),
+        )
+
+    out: dict[int, dict[int, Any]] = {}
+    for level, options in value.items():
+        if not isinstance(level, int):
+            raise ValueError(
+                f'keys in {value} must be integers corresponding to socket levels',
+            )
+        if not isinstance(options, Mapping):
+            raise ValueError(
+                f'values in {value} must be dicts.',
+            )
+        out[level] = {}
+        for opt, val in options.items():
+            if not isinstance(opt, int):
+                raise ValueError(
+                    f'keys in sub-dicts of {value} must be integers.',
+                )
+            out[level][opt] = val
+
+    return out
+
+
 def check_dict_str_str(
     value: Any,
 ) -> Optional[Dict[str, str]]: