Fix: Move DATETIMEOFFSET handling to MSSQL connection initialization

Author: mattiasthalen

sqlmesh/core/config/connection.py

@@ -1520,41 +1520,6 @@ def _mssql_engine_import_validator(cls, data: t.Any) -> t.Any:
         # Call the raw validation function directly
         return validator_func(cls, data)
 
-    @property
-    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
-        """Initialize the cursor with output converters for MSSQL-specific data types."""
-        # Only apply pyodbc-specific cursor initialization when using pyodbc driver
-        if self.driver != "pyodbc":
-            return None
-
-        def init(cursor: t.Any) -> None:
-            # Get the connection from the cursor and set the output converter
-            conn = cursor.connection
-            if hasattr(conn, "add_output_converter"):
-                # Handle SQL type -155 (DATETIMEOFFSET) which is not yet supported by pyodbc
-                # ref: https://github.com/mkleehammer/pyodbc/issues/134#issuecomment-281739794
-                def handle_datetimeoffset(dto_value: t.Any) -> t.Any:
-                    from datetime import datetime, timedelta, timezone
-                    import struct
-
-                    # Unpack the DATETIMEOFFSET binary format:
-                    # Format: <6hI2h = (year, month, day, hour, minute, second, nanoseconds, tz_hour_offset, tz_minute_offset)
-                    tup = struct.unpack("<6hI2h", dto_value)
-                    return datetime(
-                        tup[0],
-                        tup[1],
-                        tup[2],
-                        tup[3],
-                        tup[4],
-                        tup[5],
-                        tup[6] // 1000,
-                        timezone(timedelta(hours=tup[7], minutes=tup[8])),
-                    )
-
-                conn.add_output_converter(-155, handle_datetimeoffset)
-
-        return init
-
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         base_keys = {
@@ -1659,7 +1624,33 @@ def connect(**kwargs: t.Any) -> t.Callable:
             # Create the connection string
             conn_str = ";".join(conn_str_parts)
 
-            return pyodbc.connect(conn_str, autocommit=kwargs.get("autocommit", False))
+            conn = pyodbc.connect(conn_str, autocommit=kwargs.get("autocommit", False))
+
+            # Set up output converters for MSSQL-specific data types
+            if hasattr(conn, "add_output_converter"):
+                # Handle SQL type -155 (DATETIMEOFFSET) which is not yet supported by pyodbc
+                # ref: https://github.com/mkleehammer/pyodbc/issues/134#issuecomment-281739794
+                def handle_datetimeoffset(dto_value: t.Any) -> t.Any:
+                    from datetime import datetime, timedelta, timezone
+                    import struct
+
+                    # Unpack the DATETIMEOFFSET binary format:
+                    # Format: <6hI2h = (year, month, day, hour, minute, second, nanoseconds, tz_hour_offset, tz_minute_offset)
+                    tup = struct.unpack("<6hI2h", dto_value)
+                    return datetime(
+                        tup[0],
+                        tup[1],
+                        tup[2],
+                        tup[3],
+                        tup[4],
+                        tup[5],
+                        tup[6] // 1000,
+                        timezone(timedelta(hours=tup[7], minutes=tup[8])),
+                    )
+
+                conn.add_output_converter(-155, handle_datetimeoffset)
+
+            return conn
 
         return connect
 

tests/core/test_connection_config.py

@@ -1523,182 +1523,166 @@ def test_mssql_pymssql_connection_factory():
             del sys.modules["pymssql"]
 
 
-def test_mssql_cursor_init_datetimeoffset_handling():
-    """Test that the MSSQL cursor init properly handles DATETIMEOFFSET conversion."""
+def test_mssql_pyodbc_connection_datetimeoffset_handling():
+    """Test that the MSSQL pyodbc connection properly handles DATETIMEOFFSET conversion."""
     from datetime import datetime, timezone, timedelta
     import struct
-    from unittest.mock import Mock
+    from unittest.mock import Mock, patch
 
-    config = MSSQLConnectionConfig(
-        host="localhost",
-        driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
-        check_import=False,
-    )
+    with patch("pyodbc.connect") as mock_pyodbc_connect:
+        # Track calls to add_output_converter
+        converter_calls = []
 
-    # Get the cursor init function
-    cursor_init = config._cursor_init
-    assert cursor_init is not None
-
-    # Create a mock cursor and connection
-    mock_connection = Mock()
-    mock_cursor = Mock()
-    mock_cursor.connection = mock_connection
-
-    # Track calls to add_output_converter
-    converter_calls = []
-
-    def mock_add_output_converter(sql_type, converter_func):
-        converter_calls.append((sql_type, converter_func))
-
-    mock_connection.add_output_converter = mock_add_output_converter
-
-    # Call the cursor init function
-    cursor_init(mock_cursor)
-
-    # Verify that add_output_converter was called for SQL type -155 (DATETIMEOFFSET)
-    assert len(converter_calls) == 1
-    sql_type, converter_func = converter_calls[0]
-    assert sql_type == -155
-
-    # Test the converter function with actual DATETIMEOFFSET binary data
-    # Create a test DATETIMEOFFSET value: 2023-12-25 15:30:45.123456789 +05:30
-    year, month, day = 2023, 12, 25
-    hour, minute, second = 15, 30, 45
-    nanoseconds = 123456789
-    tz_hour_offset, tz_minute_offset = 5, 30
-
-    # Pack the binary data according to the DATETIMEOFFSET format
-    binary_data = struct.pack(
-        "<6hI2h",
-        year,
-        month,
-        day,
-        hour,
-        minute,
-        second,
-        nanoseconds,
-        tz_hour_offset,
-        tz_minute_offset,
-    )
+        def mock_add_output_converter(sql_type, converter_func):
+            converter_calls.append((sql_type, converter_func))
 
-    # Convert using the registered converter
-    result = converter_func(binary_data)
-
-    # Verify the result
-    expected_dt = datetime(
-        2023,
-        12,
-        25,
-        15,
-        30,
-        45,
-        123456,  # microseconds = nanoseconds // 1000
-        timezone(timedelta(hours=5, minutes=30)),
-    )
-    assert result == expected_dt
-    assert result.tzinfo == timezone(timedelta(hours=5, minutes=30))
+        # Create a mock connection that will be returned by pyodbc.connect
+        mock_connection = Mock()
+        mock_connection.add_output_converter = mock_add_output_converter
+        mock_pyodbc_connect.return_value = mock_connection
+
+        config = MSSQLConnectionConfig(
+            host="localhost",
+            driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
+            check_import=False,
+        )
 
+        # Get the connection factory and call it
+        factory_with_kwargs = config._connection_factory_with_kwargs
+        connection = factory_with_kwargs()
+
+        # Verify that add_output_converter was called for SQL type -155 (DATETIMEOFFSET)
+        assert len(converter_calls) == 1
+        sql_type, converter_func = converter_calls[0]
+        assert sql_type == -155
+
+        # Test the converter function with actual DATETIMEOFFSET binary data
+        # Create a test DATETIMEOFFSET value: 2023-12-25 15:30:45.123456789 +05:30
+        year, month, day = 2023, 12, 25
+        hour, minute, second = 15, 30, 45
+        nanoseconds = 123456789
+        tz_hour_offset, tz_minute_offset = 5, 30
+
+        # Pack the binary data according to the DATETIMEOFFSET format
+        binary_data = struct.pack(
+            "<6hI2h",
+            year,
+            month,
+            day,
+            hour,
+            minute,
+            second,
+            nanoseconds,
+            tz_hour_offset,
+            tz_minute_offset,
+        )
+
+        # Convert using the registered converter
+        result = converter_func(binary_data)
+
+        # Verify the result
+        expected_dt = datetime(
+            2023,
+            12,
+            25,
+            15,
+            30,
+            45,
+            123456,  # microseconds = nanoseconds // 1000
+            timezone(timedelta(hours=5, minutes=30)),
+        )
+        assert result == expected_dt
+        assert result.tzinfo == timezone(timedelta(hours=5, minutes=30))
 
-def test_mssql_cursor_init_negative_timezone_offset():
-    """Test DATETIMEOFFSET handling with negative timezone offset."""
+
+def test_mssql_pyodbc_connection_negative_timezone_offset():
+    """Test DATETIMEOFFSET handling with negative timezone offset at connection level."""
     from datetime import datetime, timezone, timedelta
     import struct
-    from unittest.mock import Mock
+    from unittest.mock import Mock, patch
 
-    config = MSSQLConnectionConfig(
-        host="localhost",
-        driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
-        check_import=False,
-    )
+    with patch("pyodbc.connect") as mock_pyodbc_connect:
+        converter_calls = []
 
-    cursor_init = config._cursor_init
-    mock_connection = Mock()
-    mock_cursor = Mock()
-    mock_cursor.connection = mock_connection
-
-    converter_calls = []
-
-    def mock_add_output_converter(sql_type, converter_func):
-        converter_calls.append((sql_type, converter_func))
-
-    mock_connection.add_output_converter = mock_add_output_converter
-    cursor_init(mock_cursor)
-
-    # Get the converter function
-    _, converter_func = converter_calls[0]
-
-    # Test with negative timezone offset: 2023-01-01 12:00:00.0 -08:00
-    year, month, day = 2023, 1, 1
-    hour, minute, second = 12, 0, 0
-    nanoseconds = 0
-    tz_hour_offset, tz_minute_offset = -8, 0
-
-    binary_data = struct.pack(
-        "<6hI2h",
-        year,
-        month,
-        day,
-        hour,
-        minute,
-        second,
-        nanoseconds,
-        tz_hour_offset,
-        tz_minute_offset,
-    )
+        def mock_add_output_converter(sql_type, converter_func):
+            converter_calls.append((sql_type, converter_func))
 
-    result = converter_func(binary_data)
+        mock_connection = Mock()
+        mock_connection.add_output_converter = mock_add_output_converter
+        mock_pyodbc_connect.return_value = mock_connection
 
-    expected_dt = datetime(2023, 1, 1, 12, 0, 0, 0, timezone(timedelta(hours=-8, minutes=0)))
-    assert result == expected_dt
-    assert result.tzinfo == timezone(timedelta(hours=-8))
+        config = MSSQLConnectionConfig(
+            host="localhost",
+            driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
+            check_import=False,
+        )
 
+        factory_with_kwargs = config._connection_factory_with_kwargs
+        connection = factory_with_kwargs()
 
-def test_mssql_cursor_init_no_add_output_converter():
-    """Test that cursor init gracefully handles connections without add_output_converter."""
-    from unittest.mock import Mock
+        # Get the converter function
+        _, converter_func = converter_calls[0]
+
+        # Test with negative timezone offset: 2023-01-01 12:00:00.0 -08:00
+        year, month, day = 2023, 1, 1
+        hour, minute, second = 12, 0, 0
+        nanoseconds = 0
+        tz_hour_offset, tz_minute_offset = -8, 0
+
+        binary_data = struct.pack(
+            "<6hI2h",
+            year,
+            month,
+            day,
+            hour,
+            minute,
+            second,
+            nanoseconds,
+            tz_hour_offset,
+            tz_minute_offset,
+        )
 
-    config = MSSQLConnectionConfig(
-        host="localhost",
-        driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
-        check_import=False,
-    )
+        result = converter_func(binary_data)
 
-    cursor_init = config._cursor_init
-    assert cursor_init is not None
+        expected_dt = datetime(2023, 1, 1, 12, 0, 0, 0, timezone(timedelta(hours=-8, minutes=0)))
+        assert result == expected_dt
+        assert result.tzinfo == timezone(timedelta(hours=-8))
 
-    # Create a mock cursor and connection without add_output_converter
-    mock_connection = Mock()
-    mock_cursor = Mock()
-    mock_cursor.connection = mock_connection
 
-    # Remove the add_output_converter attribute
-    if hasattr(mock_connection, "add_output_converter"):
-        delattr(mock_connection, "add_output_converter")
+def test_mssql_pyodbc_connection_no_add_output_converter():
+    """Test that connection gracefully handles pyodbc without add_output_converter."""
+    from unittest.mock import Mock, patch
 
-    # This should not raise an exception
-    cursor_init(mock_cursor)
+    with patch("pyodbc.connect") as mock_pyodbc_connect:
+        # Create a mock connection without add_output_converter
+        mock_connection = Mock()
+        # Remove the add_output_converter attribute
+        if hasattr(mock_connection, "add_output_converter"):
+            delattr(mock_connection, "add_output_converter")
+        mock_pyodbc_connect.return_value = mock_connection
 
+        config = MSSQLConnectionConfig(
+            host="localhost",
+            driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
+            check_import=False,
+        )
 
-def test_mssql_cursor_init_returns_callable_for_pyodbc():
-    """Test that _cursor_init returns a callable function for pyodbc driver."""
-    config = MSSQLConnectionConfig(
-        host="localhost",
-        driver="pyodbc",
-        check_import=False,
-    )
+        # This should not raise an exception
+        factory_with_kwargs = config._connection_factory_with_kwargs
+        connection = factory_with_kwargs()
 
-    cursor_init = config._cursor_init
-    assert cursor_init is not None
-    assert callable(cursor_init)
+        # Verify we get the connection back
+        assert connection is mock_connection
 
 
-def test_mssql_cursor_init_returns_none_for_pymssql():
-    """Test that _cursor_init returns None for pymssql driver."""
+def test_mssql_no_cursor_init_for_pymssql():
+    """Test that _cursor_init is not needed for pymssql driver."""
     config = MSSQLConnectionConfig(
         host="localhost",
         driver="pymssql",
         check_import=False,
     )
 
-    cursor_init = config._cursor_init
-    assert cursor_init is None
+    # Since we moved output converter setup to connection level,
+    # there's no cursor init needed for any driver
+    assert not hasattr(config, "_cursor_init") or config._cursor_init is None