path: root/test/test_completion_engine.py

from mycli.packages.completion_engine import suggest_type
import pytest


def sorted_dicts(dicts):
    """input is a list of dicts."""
    return sorted(tuple(x.items()) for x in dicts)


def test_select_suggests_cols_with_visible_table_scope():
    suggestions = suggest_type('SELECT  FROM tabl', 'SELECT ')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tabl']},
        {'type': 'column', 'tables': [(None, 'tabl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


def test_select_suggests_cols_with_qualified_table_scope():
    suggestions = suggest_type('SELECT  FROM sch.tabl', 'SELECT ')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tabl']},
        {'type': 'column', 'tables': [('sch', 'tabl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


@pytest.mark.parametrize('expression', [
    'SELECT * FROM tabl WHERE ',
    'SELECT * FROM tabl WHERE (',
    'SELECT * FROM tabl WHERE foo = ',
    'SELECT * FROM tabl WHERE bar OR ',
    'SELECT * FROM tabl WHERE foo = 1 AND ',
    'SELECT * FROM tabl WHERE (bar > 10 AND ',
    'SELECT * FROM tabl WHERE (bar AND (baz OR (qux AND (',
    'SELECT * FROM tabl WHERE 10 < ',
    'SELECT * FROM tabl WHERE foo BETWEEN ',
    'SELECT * FROM tabl WHERE foo BETWEEN foo AND ',
])
def test_where_suggests_columns_functions(expression):
    suggestions = suggest_type(expression, expression)
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tabl']},
        {'type': 'column', 'tables': [(None, 'tabl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


@pytest.mark.parametrize('expression', [
    'SELECT * FROM tabl WHERE foo IN (',
    'SELECT * FROM tabl WHERE foo IN (bar, ',
])
def test_where_in_suggests_columns(expression):
    suggestions = suggest_type(expression, expression)
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tabl']},
        {'type': 'column', 'tables': [(None, 'tabl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


def test_where_equals_any_suggests_columns_or_keywords():
    text = 'SELECT * FROM tabl WHERE foo = ANY('
    suggestions = suggest_type(text, text)
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tabl']},
        {'type': 'column', 'tables': [(None, 'tabl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'}])


def test_lparen_suggests_cols():
    suggestion = suggest_type('SELECT MAX( FROM tbl', 'SELECT MAX(')
    assert suggestion == [
        {'type': 'column', 'tables': [(None, 'tbl', None)]}]


def test_operand_inside_function_suggests_cols1():
    suggestion = suggest_type(
        'SELECT MAX(col1 +  FROM tbl', 'SELECT MAX(col1 + ')
    assert suggestion == [
        {'type': 'column', 'tables': [(None, 'tbl', None)]}]


def test_operand_inside_function_suggests_cols2():
    suggestion = suggest_type(
        'SELECT MAX(col1 + col2 +  FROM tbl', 'SELECT MAX(col1 + col2 + ')
    assert suggestion == [
        {'type': 'column', 'tables': [(None, 'tbl', None)]}]


def test_select_suggests_cols_and_funcs():
    suggestions = suggest_type('SELECT ', 'SELECT ')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': []},
        {'type': 'column', 'tables': []},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


@pytest.mark.parametrize('expression', [
    'SELECT * FROM ',
    'INSERT INTO ',
    'COPY ',
    'UPDATE ',
    'DESCRIBE ',
    'DESC ',
    'EXPLAIN ',
    'SELECT * FROM foo JOIN ',
])
def test_expression_suggests_tables_views_and_schemas(expression):
    suggestions = suggest_type(expression, expression)
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'table', 'schema': []},
        {'type': 'view', 'schema': []},
        {'type': 'schema'}])


@pytest.mark.parametrize('expression', [
    'SELECT * FROM sch.',
    'INSERT INTO sch.',
    'COPY sch.',
    'UPDATE sch.',
    'DESCRIBE sch.',
    'DESC sch.',
    'EXPLAIN sch.',
    'SELECT * FROM foo JOIN sch.',
])
def test_expression_suggests_qualified_tables_views_and_schemas(expression):
    suggestions = suggest_type(expression, expression)
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'table', 'schema': 'sch'},
        {'type': 'view', 'schema': 'sch'}])


def test_truncate_suggests_tables_and_schemas():
    suggestions = suggest_type('TRUNCATE ', 'TRUNCATE ')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'table', 'schema': []},
        {'type': 'schema'}])


def test_truncate_suggests_qualified_tables():
    suggestions = suggest_type('TRUNCATE sch.', 'TRUNCATE sch.')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'table', 'schema': 'sch'}])


def test_distinct_suggests_cols():
    suggestions = suggest_type('SELECT DISTINCT ', 'SELECT DISTINCT ')
    assert suggestions == [{'type': 'column', 'tables': []}]


def test_col_comma_suggests_cols():
    suggestions = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tbl']},
        {'type': 'column', 'tables': [(None, 'tbl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


def test_table_comma_suggests_tables_and_schemas():
    suggestions = suggest_type('SELECT a, b FROM tbl1, ',
                               'SELECT a, b FROM tbl1, ')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'table', 'schema': []},
        {'type': 'view', 'schema': []},
        {'type': 'schema'}])


def test_into_suggests_tables_and_schemas():
    suggestion = suggest_type('INSERT INTO ', 'INSERT INTO ')
    assert sorted_dicts(suggestion) == sorted_dicts([
        {'type': 'table', 'schema': []},
        {'type': 'view', 'schema': []},
        {'type': 'schema'}])


def test_insert_into_lparen_suggests_cols():
    suggestions = suggest_type('INSERT INTO abc (', 'INSERT INTO abc (')
    assert suggestions == [{'type': 'column', 'tables': [(None, 'abc', None)]}]


def test_insert_into_lparen_partial_text_suggests_cols():
    suggestions = suggest_type('INSERT INTO abc (i', 'INSERT INTO abc (i')
    assert suggestions == [{'type': 'column', 'tables': [(None, 'abc', None)]}]


def test_insert_into_lparen_comma_suggests_cols():
    suggestions = suggest_type('INSERT INTO abc (id,', 'INSERT INTO abc (id,')
    assert suggestions == [{'type': 'column', 'tables': [(None, 'abc', None)]}]


def test_partially_typed_col_name_suggests_col_names():
    suggestions = suggest_type('SELECT * FROM tabl WHERE col_n',
                               'SELECT * FROM tabl WHERE col_n')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'alias', 'aliases': ['tabl']},
        {'type': 'column', 'tables': [(None, 'tabl', None)]},
        {'type': 'function', 'schema': []},
        {'type': 'keyword'},
    ])


def test_dot_suggests_cols_of_a_table_or_schema_qualified_table():
    suggestions = suggest_type('SELECT tabl. FROM tabl', 'SELECT tabl.')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'column', 'tables': [(None, 'tabl', None)]},
        {'type': 'table', 'schema': 'tabl'},
        {'type': 'view', 'schema': 'tabl'},
        {'type': 'function', 'schema': 'tabl'}])


def test_dot_suggests_cols_of_an_alias():
    suggestions = suggest_type('SELECT t1. FROM tabl1 t1, tabl2 t2',
                               'SELECT t1.')
    assert sorted_dicts(suggestions) == sorted_dicts([
        {'type': 'table', 'schema': 't1'},
        {'type': 'view', 'schema': 't1'},
        {'type': 'column', 'tables': [(None, 'tabl1', 't1')]},
        {'type': 'function', 'schema': 't1'}])