QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-19 16:14:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: make user agent customizable (#32506)

Browse Source
This commit is contained in:
Ville Brofeldt
2025-03-05 16:33:24 -08:00
committed by GitHub
parent 8decc9e45f
commit c90e45a373
18 changed files with 101 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/unit_tests/db_engine_specs/test_trino.py
View File

@@ -37,7 +37,7 @@ from trino.sqlalchemy import datatype
from trino.sqlalchemy.dialect import TrinoDialect
import superset.config
from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY, USER_AGENT
from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY
from superset.db_engine_specs.exceptions import (
SupersetDBAPIConnectionError,
SupersetDBAPIDatabaseError,
@@ -81,7 +81,7 @@ def _assert_columns_equal(actual_cols, expected_cols) -> None:
@pytest.mark.parametrize(
"extra,expected",
[
({}, {"engine_params": {"connect_args": {"source": USER_AGENT}}}),
({}, {"engine_params": {"connect_args": {"source": "Apache Superset"}}}),
(
{
"first": 1,
@@ -110,7 +110,7 @@ def test_get_extra_params(extra: dict[str, Any], expected: dict[str, Any]) -> No
assert TrinoEngineSpec.get_extra_params(database) == expected
@patch("superset.utils.core.create_ssl_cert_file")
@patch("superset.db_engine_specs.trino.create_ssl_cert_file")
def test_get_extra_params_with_server_cert(mock_create_ssl_cert_file: Mock) -> None:
from superset.db_engine_specs.trino import TrinoEngineSpec
@@ -118,6 +118,7 @@ def test_get_extra_params_with_server_cert(mock_create_ssl_cert_file: Mock) -> N
database.extra = json.dumps({})
database.server_cert = "TEST_CERT"
database.db_engine_spec = TrinoEngineSpec
mock_create_ssl_cert_file.return_value = "/path/to/tls.crt"
extra = TrinoEngineSpec.get_extra_params(database)

22
tests/unit_tests/utils/test_core.py
View File

@@ -21,6 +21,7 @@ from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
from pytest_mock import MockerFixture
from superset.exceptions import SupersetException
from superset.utils.core import (
@@ -30,10 +31,12 @@ from superset.utils.core import (
generic_find_constraint_name,
generic_find_fk_constraint_name,
get_datasource_full_name,
get_user_agent,
is_test,
normalize_dttm_col,
parse_boolean_string,
QueryObjectFilterClause,
QuerySource,
remove_extra_adhoc_filters,
)
@@ -396,3 +399,22 @@ def test_get_datasource_full_name():
get_datasource_full_name("db", "table", "catalog", None)
== "[db].[catalog].[table]"
)
def test_get_user_agent(mocker: MockerFixture) -> None:
database_mock = mocker.MagicMock()
database_mock.database_name = "mydb"
current_app_mock = mocker.patch("superset.utils.core.current_app")
current_app_mock.config = {"USER_AGENT_FUNC": None}
assert get_user_agent(database_mock, QuerySource.DASHBOARD) == "Apache Superset", (
"The default user agent should be returned"
)
current_app_mock.config["USER_AGENT_FUNC"] = (
lambda database, source: f"{database.database_name} {source.name}"
)
assert get_user_agent(database_mock, QuerySource.DASHBOARD) == "mydb DASHBOARD", (
"the custom user agent function result should have been returned"
)