kaqing 2.0.98__py3-none-any.whl → 2.0.203__py3-none-any.whl

adam/sql/sql_completer.py

@@ -1,78 +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
+from sqlparse.sql import Token
 
-from adam.sql.state_machine import StateMachine, StateTo
 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
 
 __all__ = [
     "SqlCompleter",
 ]
 
-DML_COMPLETER = TermCompleter(['select', 'insert', 'delete', 'update'])
-
-def default_columns(tables: list[str]):
+def default_columns(x: list[str]):
     return 'id,x.,y.,z.'.split(',')
 
-class SqlCompleter(Completer):
-    def __init__(self, tables: Callable[[], list[str]], dml: str = None, columns: Callable[[list[str]], list[str]] = default_columns, debug = False):
-        super().__init__()
-        self.dml = dml
+class SqlVariant(Enum):
+    SQL = 'sql'
+    CQL = 'cql'
+    ATHENA = 'athena'
+
+class SqlCompleter(AutomataCompleter[Token]):
+    def tokens(self, text: str) -> list[Token]:
+        tokens = []
+
+        stmts = sqlparse.parse(text)
+        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)
+
         self.tables = tables
-        self.columns = columns
+        if 'columns' not in expandables:
+            expandables['columns'] = default_columns
+        self.expandables = expandables
+        self.variant = variant
         self.debug = debug
-        self.machine = StateMachine(debug=self.debug)
-
-    def get_completions(
-        self, document: Document, complete_event: CompleteEvent
-    ) -> Iterable[Completion]:
-        text = document.text_before_cursor.lstrip()
-        state = ''
-        if self.dml:
-            state = f'{self.dml}_'
-            text = f'{self.dml} {text}'
-
-        completer: Completer = None
-        stmts = sqlparse.parse(text)
-        if not stmts:
-            completer = DML_COMPLETER
+
+    def suggestions_completer(self, state: State, suggestions: str) -> list[str]:
+        if not suggestions:
+            return None
+
+        terms = []
+        for suggestion in suggestions.split(','):
+            terms.extend(self._terms(state, suggestion))
+
+        return TermCompleter(terms)
+
+    def _terms(self, state: State, word: str) -> list[str]:
+        terms = []
+
+        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:
+                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:
-            statement: Statement = stmts[0]
-            state: StateTo = self.machine.traverse_tokens(statement.tokens, StateTo(state))
-            if self.debug:
-                print('\n  =>', state.to_s if isinstance(state, StateTo) else '')
-
-            if not state or not state.to_s:
-                completer = DML_COMPLETER
-
-            if state and state.to_s in self.machine.suggestions:
-                terms = []
-
-                for word in self.machine.suggestions[state.to_s].strip(' ').split(','):
-                    if word == 'tables':
-                        terms.extend(self.tables())
-                    elif word == 'columns':
-                        terms.extend(self.columns([]))
-                    elif word == 'single':
-                        terms.append("'")
-                    elif word == 'comma':
-                        terms.append(",")
-                    else:
-                        terms.append(word)
-
-                if terms:
-                    completer = TermCompleter(terms)
-
-        if completer:
-            for c in completer.get_completions(document, complete_event):
-                yield c
-
-    def completions(table_names: Callable[[], list[str]], columns: Callable[[list[str]], list[str]] = default_columns):
+            terms.append(word)
+
+        return terms
+
+    def _call_expandable(self, name: str, *args):
+        if name in self.expandables:
+            c = self.expandables[name]
+            if args:
+                return c(args)
+            else:
+                return c()
+
+        return []
+
+    def completions_for_nesting(self, dml: str = None):
+        if dml:
+            return {dml: SqlCompleter(self.tables, dml, expandables=self.expandables, variant=self.variant)}
+
         return {
-            'delete': SqlCompleter(table_names, 'delete', columns=columns),
-            'insert': SqlCompleter(table_names, 'insert', columns=columns),
-            'select': SqlCompleter(table_names, 'select', columns=columns),
-            'update': SqlCompleter(table_names, 'update', columns=columns),
-        }
+            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}'