# 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 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.sql_parse import ParsedQuery from 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) ) 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)"}, {"name": "array_column", "type": "ARRAY(BIGINT)"}, ] 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)"}, {"name": "row_column.nested_obj", "type": "VARCHAR"}, {"name": "array_column", "type": "ARRAY(BIGINT)"}, ] 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"}] 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))", } ] 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))", }, {"name": "row_column.nested_obj1", "type": "VARCHAR"}, {"name": "row_column.nested_row", "type": "ROW(NESTED_OBJ2 VARCHAR)"}, {"name": "row_column.nested_row.nested_obj2", "type": "VARCHAR"}, ] 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"}, {"name": "row_column.nested_row", "type": "ROW(NESTED_OBJ2 VARCHAR)"}, {"name": "row_column.nested_row.nested_obj2", "type": "VARCHAR"}, ] 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"}, { "name": "array_column", "type": "ARRAY(ROW(NESTED_ARRAY ARRAY(ROW(NESTED_OBJ VARCHAR))))", }, ] 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"}, { "name": "array_column", "type": "ARRAY(ROW(NESTED_ARRAY ARRAY(ROW(NESTED_OBJ VARCHAR))))", }, { "name": "array_column.nested_array", "type": "ARRAY(ROW(NESTED_OBJ VARCHAR))", }, {"name": "array_column.nested_array.nested_obj", "type": "VARCHAR"}, ] 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))", }, {"name": "array_column.nested_array.nested_obj", "type": "VARCHAR"}, ] 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_convert_dttm(self): dttm = self.get_dttm() self.assertEqual( PrestoEngineSpec.convert_dttm("DATE", dttm), "from_iso8601_date('2019-01-02')", ) self.assertEqual( PrestoEngineSpec.convert_dttm("TIMESTAMP", dttm), "from_iso8601_timestamp('2019-01-02T03:04:05.678900')", ) 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_date": False}, {"name": "timestamp", "type": "BIGINT", "is_date": False}, { "name": "user", "type": "ROW(ID BIGINT, FIRST_NAME VARCHAR, LAST_NAME VARCHAR)", "is_date": 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_date": False}, {"name": "timestamp", "type": "BIGINT", "is_date": False}, { "name": "user", "type": "ROW(ID BIGINT, FIRST_NAME VARCHAR, LAST_NAME VARCHAR)", "is_date": False, }, {"name": "user.id", "type": "BIGINT"}, {"name": "user.first_name", "type": "VARCHAR"}, {"name": "user.last_name", "type": "VARCHAR"}, ] 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"}, {"name": "user.first_name", "type": "VARCHAR"}, {"name": "user.last_name", "type": "VARCHAR"}, ] 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): sqla_type = PrestoEngineSpec.get_sqla_column_type("varchar(255)") assert isinstance(sqla_type, types.VARCHAR) assert sqla_type.length == 255 sqla_type = PrestoEngineSpec.get_sqla_column_type("varchar") assert isinstance(sqla_type, types.String) assert sqla_type.length is None sqla_type = PrestoEngineSpec.get_sqla_column_type("char(10)") assert isinstance(sqla_type, types.CHAR) assert sqla_type.length == 10 sqla_type = PrestoEngineSpec.get_sqla_column_type("char") assert isinstance(sqla_type, types.CHAR) assert sqla_type.length is None sqla_type = PrestoEngineSpec.get_sqla_column_type("integer") assert isinstance(sqla_type, types.Integer) sqla_type = PrestoEngineSpec.get_sqla_column_type(None) assert sqla_type is None def test_is_readonly(): def is_readonly(sql: str) -> 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")