sqlmesh 0.191.1.dev1__tar.gz → 0.191.1.dev2__tar.gz

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.191.1.dev1
+Version: 0.191.1.dev2
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License
@@ -241,6 +241,8 @@ Provides-Extra: athena
 Requires-Dist: PyAthena[Pandas]; extra == "athena"
 Provides-Extra: azuresql
 Requires-Dist: pymssql; extra == "azuresql"
+Provides-Extra: azuresql-odbc
+Requires-Dist: pyodbc; extra == "azuresql-odbc"
 Provides-Extra: bigquery
 Requires-Dist: google-cloud-bigquery[pandas]; extra == "bigquery"
 Requires-Dist: google-cloud-bigquery-storage; extra == "bigquery"
@@ -311,6 +313,8 @@ Requires-Dist: langchain; extra == "llm"
 Requires-Dist: openai; extra == "llm"
 Provides-Extra: mssql
 Requires-Dist: pymssql; extra == "mssql"
+Provides-Extra: mssql-odbc
+Requires-Dist: pyodbc; extra == "mssql-odbc"
 Provides-Extra: mysql
 Requires-Dist: pymysql; extra == "mysql"
 Provides-Extra: mwaa

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/pyproject.toml

@@ -39,6 +39,7 @@ classifiers = [
 [project.optional-dependencies]
 athena = ["PyAthena[Pandas]"]
 azuresql = ["pymssql"]
+azuresql-odbc = ["pyodbc"]
 bigquery = [
     "google-cloud-bigquery[pandas]",
     "google-cloud-bigquery-storage"
@@ -104,6 +105,7 @@ gcppostgres = ["cloud-sql-python-connector[pg8000]>=1.8.0"]
 github = ["PyGithub~=2.5.0"]
 llm = ["langchain", "openai"]
 mssql = ["pymssql"]
+mssql-odbc = ["pyodbc"]
 mysql = ["pymysql"]
 mwaa = ["boto3"]
 postgres = ["psycopg2"]
@@ -203,6 +205,7 @@ module = [
     "databricks_cli.*",
     "mysql.*",
     "pymssql.*",
+    "pyodbc.*",
     "psycopg2.*",
     "langchain.*",
     "pytest_lazyfixture.*",

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.191.1.dev1'
-__version_tuple__ = version_tuple = (0, 191, 1, 'dev1')
+__version__ = version = '0.191.1.dev2'
+__version_tuple__ = version_tuple = (0, 191, 1, 'dev2')

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/sqlmesh/core/config/connection.py

@@ -55,11 +55,10 @@ MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
 
 
 def _get_engine_import_validator(
-    import_name: str, engine_type: str, extra_name: t.Optional[str] = None
+    import_name: str, engine_type: str, extra_name: t.Optional[str] = None, decorate: bool = True
 ) -> t.Callable:
     extra_name = extra_name or engine_type
 
-    @model_validator(mode="before")
     def validate(cls: t.Any, data: t.Any) -> t.Any:
         check_import = (
             str_to_bool(str(data.pop("check_import", True))) if isinstance(data, dict) else True
@@ -83,7 +82,7 @@ def _get_engine_import_validator(
 
         return data
 
-    return validate
+    return model_validator(mode="before")(validate) if decorate else validate
 
 
 class ConnectionConfig(abc.ABC, BaseConfig):
@@ -1422,17 +1421,50 @@ class MSSQLConnectionConfig(ConnectionConfig):
     autocommit: t.Optional[bool] = False
     tds_version: t.Optional[str] = None
 
+    # Driver options
+    driver: t.Literal["pymssql", "pyodbc"] = "pymssql"
+    # PyODBC specific options
+    driver_name: t.Optional[str] = None  # e.g. "ODBC Driver 18 for SQL Server"
+    trust_server_certificate: t.Optional[bool] = None
+    encrypt: t.Optional[bool] = None
+    # Dictionary of arbitrary ODBC connection properties
+    # See: https://learn.microsoft.com/en-us/sql/connect/odbc/dsn-connection-string-attribute
+    odbc_properties: t.Optional[t.Dict[str, t.Any]] = None
+
     concurrent_tasks: int = 4
     register_comments: bool = True
     pre_ping: bool = True
 
     type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
 
-    _engine_import_validator = _get_engine_import_validator("pymssql", "mssql")
+    @model_validator(mode="before")
+    @classmethod
+    def _mssql_engine_import_validator(cls, data: t.Any) -> t.Any:
+        if not isinstance(data, dict):
+            return data
+
+        driver = data.get("driver", "pymssql")
+
+        # Define the mapping of driver to import module and extra name
+        driver_configs = {"pymssql": ("pymssql", "mssql"), "pyodbc": ("pyodbc", "mssql-odbc")}
+
+        if driver not in driver_configs:
+            raise ValueError(f"Unsupported driver: {driver}")
+
+        import_module, extra_name = driver_configs[driver]
+
+        # Use _get_engine_import_validator with decorate=False to get the raw validation function
+        # This avoids the __wrapped__ issue in Python 3.9
+        validator_func = _get_engine_import_validator(
+            import_module, driver, extra_name, decorate=False
+        )
+
+        # Call the raw validation function directly
+        return validator_func(cls, data)
 
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
-        return {
+        base_keys = {
             "host",
             "user",
             "password",
@@ -1447,15 +1479,96 @@ class MSSQLConnectionConfig(ConnectionConfig):
             "tds_version",
         }
 
+        if self.driver == "pyodbc":
+            base_keys.update(
+                {
+                    "driver_name",
+                    "trust_server_certificate",
+                    "encrypt",
+                    "odbc_properties",
+                }
+            )
+            # Remove pymssql-specific parameters
+            base_keys.discard("tds_version")
+            base_keys.discard("conn_properties")
+
+        return base_keys
+
     @property
     def _engine_adapter(self) -> t.Type[EngineAdapter]:
         return engine_adapter.MSSQLEngineAdapter
 
     @property
     def _connection_factory(self) -> t.Callable:
-        import pymssql
+        if self.driver == "pymssql":
+            import pymssql
+
+            return pymssql.connect
+
+        import pyodbc
+
+        def connect(**kwargs: t.Any) -> t.Callable:
+            # Extract parameters for connection string
+            host = kwargs.pop("host")
+            port = kwargs.pop("port", 1433)
+            database = kwargs.pop("database", "")
+            user = kwargs.pop("user", None)
+            password = kwargs.pop("password", None)
+            driver_name = kwargs.pop("driver_name", "ODBC Driver 18 for SQL Server")
+            trust_server_certificate = kwargs.pop("trust_server_certificate", False)
+            encrypt = kwargs.pop("encrypt", True)
+            login_timeout = kwargs.pop("login_timeout", 60)
+
+            # Build connection string
+            conn_str_parts = [
+                f"DRIVER={{{driver_name}}}",
+                f"SERVER={host},{port}",
+            ]
+
+            if database:
+                conn_str_parts.append(f"DATABASE={database}")
+
+            # Add security options
+            conn_str_parts.append(f"Encrypt={'YES' if encrypt else 'NO'}")
+            if trust_server_certificate:
+                conn_str_parts.append("TrustServerCertificate=YES")
+
+            conn_str_parts.append(f"Connection Timeout={login_timeout}")
+
+            # Standard SQL Server authentication
+            if user:
+                conn_str_parts.append(f"UID={user}")
+            if password:
+                conn_str_parts.append(f"PWD={password}")
+
+            # Add any additional ODBC properties from the odbc_properties dictionary
+            if self.odbc_properties:
+                for key, value in self.odbc_properties.items():
+                    # Skip properties that we've already set above
+                    if key.lower() in (
+                        "driver",
+                        "server",
+                        "database",
+                        "uid",
+                        "pwd",
+                        "encrypt",
+                        "trustservercertificate",
+                        "connection timeout",
+                    ):
+                        continue
 
-        return pymssql.connect
+                    # Handle boolean values properly
+                    if isinstance(value, bool):
+                        conn_str_parts.append(f"{key}={'YES' if value else 'NO'}")
+                    else:
+                        conn_str_parts.append(f"{key}={value}")
+
+            # Create the connection string
+            conn_str = ";".join(conn_str_parts)
+
+            return pyodbc.connect(conn_str, autocommit=kwargs.get("autocommit", False))
+
+        return connect
 
     @property
     def _extra_engine_config(self) -> t.Dict[str, t.Any]:

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/sqlmesh/core/engine_adapter/mssql.py

@@ -219,6 +219,10 @@ class MSSQLEngineAdapter(
         assert isinstance(df, pd.DataFrame)
         temp_table = self._get_temp_table(target_table or "pandas")
 
+        # Return the superclass implementation if the connection pool doesn't support bulk_copy
+        if not hasattr(self._connection_pool.get(), "bulk_copy"):
+            return super()._df_to_source_queries(df, columns_to_types, batch_size, target_table)
+
         def query_factory() -> Query:
             # It is possible for the factory to be called multiple times and if so then the temp table will already
             # be created so we skip creating again. This means we are assuming the first call is the same result

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.191.1.dev1
+Version: 0.191.1.dev2
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License
@@ -241,6 +241,8 @@ Provides-Extra: athena
 Requires-Dist: PyAthena[Pandas]; extra == "athena"
 Provides-Extra: azuresql
 Requires-Dist: pymssql; extra == "azuresql"
+Provides-Extra: azuresql-odbc
+Requires-Dist: pyodbc; extra == "azuresql-odbc"
 Provides-Extra: bigquery
 Requires-Dist: google-cloud-bigquery[pandas]; extra == "bigquery"
 Requires-Dist: google-cloud-bigquery-storage; extra == "bigquery"
@@ -311,6 +313,8 @@ Requires-Dist: langchain; extra == "llm"
 Requires-Dist: openai; extra == "llm"
 Provides-Extra: mssql
 Requires-Dist: pymssql; extra == "mssql"
+Provides-Extra: mssql-odbc
+Requires-Dist: pyodbc; extra == "mssql-odbc"
 Provides-Extra: mysql
 Requires-Dist: pymysql; extra == "mysql"
 Provides-Extra: mwaa

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/sqlmesh.egg-info/requires.txt

@@ -26,6 +26,9 @@ PyAthena[Pandas]
 [azuresql]
 pymssql
 
+[azuresql-odbc]
+pyodbc
+
 [bigframes]
 bigframes>=1.32.0
 
@@ -116,6 +119,9 @@ lsprotocol
 [mssql]
 pymssql
 
+[mssql-odbc]
+pyodbc
+
 [mwaa]
 boto3
 

{sqlmesh-0.191.1.dev1 → sqlmesh-0.191.1.dev2}/tests/core/test_connection_config.py

@@ -4,6 +4,7 @@ import typing as t
 
 import pytest
 from _pytest.fixtures import FixtureRequest
+from unittest.mock import patch
 
 from sqlmesh.core.config.connection import (
     BigQueryConnectionConfig,
@@ -19,6 +20,7 @@ from sqlmesh.core.config.connection import (
     SnowflakeConnectionConfig,
     TrinoAuthenticationMethod,
     AthenaConnectionConfig,
+    MSSQLConnectionConfig,
     _connection_config_validator,
     _get_engine_import_validator,
 )
@@ -1127,3 +1129,266 @@ def test_engine_import_validator():
         _engine_import_validator = _get_engine_import_validator("sqlmesh", "bigquery")
 
     TestConfigC()
+
+
+def test_mssql_engine_import_validator():
+    """Test that MSSQL import validator respects driver configuration."""
+
+    # Test PyODBC driver suggests mssql-odbc extra when import fails
+    with pytest.raises(ConfigError, match=r"pip install \"sqlmesh\[mssql-odbc\]\""):
+        with patch("importlib.import_module") as mock_import:
+            mock_import.side_effect = ImportError("No module named 'pyodbc'")
+            MSSQLConnectionConfig(host="localhost", driver="pyodbc")
+
+    # Test PyMSSQL driver suggests mssql extra when import fails
+    with pytest.raises(ConfigError, match=r"pip install \"sqlmesh\[mssql\]\""):
+        with patch("importlib.import_module") as mock_import:
+            mock_import.side_effect = ImportError("No module named 'pymssql'")
+            MSSQLConnectionConfig(host="localhost", driver="pymssql")
+
+    # Test default driver (pymssql) suggests mssql extra when import fails
+    with pytest.raises(ConfigError, match=r"pip install \"sqlmesh\[mssql\]\""):
+        with patch("importlib.import_module") as mock_import:
+            mock_import.side_effect = ImportError("No module named 'pymssql'")
+            MSSQLConnectionConfig(host="localhost")  # No driver specified
+
+    # Test successful import works without error
+    with patch("importlib.import_module") as mock_import:
+        mock_import.return_value = None
+        config = MSSQLConnectionConfig(host="localhost", driver="pyodbc")
+        assert config.driver == "pyodbc"
+
+
+def test_mssql_connection_config_parameter_validation(make_config):
+    """Test MSSQL connection config parameter validation."""
+    # Test default driver is pymssql
+    config = make_config(type="mssql", host="localhost", check_import=False)
+    assert isinstance(config, MSSQLConnectionConfig)
+    assert config.driver == "pymssql"
+
+    # Test explicit pyodbc driver
+    config = make_config(type="mssql", host="localhost", driver="pyodbc", check_import=False)
+    assert isinstance(config, MSSQLConnectionConfig)
+    assert config.driver == "pyodbc"
+
+    # Test explicit pymssql driver
+    config = make_config(type="mssql", host="localhost", driver="pymssql", check_import=False)
+    assert isinstance(config, MSSQLConnectionConfig)
+    assert config.driver == "pymssql"
+
+    # Test pyodbc specific parameters
+    config = make_config(
+        type="mssql",
+        host="localhost",
+        driver="pyodbc",
+        driver_name="ODBC Driver 18 for SQL Server",
+        trust_server_certificate=True,
+        encrypt=False,
+        odbc_properties={"Authentication": "ActiveDirectoryServicePrincipal"},
+        check_import=False,
+    )
+    assert isinstance(config, MSSQLConnectionConfig)
+    assert config.driver_name == "ODBC Driver 18 for SQL Server"
+    assert config.trust_server_certificate is True
+    assert config.encrypt is False
+    assert config.odbc_properties == {"Authentication": "ActiveDirectoryServicePrincipal"}
+
+    # Test pymssql specific parameters
+    config = make_config(
+        type="mssql",
+        host="localhost",
+        driver="pymssql",
+        tds_version="7.4",
+        conn_properties=["SET ANSI_NULLS ON"],
+        check_import=False,
+    )
+    assert isinstance(config, MSSQLConnectionConfig)
+    assert config.tds_version == "7.4"
+    assert config.conn_properties == ["SET ANSI_NULLS ON"]
+
+
+def test_mssql_connection_kwargs_keys():
+    """Test _connection_kwargs_keys returns correct keys for each driver variant."""
+    # Test pymssql driver keys
+    config = MSSQLConnectionConfig(host="localhost", driver="pymssql", check_import=False)
+    pymssql_keys = config._connection_kwargs_keys
+    expected_pymssql_keys = {
+        "password",
+        "user",
+        "database",
+        "host",
+        "timeout",
+        "login_timeout",
+        "charset",
+        "appname",
+        "port",
+        "tds_version",
+        "conn_properties",
+        "autocommit",
+    }
+    assert pymssql_keys == expected_pymssql_keys
+
+    # Test pyodbc driver keys
+    config = MSSQLConnectionConfig(host="localhost", driver="pyodbc", check_import=False)
+    pyodbc_keys = config._connection_kwargs_keys
+    expected_pyodbc_keys = {
+        "password",
+        "user",
+        "database",
+        "host",
+        "timeout",
+        "login_timeout",
+        "charset",
+        "appname",
+        "port",
+        "autocommit",
+        "driver_name",
+        "trust_server_certificate",
+        "encrypt",
+        "odbc_properties",
+    }
+    assert pyodbc_keys == expected_pyodbc_keys
+
+    # Verify pyodbc keys don't include pymssql-specific parameters
+    assert "tds_version" not in pyodbc_keys
+    assert "conn_properties" not in pyodbc_keys
+
+
+def test_mssql_pyodbc_connection_string_generation():
+    """Test pyodbc.connect gets invoked with the correct ODBC connection string."""
+    with patch("pyodbc.connect") as mock_pyodbc_connect:
+        # Mock the return value to have the methods we need
+        mock_connection = mock_pyodbc_connect.return_value
+
+        # Create a pyodbc config
+        config = MSSQLConnectionConfig(
+            host="testserver.database.windows.net",
+            port=1433,
+            database="testdb",
+            user="testuser",
+            password="testpass",
+            driver="pyodbc",
+            driver_name="ODBC Driver 18 for SQL Server",
+            trust_server_certificate=True,
+            encrypt=True,
+            login_timeout=30,
+            check_import=False,
+        )
+
+        # Get the connection factory with kwargs and call it
+        factory_with_kwargs = config._connection_factory_with_kwargs
+        connection = factory_with_kwargs()
+
+        # Verify pyodbc.connect was called with the correct connection string
+        mock_pyodbc_connect.assert_called_once()
+        call_args = mock_pyodbc_connect.call_args
+
+        # Check the connection string (first argument)
+        conn_str = call_args[0][0]
+        expected_parts = [
+            "DRIVER={ODBC Driver 18 for SQL Server}",
+            "SERVER=testserver.database.windows.net,1433",
+            "DATABASE=testdb",
+            "Encrypt=YES",
+            "TrustServerCertificate=YES",
+            "Connection Timeout=30",
+            "UID=testuser",
+            "PWD=testpass",
+        ]
+
+        for part in expected_parts:
+            assert part in conn_str
+
+        # Check autocommit parameter
+        assert call_args[1]["autocommit"] is False
+
+
+def test_mssql_pyodbc_connection_string_with_odbc_properties():
+    """Test pyodbc connection string includes custom ODBC properties."""
+    with patch("pyodbc.connect") as mock_pyodbc_connect:
+        # Create a pyodbc config with custom ODBC properties
+        config = MSSQLConnectionConfig(
+            host="testserver.database.windows.net",
+            database="testdb",
+            user="client-id",
+            password="client-secret",
+            driver="pyodbc",
+            odbc_properties={
+                "Authentication": "ActiveDirectoryServicePrincipal",
+                "ClientCertificate": "/path/to/cert.pem",
+                "TrustServerCertificate": "NO",  # This should be ignored since we set it explicitly
+            },
+            trust_server_certificate=True,  # This should take precedence
+            check_import=False,
+        )
+
+        # Get the connection factory with kwargs and call it
+        factory_with_kwargs = config._connection_factory_with_kwargs
+        connection = factory_with_kwargs()
+
+        # Verify pyodbc.connect was called
+        mock_pyodbc_connect.assert_called_once()
+        conn_str = mock_pyodbc_connect.call_args[0][0]
+
+        # Check that custom ODBC properties are included
+        assert "Authentication=ActiveDirectoryServicePrincipal" in conn_str
+        assert "ClientCertificate=/path/to/cert.pem" in conn_str
+
+        # Verify that explicit trust_server_certificate takes precedence
+        assert "TrustServerCertificate=YES" in conn_str
+
+        # Should not have the conflicting property from odbc_properties
+        assert conn_str.count("TrustServerCertificate") == 1
+
+
+def test_mssql_pyodbc_connection_string_minimal():
+    """Test pyodbc connection string with minimal configuration."""
+    with patch("pyodbc.connect") as mock_pyodbc_connect:
+        config = MSSQLConnectionConfig(
+            host="localhost",
+            driver="pyodbc",
+            autocommit=True,
+            check_import=False,
+        )
+
+        factory_with_kwargs = config._connection_factory_with_kwargs
+        connection = factory_with_kwargs()
+
+        mock_pyodbc_connect.assert_called_once()
+        conn_str = mock_pyodbc_connect.call_args[0][0]
+
+        # Check basic required parts
+        assert "DRIVER={ODBC Driver 18 for SQL Server}" in conn_str
+        assert "SERVER=localhost,1433" in conn_str
+        assert "Encrypt=YES" in conn_str  # Default encrypt=True
+        assert "Connection Timeout=60" in conn_str  # Default timeout
+
+        # Check autocommit parameter
+        assert mock_pyodbc_connect.call_args[1]["autocommit"] is True
+
+
+def test_mssql_pymssql_connection_factory():
+    """Test pymssql connection factory returns correct function."""
+    # Mock the import of pymssql at the module level
+    import sys
+    from unittest.mock import MagicMock
+
+    # Create a mock pymssql module
+    mock_pymssql = MagicMock()
+    sys.modules["pymssql"] = mock_pymssql
+
+    try:
+        config = MSSQLConnectionConfig(
+            host="localhost",
+            driver="pymssql",
+            check_import=False,
+        )
+
+        factory = config._connection_factory
+
+        # Verify the factory returns pymssql.connect
+        assert factory is mock_pymssql.connect
+    finally:
+        # Clean up the mock module
+        if "pymssql" in sys.modules:
+            del sys.modules["pymssql"]