From ae8d47ed5b17fd206a804091939bc46a447d8c2a Mon Sep 17 00:00:00 2001
From: Jakub Gryko <87192257+jgryko5@users.noreply.github.com>
Date: Fri, 30 Jul 2021 16:39:41 +0200
Subject: [PATCH] Added support for autocompleting CREATE queries

---
 drivers/completer/completer.go      |  7 +++++++
 drivers/completer/completer_test.go | 22 ++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/drivers/completer/completer.go b/drivers/completer/completer.go
index 2d08641f19f..c585af99ec4 100644
--- a/drivers/completer/completer.go
+++ b/drivers/completer/completer.go
@@ -322,6 +322,13 @@ func (c completer) complete(previousWords []string, text []rune) [][]rune {
 	}
 	/* XXX: implement tab completion for DELETE ... USING */
 
+	/* Complete CREATE */
+	if tailMatches(IGNORE_CASE, previousWords, "CREATE") || tailMatches(IGNORE_CASE, previousWords, "CREATE", "TEMP|TEMPORARY") {
+		return completeFromList(text, "DATABASE", "SEQUENCE", "TABLE", "VIEW", "TEMPORARY")
+	}
+	if tailMatches(IGNORE_CASE, previousWords, "CREATE", "TABLE", "*") || tailMatches(IGNORE_CASE, previousWords, "CREATE", "TEMP|TEMPORARY", "TABLE", "*") {
+		return completeFromList(text, "(")
+	}
 	/* INSERT --- can be inside EXPLAIN, RULE, etc */
 	/* Complete INSERT with "INTO" */
 	if tailMatches(IGNORE_CASE, previousWords, "INSERT") {
diff --git a/drivers/completer/completer_test.go b/drivers/completer/completer_test.go
index cc8d8e6ceab..8b189a99f23 100644
--- a/drivers/completer/completer_test.go
+++ b/drivers/completer/completer_test.go
@@ -221,6 +221,28 @@ func TestCompleter(t *testing.T) {
 			[]string{},
 			2,
 		},
+		{
+			"type on create",
+			"CREATE ",
+			7,
+			[]string{
+				"DATABASE",
+				"TABLE",
+				"SEQUENCE",
+				"VIEW",
+				"TEMPORARY",
+			},
+			0,
+		},
+		{
+			"brackets on create table",
+			"CREATE TABLE p ",
+			15,
+			[]string{
+				"(",
+			},
+			0,
+		},
 	}
 
 	completer := NewDefaultCompleter(WithReader(mockReader{}), WithConnStrings([]string{"pg://"}))