diff --git a/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py b/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py index 7b6e42fff3..c537f5a853 100644 --- a/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py +++ b/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py @@ -1,5 +1,6 @@ from typing import Any, Union, TYPE_CHECKING import datetime +import decimal import pymssql from pymssql import OperationalError @@ -59,6 +60,38 @@ def _map_type(mssql_type_text: str) -> MYSQL_DATA_TYPE: return MYSQL_DATA_TYPE.VARCHAR +# pyodbc reports the Python type of each column in cursor.description[1] +_ODBC_PYTHON_TYPES = { + bool: MYSQL_DATA_TYPE.TINYINT, + int: MYSQL_DATA_TYPE.INT, + float: MYSQL_DATA_TYPE.FLOAT, + decimal.Decimal: MYSQL_DATA_TYPE.DECIMAL, + datetime.datetime: MYSQL_DATA_TYPE.DATETIME, + datetime.date: MYSQL_DATA_TYPE.DATE, + datetime.time: MYSQL_DATA_TYPE.TIME, + str: MYSQL_DATA_TYPE.TEXT, + bytes: MYSQL_DATA_TYPE.BINARY, + bytearray: MYSQL_DATA_TYPE.BINARY, +} + + +def _nullable_dtype(values: list[Any]) -> str | None: + """Pick a pandas nullable dtype for integer or bit values, so NULLs don't turn the column into float or object. + + Args: + values (list[Any]): the values of one column. + + Returns: + str | None: 'Int64' or 'boolean' if all non-NULL values are int or bool, otherwise None. + """ + value_types = {type(value) for value in values if value is not None} + if value_types == {int}: + return "Int64" + if value_types == {bool}: + return "boolean" + return None + + def _make_table_response( result: list[Union[dict[str, Any], tuple]], cursor: Union[pymssql.Cursor, "pyodbc.Cursor"], use_odbc: bool = False ) -> Response: @@ -85,13 +118,27 @@ def _make_table_response( # pymssql with as_dict=True returns list of dicts data_frame = pd.DataFrame(result) - for column in description: + for column_index, column in enumerate(description): column_name = column[0] column_type = column[1] column_dtype = data_frame[column_name].dtype - if use_odbc: - # For pyodbc, use type inference based on pandas dtype + if ( + result + and column_type in (pymssql.NUMBER, int, bool) + and not (pd_types.is_integer_dtype(column_dtype) or pd_types.is_bool_dtype(column_dtype)) + ): + key = column_index if use_odbc else column_name + values = [row[key] for row in result] + nullable_dtype = _nullable_dtype(values) + if nullable_dtype is not None: + data_frame[column_name] = pd.array(values, dtype=nullable_dtype) + column_dtype = data_frame[column_name].dtype + + if use_odbc and column_type in _ODBC_PYTHON_TYPES: + mysql_types.append(_ODBC_PYTHON_TYPES[column_type]) + elif use_odbc: + # Fall back to type inference based on pandas dtype if pd_types.is_integer_dtype(column_dtype): mysql_types.append(MYSQL_DATA_TYPE.INT) elif pd_types.is_float_dtype(column_dtype): diff --git a/tests/unit/handlers/test_mssql.py b/tests/unit/handlers/test_mssql.py index d7024d5135..38192a7eb8 100644 --- a/tests/unit/handlers/test_mssql.py +++ b/tests/unit/handlers/test_mssql.py @@ -932,6 +932,43 @@ def test_types_casting(self): self.assertEqual(result_value, input_value) # endregion + def test_types_casting_with_nulls(self): + """Test that NULLs in integer and bit columns don't turn them into float or object""" + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.__enter__ = MagicMock(return_value=mock_cursor) + mock_cursor.__exit__ = MagicMock(return_value=None) + + self.handler.connect = MagicMock(return_value=mock_conn) + mock_conn.cursor = MagicMock(return_value=mock_cursor) + + mock_cursor.fetchall.return_value = [ + {"n_bigint": 9007199254740993, "n_int": 3, "n_bit": True, "n_float": 1.5}, + {"n_bigint": None, "n_int": None, "n_bit": None, "n_float": None}, + ] + mock_cursor.description = [ + ("n_bigint", 3, None, None, None, None, None), + ("n_int", 3, None, None, None, None, None), + ("n_bit", 3, None, None, None, None, None), + ("n_float", 3, None, None, None, None, None), + ] + + response: TableResponse = self.handler.native_query("SELECT * FROM test_table") + + self.assertEqual( + [col.type for col in response.columns], + [MYSQL_DATA_TYPE.INT, MYSQL_DATA_TYPE.INT, MYSQL_DATA_TYPE.TINYINT, MYSQL_DATA_TYPE.FLOAT], + ) + df = response.data_frame + self.assertEqual(df["n_bigint"].dtype, "Int64") + self.assertEqual(df["n_bigint"][0], 9007199254740993) + self.assertEqual(df["n_int"].dtype, "Int64") + self.assertEqual(df["n_int"][0], 3) + self.assertEqual(df["n_bit"].dtype, "boolean") + self.assertTrue(df["n_bit"][0]) + self.assertEqual(df["n_float"].dtype, "float64") + self.assertTrue(df[["n_bigint", "n_int", "n_bit", "n_float"]].iloc[1].isna().all()) + class TestMSSQLHandlerODBC(unittest.TestCase): """Tests for MSSQL handler with ODBC connection""" @@ -1185,6 +1222,57 @@ def __getitem__(self, idx): if "pyodbc" in sys.modules: del sys.modules["pyodbc"] + def test_odbc_types_from_cursor_description(self): + """Test that ODBC column types come from cursor.description and NULLs don't turn integers into float""" + mock_pyodbc = MagicMock() + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.__enter__ = MagicMock(return_value=mock_cursor) + mock_cursor.__exit__ = MagicMock(return_value=None) + + mock_cursor.fetchall.return_value = [ + (None, None, None, None, None, None), + (9007199254740993, Decimal("10.50"), datetime.date(2026, 1, 2), datetime.time(12, 30), True, "x"), + ] + mock_cursor.description = [ + ("n_bigint", int, None, 19, 19, 0, True), + ("n_decimal", Decimal, None, 19, 19, 4, True), + ("d_date", datetime.date, None, 10, 10, 0, True), + ("d_time", datetime.time, None, 16, 16, 7, True), + ("n_bit", bool, None, 1, 1, 0, True), + ("t_text", str, None, 10, 10, 0, True), + ] + + sys.modules["pyodbc"] = mock_pyodbc + + try: + handler = SqlServerHandler("mssql_odbc", connection_data=self.connection_data) + handler.connect = MagicMock(return_value=mock_conn) + handler.is_connected = True + mock_conn.cursor = MagicMock(return_value=mock_cursor) + + response = handler.native_query("SELECT * FROM test") + + self.assertEqual( + [col.type for col in response.columns], + [ + MYSQL_DATA_TYPE.INT, + MYSQL_DATA_TYPE.DECIMAL, + MYSQL_DATA_TYPE.DATE, + MYSQL_DATA_TYPE.TIME, + MYSQL_DATA_TYPE.TINYINT, + MYSQL_DATA_TYPE.TEXT, + ], + ) + df = response.data_frame + self.assertEqual(df["n_bigint"].dtype, "Int64") + self.assertEqual(df["n_bigint"][1], 9007199254740993) + self.assertEqual(df["n_bit"].dtype, "boolean") + self.assertTrue(df.iloc[0].isna().all()) + finally: + if "pyodbc" in sys.modules: + del sys.modules["pyodbc"] + if __name__ == "__main__": unittest.main()