fix(mysql): fix mysql query schema inference

Author: cpcloud

ci/schema/mysql.sql

@@ -66,7 +66,7 @@ CREATE TABLE functional_alltypes (
     double_col DOUBLE,
     date_string_col TEXT,
     string_col TEXT,
-    timestamp_col TIMESTAMP,
+    timestamp_col DATETIME,
     year INTEGER,
     month INTEGER
 ) DEFAULT CHARACTER SET = utf8;

ibis/backends/base/sql/alchemy/datatypes.py

@@ -165,16 +165,6 @@ def _(itype, **_):
     return MapType(to_sqla_type(itype.key_type), to_sqla_type(itype.value_type))
 
 
-@to_sqla_type.register(dt.GeoSpatial)
-def _(itype, **kwargs):
-    if itype.geotype == 'geometry':
-        return ga.Geometry
-    elif itype.geotype == 'geography':
-        return ga.Geography
-    else:
-        return ga.types._GISType
-
-
 @dt.dtype.register(Dialect, sa.types.NullType)
 def sa_null(_, satype, nullable=True):
     return dt.null
@@ -185,21 +175,12 @@ def sa_boolean(_, satype, nullable=True):
     return dt.Boolean(nullable=nullable)
 
 
-@dt.dtype.register(MySQLDialect, mysql.NUMERIC)
-@dt.dtype.register(MySQLDialect, sa.NUMERIC)
+@dt.dtype.register(MySQLDialect, (sa.NUMERIC, mysql.NUMERIC))
 def sa_mysql_numeric(_, satype, nullable=True):
     # https://dev.mysql.com/doc/refman/8.0/en/fixed-point-types.html
     return dt.Decimal(satype.precision or 10, satype.scale or 0, nullable=nullable)
 
 
-@dt.dtype.register(MySQLDialect, mysql.TINYBLOB)
-@dt.dtype.register(MySQLDialect, mysql.MEDIUMBLOB)
-@dt.dtype.register(MySQLDialect, mysql.BLOB)
-@dt.dtype.register(MySQLDialect, mysql.LONGBLOB)
-def sa_mysql_blob(_, satype, nullable=True):
-    return dt.Binary(nullable=nullable)
-
-
 _FLOAT_PREC_TO_TYPE = {
     11: dt.Float16,
     24: dt.Float32,
@@ -233,15 +214,30 @@ def sa_integer(_, satype, nullable=True):
 
 @dt.dtype.register(Dialect, mysql.TINYINT)
 @dt.dtype.register(MSDialect, mssql.TINYINT)
+@dt.dtype.register(MySQLDialect, mysql.YEAR)
 def sa_mysql_tinyint(_, satype, nullable=True):
     return dt.Int8(nullable=nullable)
 
 
 @dt.dtype.register(MSDialect, mssql.BIT)
-def sa_mysql_bit(_, satype, nullable=True):
+def sa_mssql_bit(_, satype, nullable=True):
     return dt.Boolean(nullable=nullable)
 
 
+@dt.dtype.register(MySQLDialect, mysql.BIT)
+def sa_mysql_bit(_, satype, nullable=True):
+    if 1 <= (length := satype.length) <= 8:
+        return dt.Int8(nullable=nullable)
+    elif 9 <= length <= 16:
+        return dt.Int16(nullable=nullable)
+    elif 17 <= length <= 32:
+        return dt.Int32(nullable=nullable)
+    elif 33 <= length <= 64:
+        return dt.Int64(nullable=nullable)
+    else:
+        raise ValueError(f"Invalid MySQL BIT length: {length:d}")
+
+
 @dt.dtype.register(Dialect, sa.types.BigInteger)
 @dt.dtype.register(MSDialect, mssql.MONEY)
 def sa_bigint(_, satype, nullable=True):
@@ -254,6 +250,7 @@ def sa_mssql_smallmoney(_, satype, nullable=True):
 
 
 @dt.dtype.register(Dialect, sa.REAL)
+@dt.dtype.register(MySQLDialect, mysql.FLOAT)
 def sa_real(_, satype, nullable=True):
     return dt.Float32(nullable=nullable)
 
@@ -271,11 +268,6 @@ def sa_uuid(_, satype, nullable=True):
     return dt.UUID(nullable=nullable)
 
 
-@dt.dtype.register(MSDialect, (mssql.BINARY, mssql.TIMESTAMP))
-def sa_mssql_timestamp(_, satype, nullable=True):
-    return dt.Binary(nullable=nullable)
-
-
 @dt.dtype.register(PGDialect, postgresql.MACADDR)
 def sa_macaddr(_, satype, nullable=True):
     return dt.MACADDR(nullable=nullable)
@@ -292,6 +284,21 @@ def sa_json(_, satype, nullable=True):
     return dt.JSON(nullable=nullable)
 
 
+@dt.dtype.register(MySQLDialect, mysql.TIMESTAMP)
+def sa_mysql_timestamp(_, satype, nullable=True):
+    return dt.Timestamp(timezone="UTC", nullable=nullable)
+
+
+@dt.dtype.register(MySQLDialect, mysql.DATETIME)
+def sa_mysql_datetime(_, satype, nullable=True):
+    return dt.Timestamp(nullable=nullable)
+
+
+@dt.dtype.register(MySQLDialect, mysql.SET)
+def sa_mysql_set(_, satype, nullable=True):
+    return dt.Set(dt.string, nullable=nullable)
+
+
 if geospatial_supported:
 
     @dt.dtype.register(Dialect, (ga.Geometry, ga.types._GISType))
@@ -314,6 +321,15 @@ def ga_geometry(_, gatype, nullable=True):
         else:
             raise ValueError(f"Unrecognized geometry type: {t}")
 
+    @to_sqla_type.register(dt.GeoSpatial)
+    def _(itype, **kwargs):
+        if itype.geotype == 'geometry':
+            return ga.Geometry
+        elif itype.geotype == 'geography':
+            return ga.Geography
+        else:
+            return ga.types._GISType
+
 
 POSTGRES_FIELD_TO_IBIS_UNIT = {
     "YEAR": "Y",
@@ -357,6 +373,18 @@ def sa_string(_, satype, nullable=True):
 
 
 @dt.dtype.register(Dialect, sa.LargeBinary)
+@dt.dtype.register(MSDialect, (mssql.BINARY, mssql.TIMESTAMP))
+@dt.dtype.register(
+    MySQLDialect,
+    (
+        mysql.TINYBLOB,
+        mysql.MEDIUMBLOB,
+        mysql.BLOB,
+        mysql.LONGBLOB,
+        mysql.BINARY,
+        mysql.VARBINARY,
+    ),
+)
 def sa_binary(_, satype, nullable=True):
     return dt.Binary(nullable=nullable)
 

ibis/backends/mssql/tests/test_client.py

@@ -3,6 +3,7 @@
 
 import ibis
 import ibis.expr.datatypes as dt
+from ibis.backends.base.sql.alchemy.geospatial import geospatial_supported
 
 DB_TYPES = [
     # Exact numbers
@@ -44,18 +45,23 @@
     ('IMAGE', dt.binary),
     # Other data types
     ('UNIQUEIDENTIFIER', dt.uuid),
-    ('GEOMETRY', dt.geometry),
-    ('GEOGRAPHY', dt.geography),
     ('TIMESTAMP', dt.binary(nullable=False)),
 ]
 
 
+skipif_no_geospatial_deps = pytest.mark.skipif(
+    not geospatial_supported, reason="geospatial dependencies not installed"
+)
+
+
 @pytest.mark.parametrize(
     ("server_type", "expected_type"),
-    [
-        param(server_type, ibis_type, id=server_type)
-        for server_type, ibis_type in DB_TYPES
+    DB_TYPES
+    + [
+        param("GEOMETRY", dt.geometry, marks=[skipif_no_geospatial_deps]),
+        param("GEOGRAPHY", dt.geography, marks=[skipif_no_geospatial_deps]),
     ],
+    ids=str,
 )
 def test_get_schema_from_query(con, server_type, expected_type):
     raw_name = f"tmp_{ibis.util.guid()}"

ibis/backends/mysql/__init__.py

@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
-import atexit
 import contextlib
+import re
 import warnings
 from typing import Iterable, Literal
 
@@ -121,8 +121,14 @@ def begin(self):
                     warnings.warn(f"Couldn't reset MySQL timezone: {str(e)}")
 
     def _metadata(self, query: str) -> Iterable[tuple[str, dt.DataType]]:
-        with self.con.begin() as con:
-            result = con.execute(f"SELECT * FROM ({query}) _ LIMIT 0")
+        if (
+            re.search(r"^\s*SELECT\s", query, flags=re.MULTILINE | re.IGNORECASE)
+            is not None
+        ):
+            query = f"({query})"
+
+        with self.begin() as con:
+            result = con.execute(f"SELECT * FROM {query} _ LIMIT 0")
             cursor = result.cursor
             yield from (
                 (field.name, _type_from_cursor_info(descr, field))
@@ -136,15 +142,6 @@ def _get_temp_view_definition(
     ) -> str:
         return f"CREATE OR REPLACE VIEW {name} AS {definition}"
 
-    def _register_temp_view_cleanup(self, name: str, raw_name: str) -> None:
-        query = f"DROP VIEW IF EXISTS {name}"
-
-        def drop(self, raw_name: str, query: str):
-            self.con.execute(query)
-            self._temp_views.discard(raw_name)
-
-        atexit.register(drop, self, raw_name, query)
-
 
 # TODO(kszucs): unsigned integers
 

ibis/backends/mysql/datatypes.py

@@ -29,15 +29,13 @@ def _type_from_cursor_info(descr, field) -> dt.DataType:
     if typename is None:
         raise NotImplementedError(f"MySQL type code {type_code:d} is not supported")
 
-    typ = _type_mapping[typename]
-
     if typename in ("DECIMAL", "NEWDECIMAL"):
         precision = _decimal_length_to_precision(
             length=field_length,
             scale=scale,
             is_unsigned=flags.is_unsigned,
         )
-        typ = partial(typ, precision=precision, scale=scale)
+        typ = partial(_type_mapping[typename], precision=precision, scale=scale)
     elif typename == "BIT":
         if field_length <= 8:
             typ = dt.int8
@@ -61,18 +59,15 @@ def _type_from_cursor_info(descr, field) -> dt.DataType:
                 typ = dt.Binary
             else:
                 typ = dt.String
+        else:
+            typ = _type_mapping[typename]
 
     # projection columns are always nullable
     return typ(nullable=True)
 
 
 # ported from my_decimal.h:my_decimal_length_to_precision in mariadb
-def _decimal_length_to_precision(
-    *,
-    length: int,
-    scale: int,
-    is_unsigned: bool,
-) -> int:
+def _decimal_length_to_precision(*, length: int, scale: int, is_unsigned: bool) -> int:
     return length - (scale > 0) - (not (is_unsigned or not length))
 
 
@@ -115,18 +110,14 @@ def _decimal_length_to_precision(
     "FLOAT": dt.Float32,
     "DOUBLE": dt.Float64,
     "NULL": dt.Null,
-    "TIMESTAMP": lambda nullable: dt.Timestamp(
-        timezone="UTC",
-        nullable=nullable,
-    ),
+    "TIMESTAMP": lambda nullable: dt.Timestamp(timezone="UTC", nullable=nullable),
     "LONGLONG": dt.Int64,
     "INT24": dt.Int32,
     "DATE": dt.Date,
     "TIME": dt.Time,
     "DATETIME": dt.Timestamp,
-    "YEAR": dt.Int16,
+    "YEAR": dt.Int8,
     "VARCHAR": dt.String,
-    "BIT": dt.Int8,
     "JSON": dt.JSON,
     "NEWDECIMAL": dt.Decimal,
     "ENUM": dt.String,

ibis/backends/mysql/tests/test_client.py

@@ -27,7 +27,7 @@
     ("date", dt.date),
     ("time", dt.time),
     ("datetime", dt.timestamp),
-    ("year", dt.int16),
+    ("year", dt.int8),
     ("char(32)", dt.string),
     ("char byte", dt.binary),
     ("varchar(42)", dt.string),
@@ -64,7 +64,9 @@ def test_get_schema_from_query(con, mysql_type, expected_type):
     # don't need to explicitly drop the table
     con.raw_sql(f"CREATE TEMPORARY TABLE {name} (x {mysql_type})")
     expected_schema = ibis.schema(dict(x=expected_type))
+    t = con.table(raw_name)
     result_schema = con._get_schema_using_query(f"SELECT * FROM {name}")
+    assert t.schema() == expected_schema
     assert result_schema == expected_schema