mirror of
https://github.com/apache/superset.git
synced 2026-09-01 21:11:28 +00:00
253 lines
8.8 KiB
Python
253 lines
8.8 KiB
Python
# 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.
|
|
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
from pytest_mock import MockerFixture
|
|
|
|
from superset.utils import json
|
|
from superset.utils.core import GenericDataType
|
|
from tests.conftest import with_config
|
|
from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
|
|
from tests.unit_tests.fixtures.common import dttm # noqa: F401
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"target_type,expected_result",
|
|
[
|
|
("Text", "'2019-01-02 03:04:05.678900'"),
|
|
("DateTime", "'2019-01-02 03:04:05.678900'"),
|
|
("UnknownType", None),
|
|
],
|
|
)
|
|
def test_convert_dttm(
|
|
target_type: str,
|
|
expected_result: Optional[str],
|
|
dttm: datetime, # noqa: F811
|
|
) -> None:
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec as spec # noqa: N813
|
|
|
|
assert_convert_dttm(spec, target_type, expected_result, dttm)
|
|
|
|
|
|
@with_config({"VERSION_STRING": "1.0.0"})
|
|
def test_get_extra_params(mocker: MockerFixture) -> None:
|
|
"""
|
|
Test the ``get_extra_params`` method.
|
|
"""
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
|
|
|
database = mocker.MagicMock()
|
|
|
|
database.extra = {}
|
|
assert DuckDBEngineSpec.get_extra_params(database) == {
|
|
"engine_params": {
|
|
"connect_args": {"config": {"custom_user_agent": "apache-superset/1.0.0"}}
|
|
}
|
|
}
|
|
|
|
database.extra = json.dumps(
|
|
{"engine_params": {"connect_args": {"config": {"custom_user_agent": "my-app"}}}}
|
|
)
|
|
assert DuckDBEngineSpec.get_extra_params(database) == {
|
|
"engine_params": {
|
|
"connect_args": {
|
|
"config": {"custom_user_agent": "apache-superset/1.0.0 my-app"}
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
def test_build_sqlalchemy_uri() -> None:
|
|
"""Test DuckDBEngineSpec.build_sqlalchemy_uri"""
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec, DuckDBParametersType
|
|
|
|
# No database provided, default to :memory:
|
|
parameters = DuckDBParametersType()
|
|
uri = DuckDBEngineSpec.build_sqlalchemy_uri(parameters)
|
|
assert "duckdb:///:memory:" == uri
|
|
|
|
# Database provided
|
|
parameters = DuckDBParametersType(database="/path/to/duck.db")
|
|
uri = DuckDBEngineSpec.build_sqlalchemy_uri(parameters)
|
|
assert "duckdb:////path/to/duck.db" == uri
|
|
|
|
|
|
def test_md_build_sqlalchemy_uri() -> None:
|
|
"""Test MotherDuckEngineSpec.build_sqlalchemy_uri"""
|
|
from superset.db_engine_specs.duckdb import (
|
|
DuckDBParametersType,
|
|
MotherDuckEngineSpec,
|
|
)
|
|
|
|
# No access token provided, throw ValueError
|
|
parameters = DuckDBParametersType(database="my_db")
|
|
with pytest.raises(ValueError): # noqa: PT011
|
|
MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
|
|
|
# No database provided, default to "md:"
|
|
parameters = DuckDBParametersType(access_token="token") # noqa: S106
|
|
uri = MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
|
assert "duckdb:///md:?motherduck_token=token"
|
|
|
|
# Database and access_token provided
|
|
parameters = DuckDBParametersType(database="my_db", access_token="token") # noqa: S106
|
|
uri = MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
|
assert "duckdb:///md:my_db?motherduck_token=token" == uri
|
|
|
|
|
|
def test_get_parameters_from_uri() -> None:
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
|
|
|
uri = "duckdb:////path/to/duck.db"
|
|
parameters = DuckDBEngineSpec.get_parameters_from_uri(uri)
|
|
|
|
assert parameters["database"] == "/path/to/duck.db"
|
|
|
|
uri = "duckdb:///md:my_db?motherduck_token=token"
|
|
parameters = DuckDBEngineSpec.get_parameters_from_uri(uri)
|
|
|
|
assert parameters["database"] == "md:my_db"
|
|
assert parameters["access_token"] == "token" # noqa: S105
|
|
|
|
|
|
def test_column_type_recognition() -> None:
|
|
"""Test that DuckDB column types are properly recognized as numeric."""
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
|
|
|
# Test standard float/double types
|
|
numeric_types = [
|
|
"FLOAT",
|
|
"DOUBLE",
|
|
"DOUBLE PRECISION",
|
|
"REAL",
|
|
"DECIMAL(10,2)",
|
|
"NUMERIC(10,2)",
|
|
"INTEGER",
|
|
"BIGINT",
|
|
"SMALLINT",
|
|
# DuckDB-specific unsigned types
|
|
"HUGEINT",
|
|
"UBIGINT",
|
|
"UINTEGER",
|
|
"USMALLINT",
|
|
"UTINYINT",
|
|
]
|
|
|
|
for type_str in numeric_types:
|
|
col_spec = DuckDBEngineSpec.get_column_spec(type_str)
|
|
assert col_spec is not None, f"Type {type_str} should be recognized"
|
|
assert col_spec.generic_type == GenericDataType.NUMERIC, (
|
|
f"Type {type_str} should be recognized as NUMERIC, "
|
|
f"got {col_spec.generic_type}"
|
|
)
|
|
|
|
# Test that TINYINT (non-unsigned) is also recognized
|
|
# Note: TINYINT is not in the default mappings, but should be handled
|
|
col_spec = DuckDBEngineSpec.get_column_spec("TINYINT")
|
|
# TINYINT matches the pattern "^int" so it should be recognized
|
|
assert col_spec is None, "TINYINT doesn't match any patterns"
|
|
|
|
|
|
def test_fetch_data_preserves_cursor_description(mocker: MockerFixture) -> None:
|
|
"""
|
|
DuckDBEngineSpec previously overrode fetch_data to capture and restore
|
|
cursor.description around fetchall(), working around a duckdb-engine bug
|
|
that no longer reproduces at current pinned versions (duckdb-engine
|
|
0.17.0, duckdb 1.5.5). Confirm the inherited base fetch_data still works
|
|
correctly against a real cursor now that the override is gone.
|
|
"""
|
|
from sqlalchemy import create_engine
|
|
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
|
|
|
engine = create_engine("duckdb:///:memory:")
|
|
raw_conn = engine.raw_connection()
|
|
try:
|
|
cursor = raw_conn.cursor()
|
|
database = mocker.MagicMock()
|
|
DuckDBEngineSpec.execute(cursor, "SELECT 1 AS col1, 2 AS col2", database)
|
|
|
|
data = DuckDBEngineSpec.fetch_data(cursor)
|
|
|
|
assert data == [(1, 2)]
|
|
assert cursor.description is not None
|
|
assert [col[0] for col in cursor.description] == ["col1", "col2"]
|
|
finally:
|
|
raw_conn.close()
|
|
|
|
|
|
def test_extended_aggregation_func_median_stddev_var_compiles() -> None:
|
|
"""
|
|
MEDIAN/STDDEV_SAMP/VAR_SAMP compile to the expected DuckDB SQL function
|
|
calls. See `test_extended_aggregation_func_median_stddev_var_executes`
|
|
for verification against a live in-process DuckDB instance.
|
|
"""
|
|
from sqlalchemy import column
|
|
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
|
|
|
col = column("sales")
|
|
|
|
for aggregate, expected_sql in [
|
|
("MEDIAN", "median(sales)"),
|
|
("STDDEV_SAMP", "stddev_samp(sales)"),
|
|
("VAR_SAMP", "var_samp(sales)"),
|
|
]:
|
|
func = DuckDBEngineSpec.get_extended_aggregation_func(aggregate)
|
|
assert func is not None
|
|
compiled = str(func(col).compile(compile_kwargs={"literal_binds": True}))
|
|
assert compiled == expected_sql
|
|
|
|
|
|
def test_extended_aggregation_func_median_stddev_var_executes() -> None:
|
|
"""
|
|
MEDIAN/STDDEV_SAMP/VAR_SAMP execute against a live in-process DuckDB
|
|
instance and return values matching Python's `statistics` module
|
|
(sample standard deviation/variance) for the same input.
|
|
"""
|
|
import statistics
|
|
|
|
from sqlalchemy import create_engine, literal_column, select, text
|
|
|
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
|
|
|
values = [1.0, 2.0, 4.0, 8.0, 16.0]
|
|
|
|
engine = create_engine("duckdb:///:memory:")
|
|
with engine.connect() as conn:
|
|
conn.execute(text("CREATE TABLE t (sales DOUBLE)"))
|
|
conn.execute(
|
|
text("INSERT INTO t VALUES (:sales)"),
|
|
[{"sales": v} for v in values],
|
|
)
|
|
|
|
expected = {
|
|
"MEDIAN": statistics.median(values),
|
|
"STDDEV_SAMP": statistics.stdev(values),
|
|
"VAR_SAMP": statistics.variance(values),
|
|
}
|
|
|
|
for aggregate, expected_value in expected.items():
|
|
func = DuckDBEngineSpec.get_extended_aggregation_func(aggregate)
|
|
assert func is not None
|
|
query = select(func(literal_column("sales"))).select_from(text("t"))
|
|
result = conn.execute(query).scalar()
|
|
assert result == pytest.approx(expected_value)
|