summaryrefslogtreecommitdiffstats
path: root/pgcli/pgcompleter.py
blob: ccaf608608932ef696e23436f9c9325abb8bbc5a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from __future__ import print_function
from prompt_toolkit.completion import Completer, Completion
import sqlparse

class PGCompleter(Completer):
    keywords = ['ACCESS', 'ADD', 'ALL', 'ALTER TABLE', 'AND', 'ANY', 'AS',
            'ASC', 'AUDIT', 'BETWEEN', 'BY', 'CHAR', 'CHECK', 'CLUSTER',
            'COLUMN', 'COMMENT', 'COMPRESS', 'CONNECT', 'CREATE', 'CURRENT',
            'DATE', 'DECIMAL', 'DEFAULT', 'DELETE FROM', 'DESC', 'DISTINCT',
            'DROP', 'ELSE', 'EXCLUSIVE', 'EXISTS', 'FILE', 'FLOAT', 'FOR',
            'FROM', 'GRANT', 'GROUP', 'HAVING', 'IDENTIFIED', 'IMMEDIATE',
            'IN', 'INCREMENT', 'INDEX', 'INITIAL', 'INSERT INTO', 'INTEGER',
            'INTERSECT', 'INTO', 'IS', 'LEVEL', 'LIKE', 'LOCK', 'LONG',
            'MAXEXTENTS', 'MINUS', 'MLSLABEL', 'MODE', 'MODIFY', 'NOAUDIT',
            'NOCOMPRESS', 'NOT', 'NOWAIT', 'NULL', 'NUMBER', 'OF', 'OFFLINE',
            'ON', 'ONLINE', 'OPTION', 'OR', 'ORDER', 'PCTFREE', 'PRIOR',
            'PRIVILEGES', 'PUBLIC', 'RAW', 'RENAME', 'RESOURCE', 'REVOKE',
            'ROW', 'ROWID', 'ROWNUM', 'ROWS', 'SELECT', 'SESSION', 'SET',
            'SHARE', 'SIZE', 'SMALLINT', 'START', 'SUCCESSFUL', 'SYNONYM',
            'SYSDATE', 'TABLE', 'THEN', 'TO', 'TRIGGER', 'UID', 'UNION',
            'UNIQUE', 'UPDATE', 'USER', 'VALIDATE', 'VALUES', 'VARCHAR',
            'VARCHAR2', 'VIEW', 'WHENEVER', 'WHERE', 'WITH', ]

    special_commands = []

    table_names = []
    column_names = ['*']
    all_completions = set(keywords)

    def __init__(self, smart_completion=True):
        super(self.__class__, self).__init__()
        self.smart_completion = smart_completion

    def extend_special_commands(self, special_commands):
        self.special_commands.extend(special_commands)

    def extend_keywords(self, additional_keywords):
        self.keywords.extend(additional_keywords)
        self.all_completions.update(additional_keywords)

    def extend_table_names(self, table_names):
        self.table_names.extend(table_names)
        self.all_completions.update(table_names)

    def extend_column_names(self, column_names):
        self.column_names.extend(column_names)
        self.all_completions.update(column_names)

    @staticmethod
    def find_matches(text, collection):
        for item in collection:
            if item.startswith(text) or item.startswith(text.upper()):
                yield Completion(item, -len(text))

    def get_completions(self, document):

        word_before_cursor = document.get_word_before_cursor(WORD=True)

        if not self.smart_completion:
            return self.find_matches(word_before_cursor, self.all_completions)

        # If we've partially typed a word then word_before_cursor won't be an
        # empty string. In that case we want to remove the partially typed
        # string before sending it to the sqlparser. Otherwise the last token
        # will always be the partially typed string which renders the smart
        # completion useless because it will always return the list of keywords
        # as completion.

        if word_before_cursor:
            parsed = sqlparse.parse(document.text[:-len(word_before_cursor)])
        else:
            parsed = sqlparse.parse(document.text)

        last_token = ''
        if parsed:
            last_token = parsed[0].token_prev(len(parsed[0].tokens))
            last_token = last_token.value if last_token else ''

        if last_token.lower() in ('select', 'where', 'having', 'set',
                'order by', 'group by'):
            return self.find_matches(word_before_cursor, self.column_names)
        elif last_token.lower() in ('from', 'update', 'into'):
            return self.find_matches(word_before_cursor, self.table_names)
        else:
            return self.find_matches(word_before_cursor,
                    self.keywords + self.special_commands)