From e1db016a6c324c64eec0b3ad2c1c0c8d9b85d813 Mon Sep 17 00:00:00 2001
From: Karol Kostrzewa <karolkostrzew@gmail.com>
Date: Thu, 21 Jan 2021 21:53:54 +0100
Subject: [PATCH] test: presto engine spec tests (#12594)

* test get_table_names

* test _get_full_name

* add test_split_data_type

* test _show_columns

* add test_is_column_name_quoted

* test select_star

* test get_view_names

* test estimate_statement_cost

* test get_all_datasource_names

* test get_create_view

* test _extract_error_message

* fix typo
---
 tests/db_engine_specs/presto_tests.py | 300 ++++++++++++++++++++++++++
 1 file changed, 300 insertions(+)

diff --git a/tests/db_engine_specs/presto_tests.py b/tests/db_engine_specs/presto_tests.py
index 9a493d3776..9b2024dd72 100644
--- a/tests/db_engine_specs/presto_tests.py
+++ b/tests/db_engine_specs/presto_tests.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from collections import namedtuple
 from unittest import mock, skipUnless
 
 import pandas as pd
@@ -23,6 +24,7 @@ from sqlalchemy.sql import select
 
 from superset.db_engine_specs.presto import PrestoEngineSpec
 from superset.sql_parse import ParsedQuery
+from superset.utils.core import DatasourceName
 from tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 
@@ -38,6 +40,45 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
             [], PrestoEngineSpec.get_view_names(mock.ANY, mock.ANY, mock.ANY)
         )
 
+    @mock.patch("superset.db_engine_specs.presto.is_feature_enabled")
+    def test_get_view_names(self, mock_is_feature_enabled):
+        mock_is_feature_enabled.return_value = True
+        mock_execute = mock.MagicMock()
+        mock_fetchall = mock.MagicMock(return_value=[["a", "b,", "c"], ["d", "e"]])
+        database = mock.MagicMock()
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
+            mock_execute
+        )
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.fetchall = (
+            mock_fetchall
+        )
+        result = PrestoEngineSpec.get_view_names(database, mock.Mock(), None)
+        mock_execute.assert_called_once_with(
+            "SELECT table_name FROM information_schema.views", {}
+        )
+        assert result == ["a", "d"]
+
+    @mock.patch("superset.db_engine_specs.presto.is_feature_enabled")
+    def test_get_view_names_with_schema(self, mock_is_feature_enabled):
+        mock_is_feature_enabled.return_value = True
+        mock_execute = mock.MagicMock()
+        mock_fetchall = mock.MagicMock(return_value=[["a", "b,", "c"], ["d", "e"]])
+        database = mock.MagicMock()
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
+            mock_execute
+        )
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.fetchall = (
+            mock_fetchall
+        )
+        schema = "schema"
+        result = PrestoEngineSpec.get_view_names(database, mock.Mock(), schema)
+        mock_execute.assert_called_once_with(
+            "SELECT table_name FROM information_schema.views "
+            "WHERE table_schema=%(schema)s",
+            {"schema": schema},
+        )
+        assert result == ["a", "d"]
+
     def verify_presto_column(self, column, expected_results):
         inspector = mock.Mock()
         inspector.engine.dialect.identifier_preparer.quote_identifier = mock.Mock()
@@ -516,6 +557,265 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
         sqla_type = PrestoEngineSpec.get_sqla_column_type(None)
         assert sqla_type is None
 
