diff --git a/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py b/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py index 7b6e42fff3..43ab0bcdfc 100644 --- a/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py +++ b/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py @@ -150,6 +150,24 @@ def _make_table_response( return Response(RESPONSE_TYPE.TABLE, data_frame=data_frame, mysql_types=mysql_types) +def _odbc_value(value: Any) -> str: + """Escape a value for an ODBC connection string. + + A value that contains ';', '{' or '}', or that starts or ends with a space, is wrapped in braces + with every '}' doubled, so it can't end the attribute early or add other attributes. + + Args: + value (Any): the attribute value. + + Returns: + str: the value, safe to use after 'KEY='. + """ + value = str(value) + if any(char in value for char in ";{}") or value != value.strip(): + return "{" + value.replace("}", "}}") + "}" + return value + + class SqlServerHandler(MetaDatabaseHandler): """ This handler handles connection and execution of the Microsoft SQL Server statements. @@ -252,19 +270,19 @@ def _connect_odbc(self): password = self.connection_args.get("password") conn_str_parts = [ - f"DRIVER={{{driver}}}", - f"SERVER={host},{port}", - f"DATABASE={database}", - f"UID={user}", - f"PWD={password}", + "DRIVER={" + str(driver).replace("}", "}}") + "}", + f"SERVER={_odbc_value(f'{host},{port}')}", + f"DATABASE={_odbc_value(database)}", + f"UID={_odbc_value(user)}", + f"PWD={_odbc_value(password)}", ] # Add optional parameters if "encrypt" in self.connection_args: - conn_str_parts.append(f"Encrypt={self.connection_args.get('encrypt', 'yes')}") + conn_str_parts.append(f"Encrypt={_odbc_value(self.connection_args.get('encrypt', 'yes'))}") if "trust_server_certificate" in self.connection_args: conn_str_parts.append( - f"TrustServerCertificate={self.connection_args.get('trust_server_certificate', 'yes')}" + f"TrustServerCertificate={_odbc_value(self.connection_args.get('trust_server_certificate', 'yes'))}" ) if "connection_string_args" in self.connection_args: diff --git a/tests/unit/handlers/test_mssql.py b/tests/unit/handlers/test_mssql.py index d7024d5135..c63f50d4d6 100644 --- a/tests/unit/handlers/test_mssql.py +++ b/tests/unit/handlers/test_mssql.py @@ -998,6 +998,33 @@ def test_odbc_connection_string_construction(self): if "pyodbc" in sys.modules: del sys.modules["pyodbc"] + def test_odbc_connection_string_escapes_special_characters(self): + """Test that values with ';', '{', '}' or edge spaces can't split the ODBC connection string""" + mock_pyodbc = MagicMock() + mock_connect = MagicMock() + mock_pyodbc.connect = mock_connect + + sys.modules["pyodbc"] = mock_pyodbc + + connection_data = self.connection_data.copy() + connection_data["password"] = "Xy;9}Encrypt=no" + connection_data["user"] = "{admin" + connection_data["database"] = " sales " + connection_data["encrypt"] = "yes" + + try: + handler = SqlServerHandler("mssql_odbc", connection_data=connection_data) + handler.connect() + conn_str = mock_connect.call_args[0][0] + self.assertEqual( + conn_str, + "DRIVER={ODBC Driver 18 for SQL Server};SERVER=127.0.0.1,1433;DATABASE={ sales };" + "UID={{admin};PWD={Xy;9}}Encrypt=no};Encrypt=yes", + ) + finally: + if "pyodbc" in sys.modules: + del sys.modules["pyodbc"] + def test_odbc_connection_with_encryption_params(self): """Test that encryption parameters are added to connection string""" mock_pyodbc = MagicMock()