Add option to pass tls sni servername (#49)

Author: hbindu-ss

singlestoredb/config.py

@@ -134,6 +134,12 @@
     environ='SINGLESTOREDB_SSL_CIPHER',
 )
 
+register_option(
+    'tls_sni_servername', 'str', check_str, None,
+    'Sets TLS SNI servername',
+    environ='SINGLESTOREDB_TLS_SNI_SERVERNAME',
+)
+
 register_option(
     'ssl_disabled', 'bool', check_bool, False,
     'Disable SSL usage',

singlestoredb/connection.py

@@ -1298,6 +1298,7 @@ def connect(
     ssl_key: Optional[str] = None, ssl_cert: Optional[str] = None,
     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,
     ssl_verify_identity: Optional[bool] = None,
     conv: Optional[Dict[int, Callable[..., Any]]] = None,
     credential_type: Optional[str] = None,

singlestoredb/mysql/connection.py

@@ -226,6 +226,8 @@ class Connection(BaseConnection):
         Set to true to check the server certificate's validity.
     ssl_verify_identity : bool, optional
         Set to true to check the server's identity.
+    tls_sni_servername: str, optional
+        Set server host name for TLS connection
     read_default_group : str, optional
         Group to read from in the configuration file.
     autocommit : bool, optional
@@ -295,6 +297,7 @@ class Connection(BaseConnection):
     _auth_plugin_name = ''
     _closed = False
     _secure = False
+    _tls_sni_servername = None
 
     def __init__(  # noqa: C901
         self,
@@ -335,6 +338,7 @@ def __init__(  # noqa: C901
         ssl_key=None,
         ssl_verify_cert=None,
         ssl_verify_identity=None,
+        tls_sni_servername=None,
         parse_json=True,
         invalid_values=None,
         pure_python=None,
@@ -638,6 +642,7 @@ def _config(key, arg):
 
         self._is_committable = True
         self._in_sync = False
+        self._tls_sni_servername = tls_sni_servername
         self._track_env = bool(track_env) or self.host == 'singlestore.com'
         self._enable_extended_data_types = enable_extended_data_types
         if vector_data_format.lower() in ['json', 'binary']:
@@ -1364,7 +1369,10 @@ def _request_authentication(self):  # noqa: C901
         if self.ssl and self.server_capabilities & CLIENT.SSL:
             self.write_packet(data_init)
 
-            self._sock = self.ctx.wrap_socket(self._sock, server_hostname=self.host)
+            hostname = self.host
+            if self._tls_sni_servername:
+                hostname = self._tls_sni_servername
+            self._sock = self.ctx.wrap_socket(self._sock, server_hostname=hostname)
             self._rfile = self._sock.makefile('rb')
             self._secure = True