mirror of
https://github.com/apache/superset.git
synced 2026-04-19 08:04:53 +00:00
feat: improve Doris catalog support (#34140)
This commit is contained in:
@@ -19,6 +19,7 @@ from typing import Any, Optional
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from sqlalchemy import JSON, types
|
||||
from sqlalchemy.engine.url import make_url
|
||||
|
||||
@@ -81,30 +82,62 @@ def test_get_column_spec(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sqlalchemy_uri,connect_args,return_schema,return_connect_args",
|
||||
"sqlalchemy_uri, connect_args, catalog, schema, return_schema,return_connect_args",
|
||||
[
|
||||
(
|
||||
"doris://user:password@host/db1",
|
||||
{"param1": "some_value"},
|
||||
"internal.information_schema",
|
||||
None,
|
||||
None,
|
||||
"db1",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
(
|
||||
"pydoris://user:password@host/db1",
|
||||
{"param1": "some_value"},
|
||||
"internal.information_schema",
|
||||
None,
|
||||
None,
|
||||
"db1",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
(
|
||||
"doris://user:password@host/catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
"catalog1.information_schema",
|
||||
None,
|
||||
None,
|
||||
"catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
(
|
||||
"pydoris://user:password@host/catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
"catalog1.information_schema",
|
||||
None,
|
||||
None,
|
||||
"catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
(
|
||||
"pydoris://user:password@host/catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
"catalog2",
|
||||
None,
|
||||
"catalog2.db1",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
(
|
||||
"pydoris://user:password@host/catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
None,
|
||||
"db2",
|
||||
"catalog1.db2",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
(
|
||||
"pydoris://user:password@host/catalog1.db1",
|
||||
{"param1": "some_value"},
|
||||
"catalog2",
|
||||
"db2",
|
||||
"catalog2.db2",
|
||||
{"param1": "some_value"},
|
||||
),
|
||||
],
|
||||
@@ -112,6 +145,8 @@ def test_get_column_spec(
|
||||
def test_adjust_engine_params(
|
||||
sqlalchemy_uri: str,
|
||||
connect_args: dict[str, Any],
|
||||
catalog: str | None,
|
||||
schema: str | None,
|
||||
return_schema: str,
|
||||
return_connect_args: dict[str, Any],
|
||||
) -> None:
|
||||
@@ -119,18 +154,36 @@ def test_adjust_engine_params(
|
||||
|
||||
url = make_url(sqlalchemy_uri)
|
||||
returned_url, returned_connect_args = DorisEngineSpec.adjust_engine_params(
|
||||
url, connect_args
|
||||
url,
|
||||
connect_args,
|
||||
catalog,
|
||||
schema,
|
||||
)
|
||||
|
||||
assert returned_url.database == return_schema
|
||||
assert returned_connect_args == return_connect_args
|
||||
|
||||
|
||||
def test_adjust_engine_params_no_database() -> None:
|
||||
"""
|
||||
Test that we raise an exception when the database is not specified.
|
||||
"""
|
||||
from superset.db_engine_specs.doris import DorisEngineSpec
|
||||
|
||||
url = make_url("doris://user:password@host")
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Doris requires a database to be specified in the URI.",
|
||||
):
|
||||
DorisEngineSpec.adjust_engine_params(url, {})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url,expected_schema",
|
||||
[
|
||||
("doris://localhost:9030/hive.test", "test"),
|
||||
("doris://localhost:9030/hive", None),
|
||||
("doris://localhost:9030/test", "test"),
|
||||
("doris://localhost:9030/", None),
|
||||
],
|
||||
)
|
||||
def test_get_schema_from_engine_params(
|
||||
@@ -154,12 +207,14 @@ def test_get_schema_from_engine_params(
|
||||
"database_value,expected_catalog",
|
||||
[
|
||||
("catalog1.schema1", "catalog1"),
|
||||
("catalog1", "catalog1"),
|
||||
(None, None),
|
||||
("schema1", "catalog2"),
|
||||
("", "catalog2"),
|
||||
],
|
||||
)
|
||||
def test_get_default_catalog(
|
||||
database_value: Optional[str], expected_catalog: Optional[str]
|
||||
mocker: MockerFixture,
|
||||
database_value: Optional[str],
|
||||
expected_catalog: Optional[str],
|
||||
) -> None:
|
||||
"""
|
||||
Test the ``get_default_catalog`` method.
|
||||
@@ -167,8 +222,14 @@ def test_get_default_catalog(
|
||||
from superset.db_engine_specs.doris import DorisEngineSpec
|
||||
from superset.models.core import Database
|
||||
|
||||
database = Mock(spec=Database)
|
||||
database = mocker.MagicMock(spec=Database)
|
||||
database.url_object.database = database_value
|
||||
rows = [
|
||||
mocker.MagicMock(IsCurrent=False, CatalogName="catalog1"),
|
||||
mocker.MagicMock(IsCurrent=True, CatalogName="catalog2"),
|
||||
]
|
||||
with database.get_sqla_engine() as engine:
|
||||
engine.execute.return_value = rows
|
||||
|
||||
assert DorisEngineSpec.get_default_catalog(database) == expected_catalog
|
||||
|
||||
|
||||
Reference in New Issue
Block a user