Fix: Add DATETIMEOFFSET handling for MSSQL cursor initialization

Author: mattiasthalen

sqlmesh/core/config/connection.py

@@ -1520,6 +1520,41 @@ 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 = {

tests/core/test_connection_config.py

@@ -1521,3 +1521,184 @@ def test_mssql_pymssql_connection_factory():
         # Clean up the mock module
         if "pymssql" in sys.modules:
             del sys.modules["pymssql"]
+
+
+def test_mssql_cursor_init_datetimeoffset_handling():
+    """Test that the MSSQL cursor init properly handles DATETIMEOFFSET conversion."""
+    from datetime import datetime, timezone, timedelta
+    import struct
+    from unittest.mock import Mock
+
+    config = MSSQLConnectionConfig(
+        host="localhost",
+        driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
+        check_import=False,
+    )
+
+    # 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,
+    )
+
+    # 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."""
+    from datetime import datetime, timezone, timedelta
+    import struct
+    from unittest.mock import Mock
+
+    config = MSSQLConnectionConfig(
+        host="localhost",
+        driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
+        check_import=False,
+    )
+
+    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,
+    )
+
+    result = converter_func(binary_data)
+
+    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))
+
+
+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
+
+    config = MSSQLConnectionConfig(
+        host="localhost",
+        driver="pyodbc",  # DATETIMEOFFSET handling is pyodbc-specific
+        check_import=False,
+    )
+
+    cursor_init = config._cursor_init
+    assert cursor_init is not None
+
+    # 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")
+
+    # This should not raise an exception
+    cursor_init(mock_cursor)
+
+
+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,
+    )
+
+    cursor_init = config._cursor_init
+    assert cursor_init is not None
+    assert callable(cursor_init)
+
+
+def test_mssql_cursor_init_returns_none_for_pymssql():
+    """Test that _cursor_init returns None for pymssql driver."""
+    config = MSSQLConnectionConfig(
+        host="localhost",
+        driver="pymssql",
+        check_import=False,
+    )
+
+    cursor_init = config._cursor_init
+    assert cursor_init is None