Suggest datatypes from a hardcoded whitelist

2 files changed, 78 insertions, 8 deletions

diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index e1731b77..41a06d09 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -354,3 +354,15 @@ def test_auto_escaped_col_names(completer, complete_event):
         Completion(text='custom_func1', start_position=0),
         Completion(text='custom_func2', start_position=0)] +
         list(map(Completion, completer.functions)))
+
+
+@pytest.mark.parametrize('text', [
+    'SELECT 1::DOU',
+    'CREATE TABLE foo (bar DOU',
+    'CREATE FUNCTION foo (bar INT, baz DOU',
+])
+def test_suggest_datatype(text, completer, complete_event):
+    pos = len(text)
+    result = completer.get_completions(
+        Document(text=text, cursor_position=pos), complete_event)
+    assert result == [Completion(text='DOUBLE PRECISION', start_position=-3)]
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index c495e5ac..277936e9 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -366,13 +366,6 @@ def test_specials_included_for_initial_completion(initial_text):
         sorted_dicts([{'type': 'keyword'}, {'type': 'special'}])
 
 
-def test_specials_not_included_after_initial_token():
-    suggestions = suggest_type('create table foo (dt d',
-                               'create table foo (dt d')
-
-    assert sorted_dicts(suggestions) == sorted_dicts([{'type': 'keyword'}])
-
-
 def test_drop_schema_qualified_table_suggests_only_tables():
     text = 'DROP TABLE schema_name.table_name'
     suggestions = suggest_type(text, text)
@@ -388,4 +381,69 @@ def test_handle_pre_completion_comma_gracefully(text):
 
 def test_drop_schema_suggests_schemas():
     sql = 'DROP SCHEMA '
-    assert suggest_type(sql, sql) == [{'type': 'schema'}]
\ No newline at end of file
+    assert suggest_type(sql, sql) == [{'type': 'schema'}]
+    
+
+@pytest.mark.parametrize('text', [
+    'SELECT x::',
+    'SELECT x::y',
+    'SELECT (x + y)::',
+])
+def test_cast_operator_suggests_types(text):
+    assert sorted_dicts(suggest_type(text, text)) == sorted_dicts([
+        {'type': 'datatype', 'schema': []},
+        {'type': 'table', 'schema': []},
+        {'type': 'schema'}])
+
+
+@pytest.mark.parametrize('text', [
+    'SELECT foo::bar.',
+    'SELECT foo::bar.baz',
+    'SELECT (x + y)::bar.',
+])
+def test_cast_operator_suggests_schema_qualified_types(text):
+    assert sorted_dicts(suggest_type(text, text)) == sorted_dicts([
+        {'type': 'datatype', 'schema': 'bar'},
+        {'type': 'table', 'schema': 'bar'}])
+
+
+def test_alter_column_type_suggests_types():
+    q = 'ALTER TABLE foo ALTER COLUMN bar TYPE '
+    assert sorted_dicts(suggest_type(q, q)) == sorted_dicts([
+        {'type': 'datatype', 'schema': []},
+        {'type': 'table', 'schema': []},
+        {'type': 'schema'}])
+
+
+@pytest.mark.parametrize('text', [
+    'CREATE TABLE foo (bar ',
+    'CREATE TABLE foo (bar DOU',
+    'CREATE TABLE foo (bar INT, baz ',
+    'CREATE TABLE foo (bar INT, baz TEXT, qux ',
+    'CREATE FUNCTION foo (bar ',
+    'CREATE FUNCTION foo (bar INT, baz ',
+    'SELECT * FROM foo() AS bar (baz ',
+    'SELECT * FROM foo() AS bar (baz INT, qux ',
+
+    # make sure this doesnt trigger special completion
+    'CREATE TABLE foo (dt d',
+])
+def test_identifier_suggests_types_in_parentheses(text):
+    assert sorted_dicts(suggest_type(text, text)) == sorted_dicts([
+        {'type': 'datatype', 'schema': []},
+        {'type': 'table', 'schema': []},
+        {'type': 'schema'}])
+
+
+@pytest.mark.parametrize('text', [
+    'SELECT foo ',
+    'SELECT foo FROM bar ',
+    'SELECT foo AS bar ',
+    'SELECT foo bar ',
+    'SELECT * FROM foo AS bar ',
+    'SELECT * FROM foo bar ',
+    'SELECT foo FROM (SELECT bar '
+])
+def test_alias_suggests_keywords(text):
+    suggestions = suggest_type(text, text)
+    assert suggestions == [{'type': 'keyword'}]