superset2/tests/unit_tests/result_set_test.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# pylint: disable=import-outside-toplevel, unused-argument

from datetime import datetime, timezone

import numpy as np
import pandas as pd
from numpy.core.multiarray import array
from pytest_mock import MockerFixture

from superset.db_engine_specs.base import BaseEngineSpec
from superset.result_set import stringify_values, SupersetResultSet
from superset.superset_typing import DbapiResult
from superset.utils import json as superset_json


def test_column_names_as_bytes() -> None:
    """
    Test that we can handle column names as bytes.
    """
    from superset.db_engine_specs.redshift import RedshiftEngineSpec
    from superset.result_set import SupersetResultSet

    data = (
        [
            "2016-01-26",
            392.002014,
            397.765991,
            390.575012,
            392.153015,
            392.153015,
            58147000,
        ],
        [
            "2016-01-27",
            392.444,
            396.842987,
            391.782013,
            394.971985,
            394.971985,
            47424400,
        ],
    )
    description = [
        (b"date", 1043, None, None, None, None, None),
        (b"open", 701, None, None, None, None, None),
        (b"high", 701, None, None, None, None, None),
        (b"low", 701, None, None, None, None, None),
        (b"close", 701, None, None, None, None, None),
        (b"adj close", 701, None, None, None, None, None),
        (b"volume", 20, None, None, None, None, None),
    ]
    result_set = SupersetResultSet(data, description, RedshiftEngineSpec)  # type: ignore

    assert (
        result_set.to_pandas_df().to_markdown()
        == """
|    | date       |    open |    high |     low |   close |   adj close |   volume |
|---:|:-----------|--------:|--------:|--------:|--------:|------------:|---------:|
|  0 | 2016-01-26 | 392.002 | 397.766 | 390.575 | 392.153 |     392.153 | 58147000 |
|  1 | 2016-01-27 | 392.444 | 396.843 | 391.782 | 394.972 |     394.972 | 47424400 |
    """.strip()
    )


def test_stringify_with_null_integers():
    """
    Test that we can safely handle type errors when an integer column has a null value
    """

    data = [
        ("foo", "bar", pd.NA, None),
        ("foo", "bar", pd.NA, True),
        ("foo", "bar", pd.NA, None),
    ]
    numpy_dtype = [
        ("id", "object"),
        ("value", "object"),
        ("num", "object"),
        ("bool", "object"),
    ]

    array2 = np.array(data, dtype=numpy_dtype)
    column_names = ["id", "value", "num", "bool"]

    result_set = np.array([stringify_values(array2[column]) for column in column_names])

    expected = np.array(
        [
            array(["foo", "foo", "foo"], dtype=object),
            array(["bar", "bar", "bar"], dtype=object),
            array([None, None, None], dtype=object),
            array([None, "True", None], dtype=object),
        ]
    )

    assert np.array_equal(result_set, expected)


def test_stringify_with_null_timestamps():
    """
    Test that we can safely handle type errors when a timestamp column has a null value
    """

    data = [
        ("foo", "bar", pd.NaT, None),
        ("foo", "bar", pd.NaT, True),
        ("foo", "bar", pd.NaT, None),
    ]
    numpy_dtype = [
        ("id", "object"),
        ("value", "object"),
        ("num", "object"),
        ("bool", "object"),
    ]

    array2 = np.array(data, dtype=numpy_dtype)
    column_names = ["id", "value", "num", "bool"]

    result_set = np.array([stringify_values(array2[column]) for column in column_names])

    expected = np.array(
        [
            array(["foo", "foo", "foo"], dtype=object),
            array(["bar", "bar", "bar"], dtype=object),
            array([None, None, None], dtype=object),
            array([None, "True", None], dtype=object),
        ]
    )

    assert np.array_equal(result_set, expected)


def test_timezone_series(mocker: MockerFixture) -> None:
    """
    Test that we can handle timezone-aware datetimes correctly.

    This covers a regression that happened when upgrading from Pandas 1.5.3 to 2.0.3.
    """
    logger = mocker.patch("superset.result_set.logger")

    data = [[datetime(2023, 1, 1, tzinfo=timezone.utc)]]
    description = [(b"__time", "datetime", None, None, None, None, False)]
    result_set = SupersetResultSet(
        data,
        description,  # type: ignore
        BaseEngineSpec,
    )
    assert result_set.to_pandas_df().values.tolist() == [
        [pd.Timestamp("2023-01-01 00:00:00+0000", tz="UTC")]
    ]
    logger.exception.assert_not_called()


def test_get_column_description_from_empty_data_using_cursor_description(
    mocker: MockerFixture,
) -> None:
    """
    Test that we can handle get_column_decription from the cursor description
    when data is empty
    """
    logger = mocker.patch("superset.result_set.logger")

    data: DbapiResult = []
    description = [(b"__time", "datetime", None, None, None, None, 1, 0, 255)]
    result_set = SupersetResultSet(
        data,
        description,  # type: ignore
        BaseEngineSpec,
    )
    assert any(col.get("column_name") == "__time" for col in result_set.columns)
    logger.exception.assert_not_called()


