| # 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 fake_query import ( |
| FakeColumn, |
| FakeFirstValue, |
| FakeQuery, |
| FakeSelectClause, |
| FakeTable) |
| from tests.comparison.common import TableExprList |
| from tests.comparison.db_types import Char, Int |
| from tests.comparison.funcs import AggCount |
| from tests.comparison.query import ( |
| FromClause, |
| InsertClause, |
| InsertStatement, |
| OrderByClause, |
| ValuesClause, |
| ValuesRow, |
| WithClause, |
| WithClauseInlineView) |
| |
| |
| QueryTest = namedtuple( |
| # A QueryTest object contains a SELECT query and all data to verify about it |
| # as other attributes. This allows a new query to be added without need to modify |
| # tests themselves. The various tests cherry-pick which test attributes they need to |
| # verify against the Query. |
| # |
| # If you add a new test, add a new attribute, or perhaps reuse one or more |
| # existing attributes. |
| # |
| # If you add a new test case, add a new item to QUERY_TEST_CASESs array. |
| # |
| # All attributes are required. |
| 'QueryTest', |
| [ |
| # string to represent readable pytest testid |
| 'testid', |
| # Query object, formed via FakeQuery |
| 'query', |
| # textual form of query in Impala dialect |
| 'impala_query_string', |
| # textual form of query in PostgreSQL dialect |
| 'postgres_query_string', |
| # dictionary representing various item counts (see SelectItem property methods) |
| 'select_item_counts', |
| ] |
| ) |
| |
| InsertStatementTest = namedtuple('InsertStatementTest', |
| ['testid', 'query', 'impala_query_string', |
| 'postgres_query_string']) |
| |
| |
| # FakeTables must be declared for use by queries. Tables may be reused as needed for |
| # multiple FakeQueries. |
| SIMPLE_TABLE = FakeTable( |
| 'fake_table', |
| [ |
| FakeColumn('int_col', Int), |
| FakeColumn('char_col', Char), |
| ] |
| ) |
| |
| KUDU_TABLE = FakeTable( |
| 'kudu_table', |
| [ |
| FakeColumn('int_col', Int, is_primary_key=True), |
| FakeColumn('char_col', Char), |
| ] |
| ) |
| |
| FOUR_COL_KUDU_TABLE = FakeTable( |
| 'four_col_kudu_table', |
| [ |
| FakeColumn('int_col1', Int, is_primary_key=True), |
| FakeColumn('char_col1', Char, is_primary_key=True), |
| FakeColumn('int_col2', Int), |
| FakeColumn('char_col2', Char), |
| ] |
| ) |
| |
| |
| ONE_COL_KUDU_TABLE = FakeTable( |
| 'one_col_kudu_table', |
| [ |
| FakeColumn('int_col', Int, is_primary_key=True), |
| ] |
| ) |
| |
| # This can't be used inline because we need its table expressions later. |
| SIMPLE_WITH_CLAUSE = WithClause( |
| TableExprList([ |
| WithClauseInlineView( |
| FakeQuery( |
| select_clause=FakeSelectClause(SIMPLE_TABLE.cols[0]), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| 'with_view' |
| ) |
| ]) |
| ) |
| |
| # All tests involving SELECT queries should be written to use this data set. |
| SELECT_QUERY_TEST_CASES = [ |
| QueryTest( |
| testid='select col from table', |
| query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_TABLE.cols), |
| from_clause=FromClause(SIMPLE_TABLE), |
| ), |
| impala_query_string=( |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'TRIM(fake_table.char_col)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'fake_table.char_col\n' |
| 'FROM fake_table' |
| ), |
| select_item_counts={ |
| 'items': 2, |
| 'basic_items': 2, |
| 'agg_items': 0, |
| 'analytic_items': 0, |
| }, |
| ), |
| QueryTest( |
| testid='select count()', |
| query=FakeQuery( |
| select_clause=FakeSelectClause( |
| AggCount.create_from_args(SIMPLE_TABLE.cols[0])), |
| from_clause=FromClause(SIMPLE_TABLE), |
| ), |
| impala_query_string=( |
| 'SELECT\n' |
| 'COUNT(fake_table.int_col)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'SELECT\n' |
| 'COUNT(fake_table.int_col)\n' |
| 'FROM fake_table' |
| ), |
| select_item_counts={ |
| 'items': 1, |
| 'basic_items': 0, |
| 'agg_items': 1, |
| 'analytic_items': 0, |
| }, |
| ), |
| QueryTest( |
| testid='select first_value(col) over (order by col)', |
| query=FakeQuery( |
| select_clause=FakeSelectClause( |
| FakeFirstValue( |
| SIMPLE_TABLE.cols[0], |
| order_by_clause=OrderByClause([SIMPLE_TABLE.cols[0]]) |
| ), |
| ), |
| from_clause=FromClause(SIMPLE_TABLE), |
| ), |
| impala_query_string=( |
| 'SELECT\n' |
| 'FIRST_VALUE(fake_table.int_col) OVER (ORDER BY fake_table.int_col ASC)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'SELECT\n' |
| 'FIRST_VALUE(fake_table.int_col) OVER (ORDER BY fake_table.int_col ASC)\n' |
| 'FROM fake_table' |
| ), |
| select_item_counts={ |
| 'items': 1, |
| 'basic_items': 0, |
| 'agg_items': 0, |
| 'analytic_items': 1, |
| }, |
| ), |
| ] |
| |
| INSERT_QUERY_TEST_CASES = [ |
| InsertStatementTest( |
| testid='insert into table select cols', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_TABLE.cols), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'TRIM(fake_table.char_col)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'fake_table.char_col\n' |
| 'FROM fake_table' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert into table column permutations select cols', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE, column_list=KUDU_TABLE.cols), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_TABLE.cols), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table (int_col, char_col)\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'TRIM(fake_table.char_col)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table (int_col, char_col)\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'fake_table.char_col\n' |
| 'FROM fake_table' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert into table partial column permutation select 1 col', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE, |
| column_list=[KUDU_TABLE.cols[0]]), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(SIMPLE_TABLE.cols[0]), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert into table select 1 col', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(SIMPLE_TABLE.cols[0]), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert 2 value rows', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE), |
| values_clause=ValuesClause(( |
| ValuesRow((Int(1), Char('a'))), |
| ValuesRow((Int(2), Char('b'))), |
| )), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| "(1, 'a'),\n" |
| "(2, 'b')" |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| "(1, 'a' || ''),\n" |
| "(2, 'b' || '')" |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert 1 value', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE), |
| values_clause=ValuesClause(( |
| ValuesRow((Int(1),)), |
| )), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| '(1)' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| '(1)' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert value row with full column permutation', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE, column_list=KUDU_TABLE.cols), |
| values_clause=ValuesClause(( |
| ValuesRow((Int(1), Char('a'))), |
| )), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table (int_col, char_col)\n' |
| 'VALUES\n' |
| "(1, 'a')" |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table (int_col, char_col)\n' |
| 'VALUES\n' |
| "(1, 'a' || '')" |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert value row with partial column permutation', |
| query=InsertStatement( |
| insert_clause=InsertClause(KUDU_TABLE, |
| column_list=(KUDU_TABLE.cols[0],)), |
| values_clause=ValuesClause(( |
| ValuesRow((Int(1),)), |
| )), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'VALUES\n' |
| '(1)' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'VALUES\n' |
| '(1)' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert values seleted from with clause', |
| query=InsertStatement( |
| with_clause=SIMPLE_WITH_CLAUSE, |
| insert_clause=InsertClause(KUDU_TABLE, |
| column_list=(KUDU_TABLE.cols[0],)), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_WITH_CLAUSE.table_exprs[0].cols), |
| from_clause=FromClause(SIMPLE_WITH_CLAUSE.table_exprs[0]) |
| ), |
| ), |
| impala_query_string=( |
| 'WITH with_view AS (SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table)\n' |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'SELECT\n' |
| 'with_view.int_col\n' |
| 'FROM with_view' |
| ), |
| postgres_query_string=( |
| 'WITH with_view AS (SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table)\n' |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'SELECT\n' |
| 'with_view.int_col\n' |
| 'FROM with_view' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert into table select cols ignore conflicts', |
| query=InsertStatement( |
| insert_clause=InsertClause( |
| KUDU_TABLE, |
| conflict_action=InsertClause.CONFLICT_ACTION_IGNORE), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_TABLE.cols), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'TRIM(fake_table.char_col)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'fake_table.char_col\n' |
| 'FROM fake_table\n' |
| 'ON CONFLICT DO NOTHING' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert 2 value rows ignore conflicts', |
| query=InsertStatement( |
| insert_clause=InsertClause( |
| KUDU_TABLE, |
| conflict_action=InsertClause.CONFLICT_ACTION_IGNORE, |
| ), |
| values_clause=ValuesClause(( |
| ValuesRow((Int(1), Char('a'))), |
| ValuesRow((Int(2), Char('b'))), |
| )), |
| ), |
| impala_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| "(1, 'a'),\n" |
| "(2, 'b')" |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| "(1, 'a' || ''),\n" |
| "(2, 'b' || '')\n" |
| 'ON CONFLICT DO NOTHING' |
| ), |
| ), |
| InsertStatementTest( |
| testid='insert values seleted from with clause ignore conflicts', |
| query=InsertStatement( |
| with_clause=SIMPLE_WITH_CLAUSE, |
| insert_clause=InsertClause( |
| KUDU_TABLE, |
| column_list=(KUDU_TABLE.cols[0],), |
| conflict_action=InsertClause.CONFLICT_ACTION_IGNORE, |
| ), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_WITH_CLAUSE.table_exprs[0].cols), |
| from_clause=FromClause(SIMPLE_WITH_CLAUSE.table_exprs[0]) |
| ), |
| ), |
| impala_query_string=( |
| 'WITH with_view AS (SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table)\n' |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'SELECT\n' |
| 'with_view.int_col\n' |
| 'FROM with_view' |
| ), |
| postgres_query_string=( |
| 'WITH with_view AS (SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table)\n' |
| 'INSERT INTO kudu_table (int_col)\n' |
| 'SELECT\n' |
| 'with_view.int_col\n' |
| 'FROM with_view\n' |
| 'ON CONFLICT DO NOTHING' |
| ), |
| ), |
| InsertStatementTest( |
| testid='upsert into table select cols', |
| query=InsertStatement( |
| insert_clause=InsertClause( |
| KUDU_TABLE, |
| conflict_action=InsertClause.CONFLICT_ACTION_UPDATE), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*SIMPLE_TABLE.cols), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'UPSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'TRIM(fake_table.char_col)\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col,\n' |
| 'fake_table.char_col\n' |
| 'FROM fake_table\n' |
| 'ON CONFLICT (int_col)\n' |
| 'DO UPDATE SET\n' |
| 'char_col = EXCLUDED.char_col' |
| ), |
| ), |
| InsertStatementTest( |
| testid='upsert 2 value rows', |
| query=InsertStatement( |
| insert_clause=InsertClause( |
| KUDU_TABLE, |
| conflict_action=InsertClause.CONFLICT_ACTION_UPDATE, |
| ), |
| values_clause=ValuesClause(( |
| ValuesRow((Int(1), Char('a'))), |
| ValuesRow((Int(2), Char('b'))), |
| )), |
| ), |
| impala_query_string=( |
| 'UPSERT INTO kudu_table\n' |
| 'VALUES\n' |
| "(1, 'a'),\n" |
| "(2, 'b')" |
| ), |
| postgres_query_string=( |
| 'INSERT INTO kudu_table\n' |
| 'VALUES\n' |
| "(1, 'a' || ''),\n" |
| "(2, 'b' || '')\n" |
| 'ON CONFLICT (int_col)\n' |
| 'DO UPDATE SET\n' |
| 'char_col = EXCLUDED.char_col' |
| ), |
| ), |
| InsertStatementTest( |
| testid='upsert select into table with multiple pk / updatable columns', |
| query=InsertStatement( |
| insert_clause=InsertClause( |
| FOUR_COL_KUDU_TABLE, |
| conflict_action=InsertClause.CONFLICT_ACTION_UPDATE), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(*FOUR_COL_KUDU_TABLE.cols), |
| from_clause=FromClause(FOUR_COL_KUDU_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'UPSERT INTO four_col_kudu_table\n' |
| 'SELECT\n' |
| 'four_col_kudu_table.int_col1,\n' |
| 'TRIM(four_col_kudu_table.char_col1),\n' |
| 'four_col_kudu_table.int_col2,\n' |
| 'TRIM(four_col_kudu_table.char_col2)\n' |
| 'FROM four_col_kudu_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO four_col_kudu_table\n' |
| 'SELECT\n' |
| 'four_col_kudu_table.int_col1,\n' |
| 'four_col_kudu_table.char_col1,\n' |
| 'four_col_kudu_table.int_col2,\n' |
| 'four_col_kudu_table.char_col2\n' |
| 'FROM four_col_kudu_table\n' |
| 'ON CONFLICT (int_col1, char_col1)\n' |
| 'DO UPDATE SET\n' |
| 'int_col2 = EXCLUDED.int_col2,\n' |
| 'char_col2 = EXCLUDED.char_col2' |
| ), |
| ), |
| InsertStatementTest( |
| testid='upsert select into table with no updatable columns', |
| query=InsertStatement( |
| insert_clause=InsertClause( |
| ONE_COL_KUDU_TABLE, |
| conflict_action=InsertClause.CONFLICT_ACTION_UPDATE), |
| select_query=FakeQuery( |
| select_clause=FakeSelectClause(SIMPLE_TABLE.cols[0]), |
| from_clause=FromClause(SIMPLE_TABLE) |
| ), |
| ), |
| impala_query_string=( |
| 'UPSERT INTO one_col_kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table' |
| ), |
| postgres_query_string=( |
| 'INSERT INTO one_col_kudu_table\n' |
| 'SELECT\n' |
| 'fake_table.int_col\n' |
| 'FROM fake_table\n' |
| 'ON CONFLICT DO NOTHING' |
| ), |
| ), |
| ] |