# 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. from collections import namedtuple from unittest import mock, skipUnless import pandas as pd from sqlalchemy import types from sqlalchemy.engine.result import RowProxy from sqlalchemy.sql import select from superset.db_engine_specs.presto import PrestoEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.sql_parse import ParsedQuery from superset.utils.core import DatasourceName, GenericDataType from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec class TestPrestoDbEngineSpec(TestDbEngineSpec): @skipUnless(TestDbEngineSpec.is_module_installed("pyhive"), "pyhive not installed") def test_get_datatype_presto(self): self.assertEqual("STRING", PrestoEngineSpec.get_datatype("string")) def test_presto_get_view_names_return_empty_list( self, ): # pylint: disable=invalid-name self.assertEqual( [], 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() keymap = { "Column": (None, None, 0), "Type": (None, None, 1), "Null": (None, None, 2), } row = RowProxy(mock.Mock(), column, [None, None, None, None], keymap) inspector.bind.execute = mock.Mock(return_value=[row]) results = PrestoEngineSpec.get_columns(inspector, "", "") self.assertEqual(len(expected_results), len(results)) for expected_result, result in zip(expected_results, results): self.assertEqual(expected_result[0], result["name"]) self.assertEqual(expected_result[1], str(result["type"])) def test_presto_get_column(self): presto_column = ("column_name", "boolean", "") expected_results = [("column_name", "BOOLEAN")] self.verify_presto_column(presto_column, expected_results) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_get_simple_row_column(self): presto_column = ("column_name", "row(nested_obj double)", "") expected_results = [("column_name", "ROW"), ("column_name.nested_obj", "FLOAT")] self.verify_presto_column(presto_column, expected_results) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_get_simple_row_column_with_name_containing_whitespace(self): presto_column = ("column name", "row(nested_obj double)", "") expected_results = [("column name", "ROW"), ("column name.nested_obj", "FLOAT")] self.verify_presto_column(presto_column, expected_results) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_get_simple_row_column_with_tricky_nested_field_name(self): presto_column = ("column_name", 'row("Field Name(Tricky, Name)" double)', "") expected_results = [ ("column_name", "ROW"), ('column_name."Field Name(Tricky, Name)"', "FLOAT"), ] self.verify_presto_column(presto_column, expected_results) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_get_simple_array_column(self): presto_column = ("column_name", "array(double)", "") expected_results = [("column_name", "ARRAY")] self.verify_presto_column(presto_column, expected_results) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_get_row_within_array_within_row_column(self): presto_column = ( "column_name", "row(nested_array array(row(nested_row double)), nested_obj double)", "", ) expected_results = [ ("column_name", "ROW"), ("column_name.nested_array", "ARRAY"), ("column_name.nested_array.nested_row", "FLOAT"), ("column_name.nested_obj", "FLOAT"), ] self.verify_presto_column(presto_column, expected_results) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_get_array_within_row_within_array_column(self): presto_column = ( "column_name", "array(row(nested_array array(double), nested_obj double))", "", ) expected_results = [ ("column_name", "ARRAY"), ("column_name.nested_array", "ARRAY"), ("column_name.nested_obj", "FLOAT"), ] self.verify_presto_column(presto_column, expected_results) def test_presto_get_fields(self): cols = [ {"name": "column"}, {"name": "column.nested_obj"}, {"name": 'column."quoted.nested obj"'}, ] actual_results = PrestoEngineSpec._get_fields(cols) expected_results = [ {"name": '"column"', "label": "column"}, {"name": '"column"."nested_obj"', "label": "column.nested_obj"}, { "name": '"column"."quoted.nested obj"', "label": 'column."quoted.nested obj"', }, ] for actual_result, expected_result in zip(actual_results, expected_results): self.assertEqual(actual_result.element.name, expected_result["name"]) self.assertEqual(actual_result.name, expected_result["label"]) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_expand_data_with_simple_structural_columns(self): cols = [ {"name": "row_column", "type": "ROW(NESTED_OBJ VARCHAR)", "is_dttm": False}, {"name": "array_column", "type": "ARRAY(BIGINT)", "is_dttm": False}, ] data = [ {"row_column": ["a"], "array_column": [1, 2, 3]}, {"row_column": ["b"], "array_column": [4, 5, 6]}, ] actual_cols, actual_data, actual_expanded_cols = PrestoEngineSpec.expand_data( cols, data ) expected_cols = [ {"name": "row_column", "type": "ROW(NESTED_OBJ VARCHAR)", "is_dttm": False}, {"name": "row_column.nested_obj", "type": "VARCHAR", "is_dttm": False}, {"name": "array_column", "type": "ARRAY(BIGINT)", "is_dttm": False}, ] expected_data = [ {"array_column": 1, "row_column": ["a"], "row_column.nested_obj": "a"}, {"array_column": 2, "row_column": "", "row_column.nested_obj": ""}, {"array_column": 3, "row_column": "", "row_column.nested_obj": ""}, {"array_column": 4, "row_column": ["b"], "row_column.nested_obj": "b"}, {"array_column": 5, "row_column": "", "row_column.nested_obj": ""}, {"array_column": 6, "row_column": "", "row_column.nested_obj": ""}, ] expected_expanded_cols = [ {"name": "row_column.nested_obj", "type": "VARCHAR", "is_dttm": False} ] self.assertEqual(actual_cols, expected_cols) self.assertEqual(actual_data, expected_data) self.assertEqual(actual_expanded_cols, expected_expanded_cols) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_expand_data_with_complex_row_columns(self): cols = [ { "name": "row_column", "type": "ROW(NESTED_OBJ1 VARCHAR, NESTED_ROW ROW(NESTED_OBJ2 VARCHAR))", "is_dttm": False, } ] data = [{"row_column": ["a1", ["a2"]]}, {"row_column": ["b1", ["b2"]]}] actual_cols, actual_data, actual_expanded_cols = PrestoEngineSpec.expand_data( cols, data ) expected_cols = [ { "name": "row_column", "type": "ROW(NESTED_OBJ1 VARCHAR, NESTED_ROW ROW(NESTED_OBJ2 VARCHAR))", "is_dttm": False, }, {"name": "row_column.nested_obj1", "type": "VARCHAR", "is_dttm": False}, { "name": "row_column.nested_row", "type": "ROW(NESTED_OBJ2 VARCHAR)", "is_dttm": False, }, { "name": "row_column.nested_row.nested_obj2", "type": "VARCHAR", "is_dttm": False, }, ] expected_data = [ { "row_column": ["a1", ["a2"]], "row_column.nested_obj1": "a1", "row_column.nested_row": ["a2"], "row_column.nested_row.nested_obj2": "a2", }, { "row_column": ["b1", ["b2"]], "row_column.nested_obj1": "b1", "row_column.nested_row": ["b2"], "row_column.nested_row.nested_obj2": "b2", }, ] expected_expanded_cols = [ {"name": "row_column.nested_obj1", "type": "VARCHAR", "is_dttm": False}, { "name": "row_column.nested_row", "type": "ROW(NESTED_OBJ2 VARCHAR)", "is_dttm": False, }, { "name": "row_column.nested_row.nested_obj2", "type": "VARCHAR", "is_dttm": False, }, ] self.assertEqual(actual_cols, expected_cols) self.assertEqual(actual_data, expected_data) self.assertEqual(actual_expanded_cols, expected_expanded_cols) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_expand_data_with_complex_row_columns_and_null_values(self): cols = [ { "name": "row_column", "type": "ROW(NESTED_ROW ROW(NESTED_OBJ VARCHAR))", "is_dttm": False, } ] data = [ {"row_column": '[["a"]]'}, {"row_column": "[[null]]"}, {"row_column": "[null]"}, {"row_column": "null"}, ] actual_cols, actual_data, actual_expanded_cols = PrestoEngineSpec.expand_data( cols, data ) expected_cols = [ { "name": "row_column", "type": "ROW(NESTED_ROW ROW(NESTED_OBJ VARCHAR))", "is_dttm": False, }, { "name": "row_column.nested_row", "type": "ROW(NESTED_OBJ VARCHAR)", "is_dttm": False, }, { "name": "row_column.nested_row.nested_obj", "type": "VARCHAR", "is_dttm": False, }, ] expected_data = [ { "row_column": [["a"]], "row_column.nested_row": ["a"], "row_column.nested_row.nested_obj": "a", }, { "row_column": [[None]], "row_column.nested_row": [None], "row_column.nested_row.nested_obj": None, }, { "row_column": [None], "row_column.nested_row": None, "row_column.nested_row.nested_obj": "", }, { "row_column": None, "row_column.nested_row": "", "row_column.nested_row.nested_obj": "", }, ] expected_expanded_cols = [ { "name": "row_column.nested_row", "type": "ROW(NESTED_OBJ VARCHAR)", "is_dttm": False, }, { "name": "row_column.nested_row.nested_obj", "type": "VARCHAR", "is_dttm": False, }, ] self.assertEqual(actual_cols, expected_cols) self.assertEqual(actual_data, expected_data) self.assertEqual(actual_expanded_cols, expected_expanded_cols) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_expand_data_with_complex_array_columns(self): cols = [ {"name": "int_column", "type": "BIGINT", "is_dttm": False}, { "name": "array_column", "type": "ARRAY(ROW(NESTED_ARRAY ARRAY(ROW(NESTED_OBJ VARCHAR))))", "is_dttm": False, }, ] data = [ {"int_column": 1, "array_column": [[[["a"], ["b"]]], [[["c"], ["d"]]]]}, {"int_column": 2, "array_column": [[[["e"], ["f"]]], [[["g"], ["h"]]]]}, ] actual_cols, actual_data, actual_expanded_cols = PrestoEngineSpec.expand_data( cols, data ) expected_cols = [ {"name": "int_column", "type": "BIGINT", "is_dttm": False}, { "name": "array_column", "type": "ARRAY(ROW(NESTED_ARRAY ARRAY(ROW(NESTED_OBJ VARCHAR))))", "is_dttm": False, }, { "name": "array_column.nested_array", "type": "ARRAY(ROW(NESTED_OBJ VARCHAR))", "is_dttm": False, }, { "name": "array_column.nested_array.nested_obj", "type": "VARCHAR", "is_dttm": False, }, ] expected_data = [ { "array_column": [[["a"], ["b"]]], "array_column.nested_array": ["a"], "array_column.nested_array.nested_obj": "a", "int_column": 1, }, { "array_column": "", "array_column.nested_array": ["b"], "array_column.nested_array.nested_obj": "b", "int_column": "", }, { "array_column": [[["c"], ["d"]]], "array_column.nested_array": ["c"], "array_column.nested_array.nested_obj": "c", "int_column": "", }, { "array_column": "", "array_column.nested_array": ["d"], "array_column.nested_array.nested_obj": "d", "int_column": "", }, { "array_column": [[["e"], ["f"]]], "array_column.nested_array": ["e"], "array_column.nested_array.nested_obj": "e", "int_column": 2, }, { "array_column": "", "array_column.nested_array": ["f"], "array_column.nested_array.nested_obj": "f", "int_column": "", }, { "array_column": [[["g"], ["h"]]], "array_column.nested_array": ["g"], "array_column.nested_array.nested_obj": "g", "int_column": "", }, { "array_column": "", "array_column.nested_array": ["h"], "array_column.nested_array.nested_obj": "h", "int_column": "", }, ] expected_expanded_cols = [ { "name": "array_column.nested_array", "type": "ARRAY(ROW(NESTED_OBJ VARCHAR))", "is_dttm": False, }, { "name": "array_column.nested_array.nested_obj", "type": "VARCHAR", "is_dttm": False, }, ] self.assertEqual(actual_cols, expected_cols) self.assertEqual(actual_data, expected_data) self.assertEqual(actual_expanded_cols, expected_expanded_cols) def test_presto_extra_table_metadata(self): db = mock.Mock() db.get_indexes = mock.Mock(return_value=[{"column_names": ["ds", "hour"]}]) db.get_extra = mock.Mock(return_value={}) df = pd.DataFrame({"ds": ["01-01-19"], "hour": [1]}) db.get_df = mock.Mock(return_value=df) PrestoEngineSpec.get_create_view = mock.Mock(return_value=None) result = PrestoEngineSpec.extra_table_metadata(db, "test_table", "test_schema") self.assertEqual({"ds": "01-01-19", "hour": 1}, result["partitions"]["latest"]) def test_presto_where_latest_partition(self): db = mock.Mock() db.get_indexes = mock.Mock(return_value=[{"column_names": ["ds", "hour"]}]) db.get_extra = mock.Mock(return_value={}) df = pd.DataFrame({"ds": ["01-01-19"], "hour": [1]}) db.get_df = mock.Mock(return_value=df) columns = [{"name": "ds"}, {"name": "hour"}] result = PrestoEngineSpec.where_latest_partition( "test_table", "test_schema", db, select(), columns ) query_result = str(result.compile(compile_kwargs={"literal_binds": True})) self.assertEqual("SELECT \nWHERE ds = '01-01-19' AND hour = 1", query_result) def test_query_cost_formatter(self): raw_cost = [ { "inputTableColumnInfos": [ { "table": { "catalog": "hive", "schemaTable": { "schema": "default", "table": "fact_passenger_state", }, }, "columnConstraints": [ { "columnName": "ds", "typeSignature": "varchar", "domain": { "nullsAllowed": False, "ranges": [ { "low": { "value": "2019-07-10", "bound": "EXACTLY", }, "high": { "value": "2019-07-10", "bound": "EXACTLY", }, } ], }, } ], "estimate": { "outputRowCount": 9.04969899e8, "outputSizeInBytes": 3.54143678301e11, "cpuCost": 3.54143678301e11, "maxMemory": 0.0, "networkCost": 0.0, }, } ], "estimate": { "outputRowCount": 9.04969899e8, "outputSizeInBytes": 3.54143678301e11, "cpuCost": 3.54143678301e11, "maxMemory": 0.0, "networkCost": 3.54143678301e11, }, } ] formatted_cost = PrestoEngineSpec.query_cost_formatter(raw_cost) expected = [ { "Output count": "904 M rows", "Output size": "354 GB", "CPU cost": "354 G", "Max memory": "0 B", "Network cost": "354 G", } ] self.assertEqual(formatted_cost, expected) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", {"PRESTO_EXPAND_DATA": True}, clear=True, ) def test_presto_expand_data_array(self): cols = [ {"name": "event_id", "type": "VARCHAR", "is_dttm": False}, {"name": "timestamp", "type": "BIGINT", "is_dttm": False}, { "name": "user", "type": "ROW(ID BIGINT, FIRST_NAME VARCHAR, LAST_NAME VARCHAR)", "is_dttm": False, }, ] data = [ { "event_id": "abcdef01-2345-6789-abcd-ef0123456789", "timestamp": "1595895506219", "user": '[1, "JOHN", "DOE"]', } ] actual_cols, actual_data, actual_expanded_cols = PrestoEngineSpec.expand_data( cols, data ) expected_cols = [ {"name": "event_id", "type": "VARCHAR", "is_dttm": False}, {"name": "timestamp", "type": "BIGINT", "is_dttm": False}, { "name": "user", "type": "ROW(ID BIGINT, FIRST_NAME VARCHAR, LAST_NAME VARCHAR)", "is_dttm": False, }, {"name": "user.id", "type": "BIGINT", "is_dttm": False}, {"name": "user.first_name", "type": "VARCHAR", "is_dttm": False}, {"name": "user.last_name", "type": "VARCHAR", "is_dttm": False}, ] expected_data = [ { "event_id": "abcdef01-2345-6789-abcd-ef0123456789", "timestamp": "1595895506219", "user": [1, "JOHN", "DOE"], "user.id": 1, "user.first_name": "JOHN", "user.last_name": "DOE", } ] expected_expanded_cols = [ {"name": "user.id", "type": "BIGINT", "is_dttm": False}, {"name": "user.first_name", "type": "VARCHAR", "is_dttm": False}, {"name": "user.last_name", "type": "VARCHAR", "is_dttm": False}, ] self.assertEqual(actual_cols, expected_cols) self.assertEqual(actual_data, expected_data) self.assertEqual(actual_expanded_cols, expected_expanded_cols) def test_get_sqla_column_type(self): column_spec = PrestoEngineSpec.get_column_spec("varchar(255)") assert isinstance(column_spec.sqla_type, types.VARCHAR) assert column_spec.sqla_type.length == 255 self.assertEqual(column_spec.generic_type, GenericDataType.STRING) column_spec = PrestoEngineSpec.get_column_spec("varchar") assert isinstance(column_spec.sqla_type, types.String) assert column_spec.sqla_type.length is None self.assertEqual(column_spec.generic_type, GenericDataType.STRING) column_spec = PrestoEngineSpec.get_column_spec("char(10)") assert isinstance(column_spec.sqla_type, types.CHAR) assert column_spec.sqla_type.length == 10 self.assertEqual(column_spec.generic_type, GenericDataType.STRING) column_spec = PrestoEngineSpec.get_column_spec("char") assert isinstance(column_spec.sqla_type, types.CHAR) assert column_spec.sqla_type.length is None self.assertEqual(column_spec.generic_type, GenericDataType.STRING) column_spec = PrestoEngineSpec.get_column_spec("integer") assert isinstance(column_spec.sqla_type, types.Integer) self.assertEqual(column_spec.generic_type, GenericDataType.NUMERIC) column_spec = PrestoEngineSpec.get_column_spec("time") assert isinstance(column_spec.sqla_type, types.Time) self.assertEqual(column_spec.generic_type, GenericDataType.TEMPORAL) column_spec = PrestoEngineSpec.get_column_spec("timestamp") assert isinstance(column_spec.sqla_type, types.TIMESTAMP) self.assertEqual(column_spec.generic_type, GenericDataType.TEMPORAL) 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 bool: return PrestoEngineSpec.is_readonly_query(ParsedQuery(sql)) assert not is_readonly("SET hivevar:desc='Legislators'") assert not is_readonly("UPDATE t1 SET col1 = NULL") assert not is_readonly("INSERT OVERWRITE TABLE tabB SELECT a.Age FROM TableA") assert is_readonly("SHOW LOCKS test EXTENDED") assert is_readonly("EXPLAIN SELECT 1") assert is_readonly("SELECT 1") assert is_readonly("WITH (SELECT 1) bla SELECT * from bla")