+    @mock.patch(
+        "superset.utils.feature_flag_manager.FeatureFlagManager.is_feature_enabled"
+    )
+    @mock.patch("superset.db_engine_specs.base.BaseEngineSpec.get_table_names")
+    @mock.patch("superset.db_engine_specs.presto.PrestoEngineSpec.get_view_names")
+    def test_get_table_names_no_split_views_from_tables(
+        self, mock_get_view_names, mock_get_table_names, mock_is_feature_enabled
+    ):
+        mock_get_view_names.return_value = ["view1", "view2"]
+        table_names = ["table1", "table2", "view1", "view2"]
+        mock_get_table_names.return_value = table_names
+        mock_is_feature_enabled.return_value = False
+        tables = PrestoEngineSpec.get_table_names(mock.Mock(), mock.Mock(), None)
+        assert tables == table_names
+
+    @mock.patch(
+        "superset.utils.feature_flag_manager.FeatureFlagManager.is_feature_enabled"
+    )
+    @mock.patch("superset.db_engine_specs.base.BaseEngineSpec.get_table_names")
+    @mock.patch("superset.db_engine_specs.presto.PrestoEngineSpec.get_view_names")
+    def test_get_table_names_split_views_from_tables(
+        self, mock_get_view_names, mock_get_table_names, mock_is_feature_enabled
+    ):
+        mock_get_view_names.return_value = ["view1", "view2"]
+        table_names = ["table1", "table2", "view1", "view2"]
+        mock_get_table_names.return_value = table_names
+        mock_is_feature_enabled.return_value = True
+        tables = PrestoEngineSpec.get_table_names(mock.Mock(), mock.Mock(), None)
+        assert sorted(tables) == sorted(table_names)
+
+    @mock.patch(
+        "superset.utils.feature_flag_manager.FeatureFlagManager.is_feature_enabled"
+    )
+    @mock.patch("superset.db_engine_specs.base.BaseEngineSpec.get_table_names")
+    @mock.patch("superset.db_engine_specs.presto.PrestoEngineSpec.get_view_names")
+    def test_get_table_names_split_views_from_tables_no_tables(
+        self, mock_get_view_names, mock_get_table_names, mock_is_feature_enabled
+    ):
+        mock_get_view_names.return_value = []
+        table_names = []
+        mock_get_table_names.return_value = table_names
+        mock_is_feature_enabled.return_value = True
+        tables = PrestoEngineSpec.get_table_names(mock.Mock(), mock.Mock(), None)
+        assert tables == []
+
+    def test_get_full_name(self):
+        names = [
+            ("part1", "part2"),
+            ("part11", "part22"),
+        ]
+        result = PrestoEngineSpec._get_full_name(names)
+        assert result == "part1.part11"
+
+    def test_get_full_name_empty_tuple(self):
+        names = [
+            ("part1", "part2"),
+            ("", "part3"),
+            ("part4", "part5"),
+            ("", "part6"),
+        ]
+        result = PrestoEngineSpec._get_full_name(names)
+        assert result == "part1.part4"
+
+    def test_split_data_type(self):
+        data_type = "value1 value2"
+        result = PrestoEngineSpec._split_data_type(data_type, " ")
+        assert result == ["value1", "value2"]
+
+        data_type = "value1,value2"
+        result = PrestoEngineSpec._split_data_type(data_type, ",")
+        assert result == ["value1", "value2"]
+
+        data_type = '"value,1",value2'
+        result = PrestoEngineSpec._split_data_type(data_type, ",")
+        assert result == ['"value,1"', "value2"]
+
+    def test_show_columns(self):
+        inspector = mock.MagicMock()
+        inspector.engine.dialect.identifier_preparer.quote_identifier = (
+            lambda x: f'"{x}"'
+        )
+        mock_execute = mock.MagicMock(return_value=["a", "b"])
+        inspector.bind.execute = mock_execute
+        table_name = "table_name"
+        result = PrestoEngineSpec._show_columns(inspector, table_name, None)
+        assert result == ["a", "b"]
+        mock_execute.assert_called_once_with(f'SHOW COLUMNS FROM "{table_name}"')
+
+    def test_show_columns_with_schema(self):
+        inspector = mock.MagicMock()
+        inspector.engine.dialect.identifier_preparer.quote_identifier = (
+            lambda x: f'"{x}"'
+        )
+        mock_execute = mock.MagicMock(return_value=["a", "b"])
+        inspector.bind.execute = mock_execute
+        table_name = "table_name"
+        schema = "schema"
+        result = PrestoEngineSpec._show_columns(inspector, table_name, schema)
+        assert result == ["a", "b"]
+        mock_execute.assert_called_once_with(
+            f'SHOW COLUMNS FROM "{schema}"."{table_name}"'
+        )
+
+    def test_is_column_name_quoted(self):
+        column_name = "mock"
+        assert PrestoEngineSpec._is_column_name_quoted(column_name) is False
+
+        column_name = '"mock'
+        assert PrestoEngineSpec._is_column_name_quoted(column_name) is False
+
+        column_name = '"moc"k'
+        assert PrestoEngineSpec._is_column_name_quoted(column_name) is False
+
+        column_name = '"moc"k"'
+        assert PrestoEngineSpec._is_column_name_quoted(column_name) is True
+
+    @mock.patch("superset.db_engine_specs.base.BaseEngineSpec.select_star")
+    def test_select_star_no_presto_expand_data(self, mock_select_star):
+        database = mock.Mock()
+        table_name = "table_name"
+        engine = mock.Mock()
+        cols = [
+            {"col1": "val1"},
+            {"col2": "val2"},
+        ]
+        PrestoEngineSpec.select_star(database, table_name, engine, cols=cols)
+        mock_select_star.assert_called_once_with(
+            database, table_name, engine, None, 100, False, True, True, cols
+        )
+
+    @mock.patch("superset.db_engine_specs.presto.is_feature_enabled")
+    @mock.patch("superset.db_engine_specs.base.BaseEngineSpec.select_star")
+    def test_select_star_presto_expand_data(
+        self, mock_select_star, mock_is_feature_enabled
+    ):
+        mock_is_feature_enabled.return_value = True
+        database = mock.Mock()
+        table_name = "table_name"
+        engine = mock.Mock()
+        cols = [
+            {"name": "val1"},
+            {"name": "val2<?!@#$312,/'][p098"},
+            {"name": ".val2"},
+            {"name": "val2."},
+            {"name": "val.2"},
+            {"name": ".val2."},
+        ]
+        PrestoEngineSpec.select_star(
+            database, table_name, engine, show_cols=True, cols=cols
+        )
+        mock_select_star.assert_called_once_with(
+            database,
+            table_name,
+            engine,
+            None,
+            100,
+            True,
+            True,
+            True,
+            [{"name": "val1"}, {"name": "val2<?!@#$312,/'][p098"},],
+        )
+
+    def test_estimate_statement_cost(self):
+        mock_cursor = mock.MagicMock()
+        estimate_json = {"a": "b"}
+        mock_cursor.fetchone.return_value = [
+            '{"a": "b"}',
+        ]
+        result = PrestoEngineSpec.estimate_statement_cost(
+            "SELECT * FROM brth_names", mock_cursor
+        )
+        assert result == estimate_json
+
+    def test_estimate_statement_cost_invalid_syntax(self):
+        mock_cursor = mock.MagicMock()
+        mock_cursor.execute.side_effect = Exception()
+        with self.assertRaises(Exception):
+            PrestoEngineSpec.estimate_statement_cost(
+                "DROP TABLE brth_names", mock_cursor
+            )
+
+    def test_get_all_datasource_names(self):
+        df = pd.DataFrame.from_dict(
+            {"table_schema": ["schema1", "schema2"], "table_name": ["name1", "name2"]}
+        )
+        database = mock.MagicMock()
+        database.get_df.return_value = df
+        result = PrestoEngineSpec.get_all_datasource_names(database, "table")
+        expected_result = [
+            DatasourceName(schema="schema1", table="name1"),
+            DatasourceName(schema="schema2", table="name2"),
+        ]
+        assert result == expected_result
+
+    def test_get_create_view(self):
+        mock_execute = mock.MagicMock()
+        mock_fetchall = mock.MagicMock(return_value=[["a", "b,", "c"], ["d", "e"]])
+        database = mock.MagicMock()
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
+            mock_execute
+        )
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.fetchall = (
+            mock_fetchall
+        )
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.poll.return_value = (
+            False
+        )
+        schema = "schema"
+        table = "table"
+        result = PrestoEngineSpec.get_create_view(database, schema=schema, table=table)
+        assert result == "a"
+        mock_execute.assert_called_once_with(f"SHOW CREATE VIEW {schema}.{table}")
+
+    def test_get_create_view_exception(self):
+        mock_execute = mock.MagicMock(side_effect=Exception())
+        database = mock.MagicMock()
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
+            mock_execute
+        )
+        schema = "schema"
+        table = "table"
+        with self.assertRaises(Exception):
+            PrestoEngineSpec.get_create_view(database, schema=schema, table=table)
+
+    def test_get_create_view_database_error(self):
+        from pyhive.exc import DatabaseError
+
+        mock_execute = mock.MagicMock(side_effect=DatabaseError())
+        database = mock.MagicMock()
+        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
+            mock_execute
+        )
+        schema = "schema"
+        table = "table"
+        result = PrestoEngineSpec.get_create_view(database, schema=schema, table=table)
+        assert result is None
+
+    def test_extract_error_message_orig(self):
+        DatabaseError = namedtuple("DatabaseError", ["error_dict"])
+        db_err = DatabaseError(
+            {"errorName": "name", "errorLocation": "location", "message": "msg"}
+        )
+        exception = Exception()
+        exception.orig = db_err
+        result = PrestoEngineSpec._extract_error_message(exception)
+        assert result == "name at location: msg"
+
+    def test_extract_error_message_db_errr(self):
+        from pyhive.exc import DatabaseError
+
+        exception = DatabaseError({"message": "Err message"})
+        result = PrestoEngineSpec._extract_error_message(exception)
+        assert result == "Err message"
+
+    def test_extract_error_message_general_exception(self):
+        exception = Exception("Err message")
+        result = PrestoEngineSpec._extract_error_message(exception)
+        assert result == "Err message"
+
 
 def test_is_readonly():
     def is_readonly(sql: str) -> bool: