diff options
author | Amjith Ramanujam <amjith@newrelic.com> | 2014-12-18 10:07:48 -0800 |
---|---|---|
committer | Amjith Ramanujam <amjith@newrelic.com> | 2014-12-18 10:07:48 -0800 |
commit | 7d98b4a90bddeb60351d4747a72e1a4e17427c2b (patch) | |
tree | ac7a306ad4aed41075c0b27c77345c651228dab1 | |
parent | 61174dff400d1af4d6245b1fc5c033e070a2760b (diff) |
Beginnings of a table specific column completion.
-rw-r--r-- | pgcli/packages/parseutils.py | 41 | ||||
-rw-r--r-- | pgcli/pgcompleter.py | 7 | ||||
-rw-r--r-- | pgcli/pgexecute.py | 3 |
3 files changed, 45 insertions, 6 deletions
diff --git a/pgcli/packages/parseutils.py b/pgcli/packages/parseutils.py index e2eb2eee..e3c7be36 100644 --- a/pgcli/packages/parseutils.py +++ b/pgcli/packages/parseutils.py @@ -1,4 +1,7 @@ import re +import sqlparse +from sqlparse.sql import IdentifierList, Identifier +from sqlparse.tokens import Keyword, DML # This matches only alphanumerics and underscores. _LAST_WORD_RE = re.compile(r'(\w+)$') @@ -48,6 +51,40 @@ def last_word(text, include_special_chars=False): else: return '' -def extract_tables(sql): - return [] +def is_subselect(parsed): + if not parsed.is_group(): + return False + for item in parsed.tokens: + if item.ttype is DML and item.value.upper() == 'SELECT': + return True + return False + +def extract_from_part(parsed): + from_seen = False + for item in parsed.tokens: + if from_seen: + if is_subselect(item): + for x in extract_from_part(item): + yield x + elif item.ttype is Keyword: + raise StopIteration + else: + yield item + elif item.ttype is Keyword and item.value.upper() == 'FROM': + from_seen = True +def extract_table_identifiers(token_stream): + for item in token_stream: + if isinstance(item, IdentifierList): + for identifier in item.get_identifiers(): + yield identifier.get_name() + elif isinstance(item, Identifier): + yield item.get_name() + # It's a bug to check for Keyword here, but in the example + # above some tables names are identified as keywords... + elif item.ttype is Keyword: + yield item.value + +def extract_tables(sql): + stream = extract_from_part(sqlparse.parse(sql)[0]) + return list(extract_table_identifiers(stream)) diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py index 533b7f33..b4772625 100644 --- a/pgcli/pgcompleter.py +++ b/pgcli/pgcompleter.py @@ -32,9 +32,8 @@ class PGCompleter(Completer): databases = [] tables = [] # This will create a defaultdict which is initialized with a list that has - # a '*' by default. Inspired by the constant_factory example in - # https://docs.python.org/2/library/collections.html#defaultdict-examples - columns = defaultdict(itertools.repeat(['*']).next) + # a '*' by default. + columns = defaultdict(lambda: ['*']) all_completions = set(keywords) def __init__(self, smart_completion=True): @@ -88,6 +87,8 @@ class PGCompleter(Completer): scoped_cols = [] for table in scope: scoped_cols.extend(self.columns[table]) + print ('scope:', scope) + print ('scoped_cols:', scoped_cols) return self.find_matches(word_before_cursor, scoped_cols) elif category == 'columns-and-functions': scoped_cols = [] diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 6bfa5103..61c363fe 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -120,7 +120,8 @@ class PGExecute(object): def columns(self, table): with self.conn.cursor() as cur: cur.execute(self.columns_query, (table,)) - return [x[0] for x in cur.fetchall()] + cols = [x[0] for x in cur.fetchall()] + return cols def all_columns(self): columns = set() |