diff --git a/setup.py b/setup.py index fd2a9f8c80cd6..b4306a1768644 100644 --- a/setup.py +++ b/setup.py @@ -151,7 +151,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"], diff --git a/superset/db_engine_specs/db2.py b/superset/db_engine_specs/db2.py index 5f54613a4b533..db2e500b53d8f 100644 --- a/superset/db_engine_specs/db2.py +++ b/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 [] diff --git a/tests/unit_tests/db_engine_specs/test_db2.py b/tests/unit_tests/db_engine_specs/test_db2.py new file mode 100644 index 0000000000000..d7dd19ad5923c --- /dev/null +++ b/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"' + ]