Merge branch 'master' into fix-mysql-wire-protocol-column-types-42203

Address review feedback questioning the setattr()+noqa pattern in the
pymysql/mysql-connector fallback tests. ModuleType needed setattr()
because mypy flags plain attribute assignment on it as attr-defined,
which is what the noqa: B010 (ruff's set-attr-with-constant) was
suppressing. SimpleNamespace supports normal attribute construction/
access under mypy, so the fake objects can be built without setattr
or any lint suppression.

superset/db_engine_specs/base.py

@@ -1340,12 +1340,12 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
                 return cursor.fetchmany(limit)
             data = cursor.fetchall()
             description = cursor.description or []
-            # Create a mapping between column name and a mutator function to normalize
-            # values with. The first two items in the description row are
-            # the column name and type.
+            # Create a mapping between column index and a mutator function to normalize
+            # values with. The first two items in the description row are the column
+            # name and type.
             column_mutators = {
-                row[0]: func
-                for row in description
+                index: func
+                for index, row in enumerate(description)
                 if (
                     func := cls.column_type_mutators.get(
                         type(cls.get_sqla_column_type(cls.get_datatype(row[1])))
@@ -1353,11 +1353,11 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
                 )
             }
             if column_mutators:
-                indexes = {row[0]: idx for idx, row in enumerate(description)}
+                if not isinstance(data, list):
+                    data = list(data)
                 for row_idx, row in enumerate(data):
                     new_row = list(row)
-                    for col, func in column_mutators.items():
-                        col_idx = indexes[col]
+                    for col_idx, func in column_mutators.items():
                         new_row[col_idx] = func(row[col_idx])
                     data[row_idx] = tuple(new_row)
 

superset/db_engine_specs/mysql.py

@@ -249,6 +249,43 @@ class MySQLEngineSpec(BasicParametersMixin, BaseEngineSpec):
             types.VARCHAR(),
             GenericDataType.STRING,
         ),
+        # wire-protocol FIELD_TYPE names emitted by `get_datatype`, seen on
+        # SQL Lab and virtual dataset columns instead of DDL type names
+        (
+            re.compile(r"^newdecimal", re.IGNORECASE),
+            DECIMAL(),
+            GenericDataType.NUMERIC,
+        ),
+        (
+            re.compile(r"^tiny$", re.IGNORECASE),
+            TINYINT(),
+            GenericDataType.NUMERIC,
+        ),
+        (
+            re.compile(r"^short$", re.IGNORECASE),
+            types.SmallInteger(),
+            GenericDataType.NUMERIC,
+        ),
+        (
+            re.compile(r"^(blob|text)$", re.IGNORECASE),
+            types.String(),
+            GenericDataType.STRING,
+        ),
+        (
+            re.compile(r"^year$", re.IGNORECASE),
+            types.Integer(),
+            GenericDataType.NUMERIC,
+        ),
+        (
+            re.compile(r"^enum\b", re.IGNORECASE),
+            types.String(),
+            GenericDataType.STRING,
+        ),
+        (
+            re.compile(r"^set\b", re.IGNORECASE),
+            types.String(),
+            GenericDataType.STRING,
+        ),
     )
     column_type_mutators: dict[types.TypeEngine, Callable[[Any], Any]] = {
         DECIMAL: lambda val: Decimal(val) if isinstance(val, str) else val
@@ -407,22 +444,27 @@ class MySQLEngineSpec(BasicParametersMixin, BaseEngineSpec):
 
     @classmethod
     def get_datatype(cls, type_code: Any) -> Optional[str]:
-        if not cls.type_code_map:
-            # only import and store if needed at least once
-            # pylint: disable=import-outside-toplevel
-            try:
-                import MySQLdb
-
-                mysql_module = MySQLdb
-            except ImportError:
-                mysql_module = __import__("pymysql")
-
-            ft = mysql_module.constants.FIELD_TYPE
-            cls.type_code_map = {
-                getattr(ft, k): k for k in dir(ft) if not k.startswith("_")
-            }
         datatype = type_code
         if isinstance(type_code, int):
+            if not cls.type_code_map:
+                # only import and store if needed at least once
+                # pylint: disable=import-outside-toplevel
+                try:
+                    import MySQLdb
+
+                    ft = MySQLdb.constants.FIELD_TYPE
+                except ImportError:
+                    try:
+                        import pymysql  # type: ignore[import-untyped]
+
+                        ft = pymysql.constants.FIELD_TYPE
+                    except ImportError:
+                        from mysql.connector.constants import FieldType
+
+                        ft = FieldType
+                cls.type_code_map = {
+                    getattr(ft, k): k for k in dir(ft) if not k.startswith("_")
+                }
             datatype = cls.type_code_map.get(type_code)
         if datatype and isinstance(datatype, str) and datatype:
             return datatype

tests/unit_tests/db_engine_specs/test_mysql.py

@@ -18,7 +18,7 @@
 import builtins
 from datetime import datetime
 from decimal import Decimal
-from types import ModuleType
+from types import SimpleNamespace
 from typing import Any, Optional
 from unittest.mock import Mock, patch
 
@@ -73,6 +73,16 @@ from tests.unit_tests.fixtures.common import dttm  # noqa: F401
         ("DATETIME", types.DateTime, None, GenericDataType.TEMPORAL, True),
         ("TIMESTAMP", types.TIMESTAMP, None, GenericDataType.TEMPORAL, True),
         ("TIME", types.Time, None, GenericDataType.TEMPORAL, True),
+        # Wire-protocol names
+        ("VAR_STRING", types.VARCHAR, None, GenericDataType.STRING, False),
+        ("NEWDECIMAL", DECIMAL, None, GenericDataType.NUMERIC, False),
+        ("TINY", TINYINT, None, GenericDataType.NUMERIC, False),
+        ("SHORT", types.SmallInteger, None, GenericDataType.NUMERIC, False),
+        ("BLOB", types.String, None, GenericDataType.STRING, False),
+        ("TEXT", types.String, None, GenericDataType.STRING, False),
+        ("YEAR", types.Integer, None, GenericDataType.NUMERIC, False),
+        ("ENUM", types.String, None, GenericDataType.STRING, False),
+        ("SET", types.String, None, GenericDataType.STRING, False),
     ],
 )
 def test_get_column_spec(
@@ -87,6 +97,50 @@ def test_get_column_spec(
     assert_column_spec(spec, native_type, sqla_type, attrs, generic_type, is_dttm)
 
 
+def test_fetch_data_mutates_decimal_rows_in_tuple_results() -> None:
+    from superset.db_engine_specs.mysql import MySQLEngineSpec as spec  # noqa: N813
+
+    newdecimal, var_string = 246, 253
+    cursor = Mock()
+    cursor.description = [("amount", newdecimal), ("label", var_string)]
+    cursor.fetchall.return_value = (("10.50", "Ships"), ("22.30", "Planes"))
+
+    # Stub the type_code_map so this test doesn't depend on MySQLdb or
+    # pymysql being importable in the test environment.
+    original_type_code_map = spec.type_code_map
+    spec.type_code_map = {newdecimal: "NEWDECIMAL", var_string: "VAR_STRING"}
+
+    try:
+        data = spec.fetch_data(cursor)
+    finally:
+        spec.type_code_map = original_type_code_map
+
+    assert data == [(Decimal("10.50"), "Ships"), (Decimal("22.30"), "Planes")]
+
+
+def test_fetch_data_mutates_duplicate_decimal_column_names() -> None:
+    from superset.db_engine_specs.mysql import MySQLEngineSpec as spec  # noqa: N813
+
+    newdecimal, var_string = 246, 253
+    cursor = Mock()
+    cursor.description = [
+        ("amount", newdecimal),
+        ("amount", var_string),
+        ("amount", newdecimal),
+    ]
+    cursor.fetchall.return_value = [("10.50", "not a decimal", "22.30")]
+
+    original_type_code_map = spec.type_code_map
+    spec.type_code_map = {newdecimal: "NEWDECIMAL", var_string: "VAR_STRING"}
+
+    try:
+        data = spec.fetch_data(cursor)
+    finally:
+        spec.type_code_map = original_type_code_map
+
+    assert data == [(Decimal("10.50"), "not a decimal", Decimal("22.30"))]
+
+
 @pytest.mark.parametrize(
     "target_type,expected_result",
     [
@@ -269,7 +323,7 @@ def test_column_type_mutator(
     assert spec.fetch_data(mock_cursor) == expected_result
 
 
-def test_get_datatype_pymysql_fallback():
+def test_get_datatype_pymysql_fallback() -> None:
     """get_datatype() falls back to pymysql when MySQLdb is not installed."""
     from superset.db_engine_specs.mysql import MySQLEngineSpec
 
@@ -279,15 +333,9 @@ def test_get_datatype_pymysql_fallback():
 
     try:
         # Build a fake pymysql module with constants.FIELD_TYPE
-        fake_field_type = ModuleType("pymysql.constants.FIELD_TYPE")
-        fake_field_type.TINY = 1
-        fake_field_type.VARCHAR = 15
-
-        fake_constants = ModuleType("pymysql.constants")
-        fake_constants.FIELD_TYPE = fake_field_type
-
-        fake_pymysql = ModuleType("pymysql")
-        fake_pymysql.constants = fake_constants
+        fake_field_type = SimpleNamespace(TINY=1, VARCHAR=15)
+        fake_constants = SimpleNamespace(FIELD_TYPE=fake_field_type)
+        fake_pymysql = SimpleNamespace(constants=fake_constants)
 
         original_import = builtins.__import__
 
@@ -308,6 +356,31 @@ def test_get_datatype_pymysql_fallback():
         MySQLEngineSpec.type_code_map = original_type_code_map
 
 
+def test_get_datatype_mysqlconnector_fallback() -> None:
+    """get_datatype() supports mysql-connector-python without PyMySQL."""
+    from superset.db_engine_specs.mysql import MySQLEngineSpec
+
+    original_type_code_map = MySQLEngineSpec.type_code_map
+    MySQLEngineSpec.type_code_map = {}
+
+    try:
+        fake_field_type = SimpleNamespace(NEWDECIMAL=246)
+        fake_constants = SimpleNamespace(FieldType=fake_field_type)
+        original_import = builtins.__import__
+
+        def mock_import(name: str, *args: Any, **kwargs: Any) -> Any:
+            if name in {"MySQLdb", "pymysql"}:
+                raise ImportError(f"No module named '{name}'")
+            if name == "mysql.connector.constants":
+                return fake_constants
+            return original_import(name, *args, **kwargs)
+
+        with patch("builtins.__import__", side_effect=mock_import):
+            assert MySQLEngineSpec.get_datatype(246) == "NEWDECIMAL"
+    finally:
+        MySQLEngineSpec.type_code_map = original_type_code_map
+
+
 @pytest.mark.parametrize(
     ("grain", "expected_expression"),
     [