feat(doris): add catalog support for Apache Doris (#31580)

Co-authored-by: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Co-authored-by: Beto Dealmeida <roberto@dealmeida.net>
This commit is contained in:
Jiwen liu
2025-01-07 02:21:00 +08:00
committed by GitHub
parent a986a61b5f
commit 72e87d6980
2 changed files with 122 additions and 22 deletions

View File

@@ -16,6 +16,7 @@
# under the License.
from typing import Any, Optional
from unittest.mock import Mock
import pytest
from sqlalchemy import JSON, types
@@ -85,25 +86,25 @@ def test_get_column_spec(
(
"doris://user:password@host/db1",
{"param1": "some_value"},
"db1",
"internal.information_schema",
{"param1": "some_value"},
),
(
"pydoris://user:password@host/db1",
{"param1": "some_value"},
"db1",
"internal.information_schema",
{"param1": "some_value"},
),
(
"doris://user:password@host/catalog1.db1",
{"param1": "some_value"},
"catalog1.db1",
"catalog1.information_schema",
{"param1": "some_value"},
),
(
"pydoris://user:password@host/catalog1.db1",
{"param1": "some_value"},
"catalog1.db1",
"catalog1.information_schema",
{"param1": "some_value"},
),
],
@@ -120,11 +121,21 @@ def test_adjust_engine_params(
returned_url, returned_connect_args = DorisEngineSpec.adjust_engine_params(
url, connect_args
)
assert returned_url.database == return_schema
assert returned_connect_args == return_connect_args
def test_get_schema_from_engine_params() -> None:
@pytest.mark.parametrize(
"url,expected_schema",
[
("doris://localhost:9030/hive.test", "test"),
("doris://localhost:9030/hive", None),
],
)
def test_get_schema_from_engine_params(
url: str, expected_schema: Optional[str]
) -> None:
"""
Test the ``get_schema_from_engine_params`` method.
"""
@@ -132,16 +143,74 @@ def test_get_schema_from_engine_params() -> None:
assert (
DorisEngineSpec.get_schema_from_engine_params(
make_url("doris://localhost:9030/hive.test"),
make_url(url),
{},
)
== "test"
== expected_schema
)
assert (
DorisEngineSpec.get_schema_from_engine_params(
make_url("doris://localhost:9030/hive"),
{},
)
is None
)
@pytest.mark.parametrize(
"database_value,expected_catalog",
[
("catalog1.schema1", "catalog1"),
("catalog1", "catalog1"),
(None, None),
],
)
def test_get_default_catalog(
database_value: Optional[str], expected_catalog: Optional[str]
) -> None:
"""
Test the ``get_default_catalog`` method.
"""
from superset.db_engine_specs.doris import DorisEngineSpec
from superset.models.core import Database
database = Mock(spec=Database)
database.url_object.database = database_value
assert DorisEngineSpec.get_default_catalog(database) == expected_catalog
@pytest.mark.parametrize(
"mock_catalogs,expected_result",
[
(
[
Mock(CatalogName="catalog1"),
Mock(CatalogName="catalog2"),
Mock(CatalogName="catalog3"),
],
{"catalog1", "catalog2", "catalog3"},
),
(
[Mock(CatalogName="single_catalog")],
{"single_catalog"},
),
(
[],
set(),
),
],
)
def test_get_catalog_names(
mock_catalogs: list[Mock], expected_result: set[str]
) -> None:
"""
Test the ``get_catalog_names`` method.
"""
from superset.db_engine_specs.doris import DorisEngineSpec
from superset.models.core import Database
database = Mock(spec=Database)
inspector = Mock()
inspector.bind.execute.return_value = mock_catalogs
catalogs = DorisEngineSpec.get_catalog_names(database, inspector)
# Verify the SQL query
inspector.bind.execute.assert_called_once_with("SHOW CATALOGS")
# Verify the returned catalog names
assert catalogs == expected_result