fix(mssql-python): Allow login_attempts kwarg and handle odbc_property overrides.

Signed-off-by: walanguzzi <walanguzzi@outlook.com>

Author: walanguzzi

sqlmesh/core/config/connection.py

@@ -1519,6 +1519,7 @@ class MSSQLConnectionConfig(ConnectionConfig):
     database: t.Optional[str] = ""
     timeout: t.Optional[int] = 0
     login_timeout: t.Optional[int] = 60
+    login_attempts: t.Optional[int] = 1
     charset: t.Optional[str] = "UTF-8"
     appname: t.Optional[str] = None
     port: t.Optional[int] = 1433
@@ -1610,6 +1611,7 @@ def _connection_kwargs_keys(self) -> t.Set[str]:
                     "trust_server_certificate",
                     "encrypt",
                     "odbc_properties",
+                    "login_attempts",
                 }
             )
             # Remove pymssql-specific parameters
@@ -1648,8 +1650,9 @@ def connect_mssql_python(**kwargs: t.Any) -> t.Callable:
                 authentication = kwargs.pop("authentication", None)
                 trust_server_certificate = kwargs.pop("trust_server_certificate", False)
                 encrypt = kwargs.pop("encrypt", True)
+                timeout = kwargs.pop("timeout", 0)
                 login_timeout = kwargs.pop("login_timeout", 59)
-                login_attempts = kwargs.pop("login_attempts", 1)  # TODO: document
+                login_attempts = kwargs.pop("login_attempts", 1)
 
                 # Build connection string
                 conn_str_parts = [
@@ -1664,6 +1667,11 @@ def connect_mssql_python(**kwargs: t.Any) -> t.Callable:
                 if trust_server_certificate:
                     conn_str_parts.append("TrustServerCertificate=yes")
 
+                # `Connection Timeout=` is not a valid option so we leverage `ConnectRetry*`.
+                # See the following:
+                # - https://github.com/microsoft/mssql-python/issues/339
+                # - https://github.com/microsoft/mssql-python/wiki/Connection-to-SQL-Database
+                # - https://github.com/microsoft/mssql-python/wiki/Connection#timeout
                 conn_str_parts.append(f"ConnectRetryCount={login_attempts}")
                 conn_str_parts.append(f"ConnectRetryInterval={min(int(login_timeout), 60)}")
 
@@ -1687,6 +1695,8 @@ def connect_mssql_python(**kwargs: t.Any) -> t.Callable:
                             "pwd",
                             "encrypt",
                             "trustservercertificate",
+                            "connectretrycount",
+                            "connectretryinterval",
                             "connection timeout",
                         ):
                             continue
@@ -1700,7 +1710,11 @@ def connect_mssql_python(**kwargs: t.Any) -> t.Callable:
                 # Create the connection
                 conn_str = ";".join(conn_str_parts)
 
-                conn = mssql_python.connect(conn_str, autocommit=kwargs.get("autocommit", False))
+                conn = mssql_python.connect(
+                    conn_str,
+                    autocommit=kwargs.get("autocommit", False),
+                    timeout=timeout,
+                )
 
                 # TODO: Remove this output converter as DATETIMEOFFSET
                 # should be handled natively by `mssql-python`.

tests/core/test_connection_config.py

@@ -1671,6 +1671,7 @@ def test_mssql_connection_kwargs_keys():
             "host",
             "timeout",
             "login_timeout",
+            "login_attempts",
             "charset",
             "appname",
             "port",
@@ -1862,10 +1863,16 @@ def test_mssql_mssql_python_connection_string_with_odbc_properties():
             odbc_properties={
                 "Authentication": "ActiveDirectoryServicePrincipal",
                 "ClientCertificate": "/path/to/cert.pem",
-                "TrustServerCertificate": "NO",  # This should be ignored since we set it explicitly
+                # These should be ignored since we set them explicitly:
+                "TrustServerCertificate": "NO",
+                "ConnectRetryCount": 3,
+                "ConnectRetryInterval": 15,
             },
-            trust_server_certificate=True,  # This should take precedence
             check_import=False,
+            # These should take precedence:
+            trust_server_certificate=True,
+            login_timeout=10,
+            login_attempts=2,
         )
 
         # Get the connection factory with kwargs and call it
@@ -1880,11 +1887,15 @@ def test_mssql_mssql_python_connection_string_with_odbc_properties():
         assert "Authentication=ActiveDirectoryServicePrincipal" in conn_str
         assert "ClientCertificate=/path/to/cert.pem" in conn_str
 
-        # Verify that explicit trust_server_certificate takes precedence
+        # Verify that explicit kwargs take precedence over odbc_properties
         assert "TrustServerCertificate=yes" in conn_str
+        assert "ConnectRetryCount=2" in conn_str
+        assert "ConnectRetryInterval=10" in conn_str
 
-        # Should not have the conflicting property from odbc_properties
+        # Should not have conflicting properties from odbc_properties
         assert conn_str.count("TrustServerCertificate") == 1
+        assert conn_str.count("ConnectRetryCount") == 1
+        assert conn_str.count("ConnectRetryInterval") == 1
 
 
 @pytest.mark.xfail(not SUPPORTS_MSSQL_PYTHON_DRIVER, reason="mssql-python driver not supported")