def test_empty_column_names_get_synthetic_names() -> None:
    """
    SQL Server returns an empty-string column name in cursor.description for
    any un-aliased expression (e.g. ``SELECT COUNT(*) FROM t``).  An empty
    field name is illegal in NumPy structured arrays and PyArrow tables.

    SupersetResultSet must replace empty column names with synthetic names
    so queries like ``SELECT COUNT(*) FROM t`` succeed on MSSQL.

    Regression test for https://github.com/apache/superset/issues/23848
    """
    data = [(42,)]
    description = [("", 3, None, None, None, None, None)]
    result_set = SupersetResultSet(data, description, BaseEngineSpec)  # type: ignore

    assert result_set.columns[0]["column_name"] == "_col_0"
    df = result_set.to_pandas_df()
    assert list(df.columns) == ["_col_0"]
    assert df["_col_0"].iloc[0] == 42


def test_multiple_empty_column_names_get_unique_synthetic_names() -> None:
    """
    When several columns have empty names (e.g. ``SELECT COUNT(*), SUM(x)``
    on MSSQL), each must receive a distinct synthetic name.
    """
    data = [(10, 20)]
    description = [
        ("", 3, None, None, None, None, None),
        ("", 3, None, None, None, None, None),
    ]
    result_set = SupersetResultSet(data, description, BaseEngineSpec)  # type: ignore

    col_names = [c["column_name"] for c in result_set.columns]
    assert len(col_names) == 2
    assert len(set(col_names)) == 2  # all unique
    df = result_set.to_pandas_df()
    assert df.iloc[0].tolist() == [10, 20]


def test_empty_column_names_do_not_rename_explicit_synthetic_names() -> None:
    """
    Synthetic names assigned to empty columns must not collide with explicit
    user-selected names that already look like Superset fallbacks.
    """
    data = [(10, 20)]
    description = [
        ("", 3, None, None, None, None, None),
        ("_col_0", 3, None, None, None, None, None),
    ]
    result_set = SupersetResultSet(data, description, BaseEngineSpec)  # type: ignore

    col_names = [c["column_name"] for c in result_set.columns]
    assert col_names == ["_col_1", "_col_0"]
    df = result_set.to_pandas_df()
    assert list(df.columns) == ["_col_1", "_col_0"]
    assert df.iloc[0].tolist() == [10, 20]


def test_stringify_values_dict_and_list_produce_valid_json() -> None:
    """
    ClickHouse native JSON and Map types return Python dicts. When stringified for
    Arrow array storage they must produce valid double-quoted JSON strings, not
    Python's single-quoted repr. Single-quoted strings pass the cheap '{' prefix
    check in the frontend's safeJsonObjectParse but then fail JSONbig.parse(),
    so the SQL Lab cell viewer never activates.
    """
    data = np.array(
        [
            {"key": "value", "nested": {"a": 1}},
            # str() gives ['a', 'b'] (single-quoted, invalid JSON);
            # json.dumps gives ["a", "b"] (double-quoted, valid JSON).
            ["a", "b"],
            {"items": [1, 2, 3], "d": "Hello, World!"},
            None,
        ],
        dtype=object,
    )
    result = stringify_values(data)

    # Must be valid JSON strings (double-quoted), not Python repr (single-quoted)
    assert result[0] == '{"key": "value", "nested": {"a": 1}}'
    assert result[1] == '["a", "b"]'
    assert result[2] == '{"items": [1, 2, 3], "d": "Hello, World!"}'
    assert result[3] is None

    # Parseable by a JSON parser — confirms the frontend's JSON.parse would succeed
    parsed = superset_json.loads(result[0])
    assert parsed == {"key": "value", "nested": {"a": 1}}
    parsed = superset_json.loads(result[1])
    assert parsed == ["a", "b"]


def test_clickhouse_json_column_in_pa_table_is_valid_json() -> None:
    """
    Verify that ClickHouse-style heterogeneous dict columns produce valid JSON
    strings in the Arrow table used by the msgpack serialization path.

    When clickhouse-connect returns Python dicts for JSON/Map type columns,
    SupersetResultSet must serialize them with json.dumps (not str()) so that
    the SQL Lab grid's cell viewer can call JSON.parse on the value.
    """
    data = [
        (1, {"a": {"b": 42}, "c": [1, 2, 3], "d": "Hello, World!"}),
        (2, {"e": 5}),
        (3, None),
    ]
    description = [
        ("id", 3, None, None, None, None, None),
        ("json_col", None, None, None, None, None, None),
    ]
    result_set = SupersetResultSet(data, description, BaseEngineSpec)  # type: ignore

    df_from_pa = SupersetResultSet.convert_table_to_df(result_set.pa_table)

    val0 = df_from_pa["json_col"].iloc[0]
    val1 = df_from_pa["json_col"].iloc[1]

    # Values in pa_table must be valid JSON strings (parseable by JSON.parse)
    assert isinstance(val0, str)
    assert isinstance(val1, str)

    # Double-quoted JSON, not single-quoted Python repr
    parsed0 = superset_json.loads(val0)
    assert parsed0 == {"a": {"b": 42}, "c": [1, 2, 3], "d": "Hello, World!"}
    parsed1 = superset_json.loads(val1)
    assert parsed1 == {"e": 5}


def test_stringify_values_non_serializable_dict_falls_back_to_str() -> None:
    """
    When a dict/list contains a value that json.dumps cannot serialize (e.g. bytes),
    stringify_values must fall back to str() rather than raising TypeError and crashing
    the result-set construction path.
    """

    class _Unserializable:
        def __repr__(self) -> str:
            return "unserializable"

    data = np.array(
        [{"key": _Unserializable()}],
        dtype=object,
    )
    # Must not raise — falls back to str()
    result = stringify_values(data)
    assert result[0] == str({"key": _Unserializable()})