kaqing 2.0.52__py3-none-any.whl → 2.0.184__py3-none-any.whl

adam/sql/sql_completer.py

@@ -1,403 +1,118 @@
-from typing import Callable, Iterable
-from prompt_toolkit.completion import CompleteEvent, Completer, Completion
-from prompt_toolkit.document import Document
+from enum import Enum
+from typing import Callable
+
 import sqlparse
-from sqlparse.sql import Statement, Token
-from sqlparse import tokens as T
+from sqlparse.sql import Token
 
 from adam.sql.term_completer import TermCompleter
+from adam.utils_repl.automata_completer import AutomataCompleter
+from adam.sql.sql_state_machine import AthenaStateMachine, CqlStateMachine, SqlStateMachine
+from adam.utils_repl.state_machine import State
 
-class SqlCompleter(Completer):
-    def __init__(self, tables: Callable[[], list[str]], dml: str = None, debug = False):
-        super().__init__()
-        self.dml = dml
-        self.tables = tables
-        self.debug = debug
+__all__ = [
+    "SqlCompleter",
+]
 
-    def get_completions(
-        self, document: Document, complete_event: CompleteEvent
-    ) -> Iterable[Completion]:
-        text = document.text_before_cursor.lstrip()
-        if self.dml:
-            state = f'{self.dml}_'
-            text = f'{self.dml} {text}'
+def default_columns(x: list[str]):
+    return 'id,x.,y.,z.'.split(',')
+
+class SqlVariant(Enum):
+    SQL = 'sql'
+    CQL = 'cql'
+    ATHENA = 'athena'
+
+class SqlCompleter(AutomataCompleter[Token]):
+    def tokens(self, text: str) -> list[Token]:
+        tokens = []
 
-        completer = None
         stmts = sqlparse.parse(text)
-        if not stmts:
-            completer = TermCompleter(['select', 'insert', 'delete', 'update'])
-        else:
-            statement: Statement = stmts[0]
-            state = self.traverse_tokens(text, statement.tokens)
-            if self.debug:
-                print('\n  =>', state)
-            if state == 'dml_incomplete':
-                completer = TermCompleter(['select', 'insert', 'delete', 'update'])
+        if stmts:
+            for stmt in stmts:
+                tokens.extend(stmt.tokens)
+
+        return tokens
+
+    def __init__(self,
+                 tables: Callable[[], list[str]],
+                 dml: str = None,
+                 expandables: dict = {},
+                 variant: SqlVariant = SqlVariant.SQL,
+                 debug = False):
+        machine = SqlStateMachine(debug=debug)
+        if variant == SqlVariant.CQL:
+            machine = CqlStateMachine(debug=debug)
+        elif variant == SqlVariant.ATHENA:
+            machine = AthenaStateMachine(debug=debug)
+        super().__init__(machine, dml, debug)
 
-            elif state == 'select_':
-                completer = TermCompleter(['*'])
-            elif state == 'select_a':
-                completer = TermCompleter(['from'])
-            elif state == 'select_a,':
-                completer = TermCompleter(['*'])
-            elif state == 'select_a_':
-                completer = TermCompleter(['from'])
-            elif state == "select_a_from_":
-                completer = TermCompleter(self.tables())
-            elif state == "select_a_from_x_":
-                completer = TermCompleter(['where', 'group', 'limit'])
-            elif state == "select_a_from_x,":
-                completer = TermCompleter(self.tables())
-            elif state == "select_a_from_x_where_":
-                completer = TermCompleter(['id'])
-            elif state == "select_a_from_x_where_id":
-                completer = TermCompleter(['=', '<', '<=', '>', '>=', '<>', 'like'])
-            elif state == "select_a_from_x_where_id=":
-                completer = TermCompleter(["'"])
-            elif state == "select_a_from_x_where_id=v_":
-                completer = TermCompleter(['and', 'or', 'group', 'limit'])
-            elif state == "select_a_from_x_where_id=v_limit_":
-                completer = TermCompleter(['1'])
-            elif state == "select_a_from_x_group_":
-                completer = TermCompleter(['by'])
-            elif state == "select_a_from_x_group_by_":
-                completer = TermCompleter(['id'])
-            elif state == "select_a_from_x_group_by_a,":
-                completer = TermCompleter(['id'])
-            elif state == "select_a_from_x_group_by_a_":
-                completer = TermCompleter(['limit'])
-            elif state == "select_a_from_x_group_by_a_limit_":
-                completer = TermCompleter(['1'])
+        self.tables = tables
+        if 'columns' not in expandables:
+            expandables['columns'] = default_columns
+        self.expandables = expandables
+        self.variant = variant
+        self.debug = debug
 
-            elif state == "insert_":
-                completer = TermCompleter(['into'])
-            elif state == "insert_into_":
-                completer = TermCompleter(self.tables())
-            elif state == "insert_into_x_":
-                completer = TermCompleter(['values'])
-            elif state == "insert_into_x(":
-                completer = TermCompleter(['id'])
-            elif state == "insert_into_x(a,":
-                completer = TermCompleter(['id'])
-            elif state == "insert_into_x(a)_":
-                completer = TermCompleter(['values('])
-            elif state == "insert_into_x_values":
-                completer = TermCompleter(['('])
-            elif state == "insert_into_x_values(":
-                completer = TermCompleter(["'"])
+    def suggestions_completer(self, state: State, suggestions: str) -> list[str]:
+        if not suggestions:
+            return None
 
-            elif state == "update_":
-                completer = TermCompleter(self.tables())
-            elif state == "update_x_":
-                completer = TermCompleter(['set'])
-            elif state in ["update_x_set_", "update_x_set_a=v,"]:
-                completer = TermCompleter(['id'])
-            elif state == "update_x_set_a":
-                completer = TermCompleter(['='])
-            elif state == "update_x_set_a=":
-                completer = TermCompleter(["'"])
-            elif state == "update_x_set_a=v_":
-                completer = TermCompleter(['where'])
-            elif state == "update_x_set_a=v_where_":
-                completer = TermCompleter(['id'])
-            elif state == "update_x_set_a=v_where_id":
-                completer = TermCompleter(['='])
-            elif state == "update_x_set_a=v_where_id=v_":
-                completer = TermCompleter(['and', 'or'])
+        terms = []
+        for suggestion in suggestions.split(','):
+            terms.extend(self._terms(state, suggestion))
 
-            elif state == "delete_":
-                completer = TermCompleter(['from'])
-            elif state == "delete_from_":
-                completer = TermCompleter(self.tables())
-            elif state == "delete_from_x_":
-                completer = TermCompleter(['where'])
-            elif state == "delete_from_x_where_":
-                completer = TermCompleter(['id'])
-            elif state == "delete_from_x_where_id":
-                completer = TermCompleter(['='])
-            elif state == "delete_from_x_where_id=":
-                completer = TermCompleter(["'"])
-            elif state == "delete_from_x_where_id=v_":
-                completer = TermCompleter(['and', 'or'])
+        return TermCompleter(terms)
 
-        if completer:
-            for c in completer.get_completions(document, complete_event):
-                yield c
+    def _terms(self, state: State, word: str) -> list[str]:
+        terms = []
 
-    def traverse_tokens(self, text: str, tokens: list[Token], state: str = None, indent=0):
-        # state: str = None
-        for token in tokens:
-            if self.debug:
-                if token.ttype == T.Whitespace:
-                    print('_ ', end='')
-                elif token.ttype in [T.DML, T.Wildcard, T.Punctuation]:
-                    print(f'{token.value} ', end='')
-                elif token.ttype:
-                    tks = str(token.ttype).split('.')
-                    typ = tks[len(tks) - 1]
-                    if ' ' in token.value:
-                        print(f'"{token.value}:{typ}" ', end='')
-                    else:
-                        print(f'{token.value}:{typ} ', end='')
-                # print("  " * indent + f"Token: {token.value}, Type: {token.ttype}@{token.ttype.__class__}")
-            node: str = None
-            if token.is_group:
-                state = self.traverse_tokens(text, token.tokens, state, indent + 1)
+        if word.startswith('`') and word.endswith('`'):
+            terms.append(word.strip('`'))
+        elif word == 'tables':
+            terms.extend(self.tables())
+        elif word == 'columns':
+            if 'last_name' in state.context and (n := state.context['last_name']):
+                if 'last_namespace' in state.context and (ns := state.context['last_namespace']):
+                    n = f'{ns}.{n}'
+                terms.extend(self._call_expandable(word, [n]))
             else:
-                if not state:
-                    if token.ttype == T.Keyword.DML and token.value.lower() == 'select':
-                        state = 'select'
-                    if token.ttype == T.Keyword.DML and token.value.lower() == 'insert':
-                        state = 'insert'
-                    if token.ttype == T.Keyword.DML and token.value.lower() == 'update':
-                        state = 'update'
-                    if token.ttype == T.Keyword.DML and token.value.lower() == 'delete':
-                        state = 'delete'
-                    elif token.ttype == T.Name:
-                        state = 'dml_incomplete'
-
-                elif state == 'select':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_'
-                elif state == 'select_':
-                    if token.ttype == T.Name or token.ttype == T.Wildcard:
-                        state = 'select_a'
-                elif state == 'select_a':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_'
-                    elif token.ttype == T.Punctuation and token.value == ',':
-                        state = 'select_a,'
-                elif state == 'select_a,':
-                    if token.ttype == T.Name or token.ttype == T.Wildcard:
-                        state = 'select_a'
-                elif state == 'select_a_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'from':
-                        state = 'select_a_from'
-                elif state == 'select_a_from':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_'
-                elif state == 'select_a_from_':
-                    if token.ttype == T.Name:
-                        state = 'select_a_from_x'
-                elif state == 'select_a_from_x':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_'
-                    elif token.ttype == T.Punctuation and token.value == ',':
-                        state = 'select_a_from_x,'
-                elif state == 'select_a_from_x,':
-                    if token.ttype == T.Name:
-                        state = 'select_a_from_x'
-                elif state == 'select_a_from_x_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'where':
-                        state = 'select_a_from_x_where'
-                    elif token.ttype == T.Keyword and token.value.lower() == 'limit':
-                        state = 'select_a_from_x_where_id=v_limit'
-                    elif token.ttype == T.Keyword and token.value.lower() == 'group':
-                        state = 'select_a_from_x_group'
-                    elif token.ttype == T.Keyword and token.value.lower() == 'group by':
-                        state = 'select_a_from_x_group_by'
-                elif state == 'select_a_from_x_where':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_where_'
-                elif state == 'select_a_from_x_where_':
-                    if token.ttype == T.Name:
-                        state = 'select_a_from_x_where_id'
-                elif state == 'select_a_from_x_where_id':
-                    if token.ttype == T.Operator.Comparison:
-                        state = 'select_a_from_x_where_id='
-                elif state == 'select_a_from_x_where_id=':
-                    if token.ttype in [T.Literal.String.Single, T.Name]:
-                        state = 'select_a_from_x_where_id=v'
-                elif state == 'select_a_from_x_where_id=v':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_where_id=v_'
-                elif state == 'select_a_from_x_where_id=v_':
-                    if token.ttype == T.Keyword and token.value.lower() in ['and', 'or']:
-                        state = 'select_a_from_x_where'
-                    elif token.ttype == T.Keyword and token.value.lower() == 'group':
-                        state = 'select_a_from_x_group'
-                    elif token.ttype == T.Keyword and token.value.lower() == 'limit':
-                        state = 'select_a_from_x_where_id=v_limit'
-                elif state == 'select_a_from_x_group':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_group_'
-                elif state == 'select_a_from_x_group_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'by':
-                        state = 'select_a_from_x_group_by'
-                elif state == 'select_a_from_x_group_by':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_group_by_'
-                elif state == 'select_a_from_x_group_by_':
-                    if token.ttype == T.Name:
-                        state = 'select_a_from_x_group_by_a'
-                elif state == 'select_a_from_x_group_by_a':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_group_by_a_'
-                    elif token.ttype == T.Punctuation and token.value == ',':
-                        state = 'select_a_from_x_group_by_a,'
-                elif state == 'select_a_from_x_group_by_a,':
-                    if token.ttype == T.Name:
-                        state = 'select_a_from_x_group_by_a'
-                elif state == 'select_a_from_x_group_by_a_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'limit':
-                        state = 'select_a_from_x_where_id=v_limit'
-                elif state == 'select_a_from_x_where_id=v_limit':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'select_a_from_x_where_id=v_limit_'
+                terms.extend(self._call_expandable(word, []))
+        elif word == 'partition-columns':
+            terms.extend(self._call_expandable(word, []))
+        elif word == 'table-props':
+            terms.extend(self._call_expandable(word).keys())
+        elif word == 'table-prop-values':
+            if 'last_name' in state.context and state.context['last_name']:
+                table_props = self._call_expandable('table-props')
+                terms.extend(table_props[state.context['last_name']])
+        elif word == 'single':
+            terms.append("'")
+        elif word == 'comma':
+            terms.append(",")
+        elif word in self.machine.expandable_names():
+            terms.extend(self._call_expandable(word))
+        else:
+            terms.append(word)
 
-                elif state == 'insert':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'insert_'
-                elif state == 'insert_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'into':
-                        state = 'insert_into'
-                elif state == 'insert_into':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'insert_into_'
-                elif state == 'insert_into_':
-                    if token.ttype == T.Name:
-                        state = 'insert_into_x'
-                elif state == 'insert_into_x':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'insert_into_x_'
-                    elif token.ttype == T.Punctuation and token.value == '(':
-                        state = 'insert_into_x('
-                elif state == 'insert_into_x_':
-                    if token.ttype == T.Punctuation and token.value == '(':
-                        state = 'insert_into_x('
-                    elif token.ttype == T.Keyword and token.value.lower() == 'values':
-                        state = 'insert_into_x_values'
-                elif state == 'insert_into_x(':
-                    if token.ttype == T.Name:
-                        state = 'insert_into_x(a'
-                elif state == 'insert_into_x(a':
-                    if token.ttype == T.Punctuation and token.value == ',':
-                        state = 'insert_into_x(a,'
-                    elif token.ttype == T.Punctuation and token.value == ')':
-                        state = 'insert_into_x(a)'
-                elif state == 'insert_into_x(a,':
-                    if token.ttype == T.Name:
-                        state = 'insert_into_x(a'
-                elif state == 'insert_into_x(a)':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'insert_into_x(a)_'
-                elif state == 'insert_into_x(a)_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'values':
-                        state = 'insert_into_x_values'
-                elif state == 'insert_into_x_values':
-                    if token.ttype == T.Punctuation and token.value == '(':
-                        state = 'insert_into_x_values('
-                elif state == 'insert_into_x_values(':
-                    if token.ttype in [T.Literal.String.Single, T.Name]:
-                        state = 'insert_into_x_values(v'
-                elif state == 'insert_into_x_values(v':
-                    if token.ttype == T.Punctuation and token.value == ',':
-                        state = 'insert_into_x_values(v,'
-                    elif token.ttype == T.Punctuation and token.value == ')':
-                        state = 'insert_into_x_values(v)'
-                elif state == 'insert_into_x_values(v,':
-                    if token.ttype in [T.Literal.String.Single, T.Name]:
-                        state = 'insert_into_x_values(v'
+        return terms
 
