This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 8c3b985c22c fix(db2): set current_schema to the catalog name of the
selected schema (#44706)
8c3b985c22c is described below
commit 8c3b985c22c34b9e9d78629666034f9946efacbb
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Thu Oct 1 00:02:15 2026 +1000
fix(db2): set current_schema to the catalog name of the selected schema
(#44706)
---
superset/db_engine_specs/db2.py | 12 ++++++-
tests/unit_tests/db_engine_specs/test_db2.py | 52 ++++++++++++++++++++++++++--
2 files changed, 60 insertions(+), 4 deletions(-)
diff --git a/superset/db_engine_specs/db2.py b/superset/db_engine_specs/db2.py
index dc66dc31923..393d863962b 100644
--- a/superset/db_engine_specs/db2.py
+++ b/superset/db_engine_specs/db2.py
@@ -150,8 +150,18 @@ class Db2EngineSpec(BaseEngineSpec):
any tables with unqualified names. If the schema is not set by SQL Lab
it could
be anything, and we would have to block users from running any queries
referencing tables without an explicit schema.
+
+ The schema name is denormalized like reflection does, so a schema
created
+ quoted and lower case (``CREATE SCHEMA "lowonly"``) resolves to its
+ upper-case name (``LOWONLY``), matching ``get_table_names``. Such
schemas
+ are not reachable through unqualified names in SQL Lab.
"""
if not schema:
return []
- escaped = schema.replace('"', '""')
+ # Schema names come from the inspector, where unquoted (upper-case) DB2
+ # names are normalized to lower case. Quoting that name as-is would
+ # select a different, usually non-existent, schema, so convert it back
+ # to the name stored in the catalog first.
+ name = cls.denormalize_name(database.get_dialect(), schema)
+ escaped = name.replace('"', '""')
return [f'set current_schema "{escaped}"']
diff --git a/tests/unit_tests/db_engine_specs/test_db2.py
b/tests/unit_tests/db_engine_specs/test_db2.py
index ad87c4fc737..9a19638dec3 100644
--- a/tests/unit_tests/db_engine_specs/test_db2.py
+++ b/tests/unit_tests/db_engine_specs/test_db2.py
@@ -15,8 +15,12 @@
# specific language governing permissions and limitations
# under the License.
-import pytest # noqa: F401
+from typing import Callable
+
+import pytest
from pytest_mock import MockerFixture
+from sqlalchemy.engine import Dialect, make_url
+from sqlalchemy.engine.default import DefaultDialect
from sqlglot import parse_one
from sqlglot.errors import ParseError
@@ -88,23 +92,65 @@ def test_get_table_comment_unexpected_error(mocker:
MockerFixture):
)
-def test_get_prequeries(mocker: MockerFixture) -> None:
+def _ibm_db_sa_dialect() -> Dialect:
+ """The real DB2 dialect; building it needs ``ibm_db_sa`` but no server."""
+ pytest.importorskip("ibm_db_sa")
+ return make_url("db2+ibm_db://u:p@h/d").get_dialect()()
+
+
+def _sqlalchemy_normalizing_dialect() -> Dialect:
+ """
+ SQLAlchemy's own name normalization, which ``ibm_db_sa`` matches. CI does
+ not install ``ibm_db_sa`` (no Linux arm64 wheel), so this keeps the
+ assertions running there.
+ """
+ dialect = DefaultDialect()
+ dialect.requires_name_normalize = True
+ return dialect
+
+
[email protected](
+ "dialect_factory",
+ [_ibm_db_sa_dialect, _sqlalchemy_normalizing_dialect],
+ ids=["ibm_db_sa", "sqlalchemy"],
+)
+def test_get_prequeries(
+ mocker: MockerFixture, dialect_factory: Callable[[], Dialect]
+) -> None:
"""
Test the ``get_prequeries`` method.
"""
from superset.db_engine_specs.db2 import Db2EngineSpec
database = mocker.MagicMock()
+ database.get_dialect.return_value = dialect_factory()
assert Db2EngineSpec.get_prequeries(database) == []
assert Db2EngineSpec.get_prequeries(database, schema="my_schema") == [
- 'set current_schema "my_schema"'
+ 'set current_schema "MY_SCHEMA"'
+ ]
+ assert Db2EngineSpec.get_prequeries(database, schema="MixedCase") == [
+ 'set current_schema "MixedCase"'
]
assert Db2EngineSpec.get_prequeries(database, schema='evil"; SELECT 1--')
== [
'set current_schema "evil""; SELECT 1--"'
]
+def test_get_prequeries_without_name_normalization(mocker: MockerFixture) ->
None:
+ """Preserve schema names when the dialect does not request
normalization."""
+ from superset.db_engine_specs.db2 import Db2EngineSpec
+
+ database = mocker.MagicMock()
+ dialect = database.get_dialect.return_value
+ dialect.requires_name_normalize = False
+
+ assert Db2EngineSpec.get_prequeries(database, schema="my_schema") == [
+ 'set current_schema "my_schema"'
+ ]
+ dialect.denormalize_name.assert_not_called()
+
+
@pytest.mark.parametrize(
("grain", "expected_expression"),
[
