fix(db2): Improving support for ibm db2 connections (#26744)

setup.py

@@ -152,7 +152,7 @@ def get_git_sha() -> str:
             "databricks-sql-connector>=2.0.2, <3",
             "sqlalchemy-databricks>=0.2.0",
         ],
-        "db2": ["ibm-db-sa>=0.3.5, <0.4"],
+        "db2": ["ibm-db-sa>0.3.8, <=0.4.0"],
         "dremio": ["sqlalchemy-dremio>=1.1.5, <1.3"],
         "drill": ["sqlalchemy-drill==0.1.dev"],
         "druid": ["pydruid>=0.6.5,<0.7"],

superset/db_engine_specs/db2.py

@@ -14,9 +14,16 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import logging
+from typing import Optional, Union
+
+from sqlalchemy.engine.reflection import Inspector
+
 from superset.constants import TimeGrain
 from superset.db_engine_specs.base import BaseEngineSpec, LimitMethod
 
+logger = logging.getLogger(__name__)
+
 
 class Db2EngineSpec(BaseEngineSpec):
     engine = "db2"
@@ -26,6 +33,8 @@ class Db2EngineSpec(BaseEngineSpec):
     force_column_alias_quotes = True
     max_column_name_length = 30
 
+    supports_dynamic_schema = True
+
     _time_grain_expressions = {
         None: "{col}",
         TimeGrain.SECOND: "CAST({col} as TIMESTAMP) - MICROSECOND({col}) MICROSECONDS",
@@ -52,3 +61,49 @@ class Db2EngineSpec(BaseEngineSpec):
     @classmethod
     def epoch_to_dttm(cls) -> str:
         return "(TIMESTAMP('1970-01-01', '00:00:00') + {col} SECONDS)"
+
+    @classmethod
+    def get_table_comment(
+        cls, inspector: Inspector, table_name: str, schema: Union[str, None]
+    ) -> Optional[str]:
+        """
+        Get comment of table from a given schema
+
+        Ibm Db2 return comments as tuples, so we need to get the first element
+
+        :param inspector: SqlAlchemy Inspector instance
+        :param table_name: Table name
+        :param schema: Schema name. If omitted, uses default schema for database
+        :return: comment of table
+        """
+        comment = None
+        try:
+            table_comment = inspector.get_table_comment(table_name, schema)
+            comment = table_comment.get("text")
+            return comment[0]
+        except IndexError:
+            return comment
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.error("Unexpected error while fetching table comment", exc_info=True)
+            logger.exception(ex)
+            return comment
+
+    @classmethod
+    def get_prequeries(
+        cls,
+        catalog: Union[str, None] = None,
+        schema: Union[str, None] = None,
+    ) -> list[str]:
+        """
+        Set the search path to the specified schema.
+
+        This is important for two reasons: in SQL Lab it will allow queries to run in
+        the schema selected in the dropdown, resolving unqualified table names to the
+        expected schema.
+
+        But more importantly, in SQL Lab this is used to check if the user has access to
+        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.
+        """
+        return [f'set current_schema "{schema}"'] if schema else []

tests/unit_tests/db_engine_specs/test_db2.py

@@ -0,0 +1,75 @@
+# 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.
+
+import pytest
+from pytest_mock import MockerFixture
+
+
+def test_epoch_to_dttm() -> None:
+    """
+    Test the `epoch_to_dttm` method.
+    """
+    from superset.db_engine_specs.db2 import Db2EngineSpec
+
+    assert (
+        Db2EngineSpec.epoch_to_dttm().format(col="epoch_dttm")
+        == "(TIMESTAMP('1970-01-01', '00:00:00') + epoch_dttm SECONDS)"
+    )
+
+
+def test_get_table_comment(mocker: MockerFixture):
+    """
+    Test the `get_table_comment` method.
+    """
+    from superset.db_engine_specs.db2 import Db2EngineSpec
+
+    mock_inspector = mocker.MagicMock()
+    mock_inspector.get_table_comment.return_value = {
+        "text": ("This is a table comment",)
+    }
+
+    assert (
+        Db2EngineSpec.get_table_comment(mock_inspector, "my_table", "my_schema")
+        == "This is a table comment"
+    )
+
+
+def test_get_table_comment_empty(mocker: MockerFixture):
+    """
+    Test the `get_table_comment` method
+    when no comment is returned.
+    """
+    from superset.db_engine_specs.db2 import Db2EngineSpec
+
+    mock_inspector = mocker.MagicMock()
+    mock_inspector.get_table_comment.return_value = {}
+
+    assert (
+        Db2EngineSpec.get_table_comment(mock_inspector, "my_table", "my_schema") == None
+    )
+
+
+def test_get_prequeries() -> None:
+    """
+    Test the ``get_prequeries`` method.
+    """
+    from superset.db_engine_specs.db2 import Db2EngineSpec
+
+    assert Db2EngineSpec.get_prequeries() == []
+    assert Db2EngineSpec.get_prequeries(schema="my_schema") == [
+        'set current_schema "my_schema"'
+    ]