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

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 []