From 5f0b58ed6e0a64591f42ef0282a3d2cc90a087d6 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Fri, 14 Nov 2025 18:01:22 +0800 Subject: [PATCH 1/2] fix(ibis): handle JSON result for MySQL --- ibis-server/app/model/connector.py | 35 ++++++++++ .../routers/v3/connector/mysql/conftest.py | 52 +++++++++++++++ .../v3/connector/mysql/test_functions.py | 27 ++------ .../routers/v3/connector/mysql/test_query.py | 64 +++++++++++++++++++ 4 files changed, 157 insertions(+), 21 deletions(-) create mode 100644 ibis-server/tests/routers/v3/connector/mysql/test_query.py diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py index 9c00115c2..480caf0a7 100644 --- a/ibis-server/app/model/connector.py +++ b/ibis-server/app/model/connector.py @@ -88,6 +88,8 @@ def __init__(self, data_source: DataSource, connection_info: ConnectionInfo): self._connector = RedshiftConnector(connection_info) elif data_source == DataSource.postgres: self._connector = PostgresConnector(connection_info) + elif data_source == DataSource.mysql: + self._connector = MySqlConnector(connection_info) else: self._connector = SimpleConnector(data_source, connection_info) @@ -294,6 +296,39 @@ def close(self) -> None: self.connection = None +class MySqlConnector(SimpleConnector): + def __init__(self, connection_info: ConnectionInfo): + super().__init__(DataSource.mysql, connection_info) + + def _handle_pyarrow_unsupported_type(self, ibis_table: Table, **kwargs) -> Table: + result_table = ibis_table + for name, dtype in ibis_table.schema().items(): + if isinstance(dtype, Decimal): + # Round decimal columns to a specified scale + result_table = self._round_decimal_columns( + result_table=result_table, col_name=name, **kwargs + ) + elif isinstance(dtype, UUID): + # Convert UUID to string for compatibility + result_table = self._cast_uuid_columns( + result_table=result_table, col_name=name + ) + elif isinstance(dtype, dt.JSON): + # ibis doesn't handle JSON type for MySQL properly. + # We need to convert JSON columns to string for compatibility manually. + result_table = self._cast_json_columns( + result_table=result_table, col_name=name + ) + + return result_table + + def _cast_json_columns(self, result_table: Table, col_name: str) -> Table: + col = result_table[col_name] + # Convert JSON to string for compatibility + casted_col = col.cast("string") + return result_table.mutate(**{col_name: casted_col}) + + class MSSqlConnector(SimpleConnector): def __init__(self, connection_info: ConnectionInfo): super().__init__(DataSource.mssql, connection_info) diff --git a/ibis-server/tests/routers/v3/connector/mysql/conftest.py b/ibis-server/tests/routers/v3/connector/mysql/conftest.py index 54efe958f..feae1d1ac 100644 --- a/ibis-server/tests/routers/v3/connector/mysql/conftest.py +++ b/ibis-server/tests/routers/v3/connector/mysql/conftest.py @@ -1,8 +1,12 @@ import pathlib import pytest +import sqlalchemy from testcontainers.mysql import MySqlContainer +from app.config import get_config +from tests.conftest import file_path + pytestmark = pytest.mark.mysql base_url = "/v3/connector/mysql" @@ -18,10 +22,58 @@ def pytest_collection_modifyitems(items): @pytest.fixture(scope="session") def mysql(request) -> MySqlContainer: mysql = MySqlContainer(image="mysql:8.0.40", dialect="pymysql").start() + connection_url = mysql.get_connection_url() + engine = sqlalchemy.create_engine(connection_url) + with engine.connect() as conn: + conn.execute( + sqlalchemy.text( + """ + CREATE TABLE json_test ( + id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY, + object_col JSON NOT NULL, + array_col JSON NOT NULL, + CHECK (JSON_TYPE(object_col) = 'OBJECT'), + CHECK (JSON_TYPE(array_col) = 'ARRAY') + ) ENGINE=InnoDB; + """ + ) + ) + conn.execute( + sqlalchemy.text( + """ + INSERT INTO json_test (object_col, array_col) VALUES + ('{"name": "Alice", "age": 30, "city": "New York"}', '["apple", "banana", "cherry"]'), + ('{"name": "Bob", "age": 25, "city": "Los Angeles"}', '["dog", "cat", "mouse"]'), + ('{"name": "Charlie", "age": 35, "city": "Chicago"}', '["red", "green", "blue"]'); + """ + ) + ) + conn.commit() + request.addfinalizer(mysql.stop) return mysql +function_list_path = file_path("../resources/function_list") +white_function_list_path = file_path("../resources/white_function_list") + + +@pytest.fixture(autouse=True) +def set_remote_function_list_path(): + config = get_config() + config.set_remote_function_list_path(function_list_path) + yield + config.set_remote_function_list_path(None) + + +@pytest.fixture(autouse=True) +def set_remote_white_function_list_path(): + config = get_config() + config.set_remote_white_function_list_path(white_function_list_path) + yield + config.set_remote_white_function_list_path(None) + + @pytest.fixture(scope="module") def connection_info(mysql: MySqlContainer) -> dict[str, str]: return { diff --git a/ibis-server/tests/routers/v3/connector/mysql/test_functions.py b/ibis-server/tests/routers/v3/connector/mysql/test_functions.py index aef7404de..a34983bed 100644 --- a/ibis-server/tests/routers/v3/connector/mysql/test_functions.py +++ b/ibis-server/tests/routers/v3/connector/mysql/test_functions.py @@ -4,8 +4,12 @@ import pytest from app.config import get_config -from tests.conftest import DATAFUSION_FUNCTION_COUNT, file_path -from tests.routers.v3.connector.mysql.conftest import base_url +from tests.conftest import DATAFUSION_FUNCTION_COUNT +from tests.routers.v3.connector.mysql.conftest import ( + base_url, + function_list_path, + white_function_list_path, +) manifest = { "dataSource": "mysql", @@ -25,31 +29,12 @@ ], } -function_list_path = file_path("../resources/function_list") -white_function_list_path = file_path("../resources/white_function_list") - @pytest.fixture(scope="module") def manifest_str(): return base64.b64encode(orjson.dumps(manifest)).decode("utf-8") -@pytest.fixture(autouse=True) -def set_remote_function_list_path(): - config = get_config() - config.set_remote_function_list_path(function_list_path) - yield - config.set_remote_function_list_path(None) - - -@pytest.fixture(autouse=True) -def set_remote_white_function_list_path(): - config = get_config() - config.set_remote_white_function_list_path(white_function_list_path) - yield - config.set_remote_white_function_list_path(None) - - async def test_function_list(client): config = get_config() diff --git a/ibis-server/tests/routers/v3/connector/mysql/test_query.py b/ibis-server/tests/routers/v3/connector/mysql/test_query.py new file mode 100644 index 000000000..2a9bcfcc1 --- /dev/null +++ b/ibis-server/tests/routers/v3/connector/mysql/test_query.py @@ -0,0 +1,64 @@ +import pytest + +from app.dependencies import X_WREN_FALLBACK_DISABLE +from tests.routers.v3.connector.mysql.conftest import base_url + +manifest = { + "catalog": "wren", + "schema": "public", + "models": [ + { + "name": "json_test", + "tableReference": { + "table": "json_test", + }, + "columns": [ + {"name": "id", "type": "bigint"}, + {"name": "object_col", "type": "json"}, + {"name": "array_col", "type": "json"}, + ], + }, + { + "name": "orders", + "tableReference": { + "table": "orders", + }, + "columns": [ + {"name": "o_orderkey", "type": "integer"}, + {"name": "o_orderdate", "type": "date"}, + ], + }, + ], + "dataSource": "mysql", +} + + +@pytest.fixture(scope="module") +async def manifest_str(web_server): + return await web_server.register_mdl(manifest) + + +async def test_json_query(client, manifest_str, connection_info): + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT object_col, array_col FROM wren.public.json_test limit 1", + }, + headers={ + X_WREN_FALLBACK_DISABLE: "true", + }, + ) + assert response.status_code == 200 + result = response.json() + assert result["data"] == [ + [ + '{"age": 30, "city": "New York", "name": "Alice"}', + '["apple", "banana", "cherry"]', + ] + ] + assert result["dtypes"] == { + "object_col": "string", + "array_col": "string", + } From bc6d8a8a4edde2ae3ae3f09234fee9b893a8a9d1 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Fri, 14 Nov 2025 18:21:45 +0800 Subject: [PATCH 2/2] fix test --- ibis-server/tests/routers/v3/connector/mysql/test_query.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ibis-server/tests/routers/v3/connector/mysql/test_query.py b/ibis-server/tests/routers/v3/connector/mysql/test_query.py index 2a9bcfcc1..6f5e8de09 100644 --- a/ibis-server/tests/routers/v3/connector/mysql/test_query.py +++ b/ibis-server/tests/routers/v3/connector/mysql/test_query.py @@ -1,3 +1,6 @@ +import base64 + +import orjson import pytest from app.dependencies import X_WREN_FALLBACK_DISABLE @@ -34,8 +37,8 @@ @pytest.fixture(scope="module") -async def manifest_str(web_server): - return await web_server.register_mdl(manifest) +async def manifest_str(): + return base64.b64encode(orjson.dumps(manifest)).decode("utf-8") async def test_json_query(client, manifest_str, connection_info):