-                elif state == 'update':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'update_'
-                elif state == 'update_':
-                    if token.ttype == T.Name:
-                        state = 'update_x'
-                elif state == 'update_x':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'update_x_'
-                elif state == 'update_x_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'set':
-                        state = 'update_x_set'
-                elif state == 'update_x_set':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'update_x_set_'
-                elif state == 'update_x_set_':
-                    if token.ttype == T.Name:
-                        state = 'update_x_set_a'
-                elif state == 'update_x_set_a':
-                    if token.ttype == T.Operator.Comparison:
-                        state = 'update_x_set_a='
-                elif state == 'update_x_set_a=':
-                    if token.ttype in [T.Literal.String.Single, T.Name]:
-                        state = 'update_x_set_a=v'
-                elif state == 'update_x_set_a=v':
-                    if token.ttype == T.Punctuation and token.value == ',':
-                        state = 'update_x_set_a=v,'
-                    elif token.ttype == T.Text.Whitespace:
-                        state = 'update_x_set_a=v_'
-                elif state == 'update_x_set_a=v,':
-                    if token.ttype == T.Name:
-                        state = 'update_x_set_a'
-                elif state == 'update_x_set_a=v_':
-                    if token.ttype == T.Punctuation and token.value == ',':
-                        state = 'update_x_set_a=v,'
-                    elif token.ttype == T.Keyword and token.value.lower() == 'where':
-                        state = 'update_x_set_a=v_where'
-                elif state == 'update_x_set_a=v_where':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'update_x_set_a=v_where_'
-                elif state == 'update_x_set_a=v_where_':
-                    if token.ttype == T.Name:
-                        state = 'update_x_set_a=v_where_id'
-                elif state == 'update_x_set_a=v_where_id':
-                    if token.ttype == T.Operator.Comparison:
-                        state = 'update_x_set_a=v_where_id='
-                elif state == 'update_x_set_a=v_where_id=':
-                    if token.ttype in [T.Literal.String.Single, T.Name]:
-                        state = 'update_x_set_a=v_where_id=v'
-                elif state == 'update_x_set_a=v_where_id=v':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'update_x_set_a=v_where_id=v_'
-                elif state == 'update_x_set_a=v_where_id=v_':
-                    if token.ttype == T.Keyword and token.value.lower() in ['and', 'or']:
-                        state = 'update_x_set_a=v_where'
+    def _call_expandable(self, name: str, *args):
+        if name in self.expandables:
+            c = self.expandables[name]
+            if args:
+                return c(args)
+            else:
+                return c()
 
-                elif state == 'delete':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'delete_'
-                elif state == 'delete_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'from':
-                        state = 'delete_from'
-                elif state == 'delete_from':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'delete_from_'
-                elif state == 'delete_from_':
-                    if token.ttype == T.Name:
-                        state = 'delete_from_x'
-                elif state == 'delete_from_x':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'delete_from_x_'
-                elif state == 'delete_from_x_':
-                    if token.ttype == T.Keyword and token.value.lower() == 'where':
-                        state = 'delete_from_x_where'
-                elif state == 'delete_from_x_where':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'delete_from_x_where_'
-                elif state == 'delete_from_x_where_':
-                    if token.ttype == T.Name:
-                        state = 'delete_from_x_where_id'
-                elif state == 'delete_from_x_where_id':
-                    if token.ttype == T.Operator.Comparison:
-                        state = 'delete_from_x_where_id='
-                elif state == 'delete_from_x_where_id=':
-                    if token.ttype in [T.Literal.String.Single, T.Name]:
-                        state = 'delete_from_x_where_id=v'
-                elif state == 'delete_from_x_where_id=v':
-                    if token.ttype == T.Text.Whitespace:
-                        state = 'delete_from_x_where_id=v_'
-                elif state == 'delete_from_x_where_id=v_':
-                    if token.ttype == T.Keyword and token.value.lower() in ['and', 'or']:
-                        state = 'delete_from_x_where'
+        return []
 
-        return state
+    def completions_for_nesting(self, dml: str = None):
+        if dml:
+            return {dml: SqlCompleter(self.tables, dml, expandables=self.expandables, variant=self.variant)}
 
-    def completions(table_names: Callable[[], list[str]]):
         return {
-            'delete': SqlCompleter(table_names, 'delete'),
-            'insert': SqlCompleter(table_names, 'insert'),
-            'select': SqlCompleter(table_names, 'select'),
-            'update': SqlCompleter(table_names, 'update'),
-        }
+            word: SqlCompleter(self.tables, word, expandables=self.expandables, variant=self.variant)
+            for word in self.machine.suggestions[''].strip(' ').split(',')
+        }
+
+    def __str__(self):
+        return f'{self.variant}, {self.first_term}'