kaqing 2.0.103__tar.gz → 2.0.105__tar.gz

{kaqing-2.0.103 → kaqing-2.0.105}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaqing
-Version: 2.0.103
+Version: 2.0.105
 Summary: UNKNOWN
 Home-page: UNKNOWN
 License: UNKNOWN

{kaqing-2.0.103 → kaqing-2.0.105}/adam/commands/audit/audit.py

@@ -1,12 +1,13 @@
 import click
 
 from adam.commands.audit.audit_repair_tables import AuditRepairTables
+from adam.commands.audit.audit_run import AuditRun
 from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.sql.sql_completer import SqlCompleter
 from adam.utils import log2
-from adam.utils_athena import audit_column_names, audit_table_names, run_audit_query
+from adam.utils_audits import Audits
 
 class Audit(Command):
     COMMAND = 'audit'
@@ -41,7 +42,7 @@ class Audit(Command):
             else:
                 log2(sql)
 
-            run_audit_query(sql)
+            Audits.run_audit_query(sql)
 
         return state
 
@@ -51,23 +52,23 @@ class Audit(Command):
                 Config().wait_log(f'Inspecting audit database schema...')
                 self.schema_read = True
                 # warm up the caches first time when l: drive is accessed
-                audit_column_names()
-                audit_column_names(partition_cols_only=True)
+                Audits.audit_column_names()
+                Audits.audit_column_names(partition_cols_only=True)
 
             return super().completion(state) | SqlCompleter(
-                lambda: audit_table_names(),
-                columns=lambda table: audit_column_names(),
-                partition_columns=lambda table: audit_column_names(partition_cols_only=True),
+                lambda: Audits.audit_table_names(),
+                columns=lambda table: Audits.audit_column_names(),
+                partition_columns=lambda table: Audits.audit_column_names(partition_cols_only=True),
                 variant='athena'
             ).completions_for_nesting()
 
         return {}
 
     def cmd_list():
-        return [AuditRepairTables()]
+        return [AuditRepairTables(), AuditRun()]
 
     def help(self, _: ReplState):
-        return f'[{Audit.COMMAND}] <sql-statements>\t run SQL queries on Authena audit database'
+        return f'[{Audit.COMMAND}] [<sql-statements>]\t run SQL queries on Authena audit database'
 
 class AuditCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):

{kaqing-2.0.103 → kaqing-2.0.105}/adam/commands/audit/audit_repair_tables.py

@@ -1,12 +1,11 @@
 import concurrent
 import time
-import requests
 
 from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.utils import log, log2
-from adam.utils_athena import audit_query, run_audit_query
+from adam.utils_audits import AuditMeta, Audits
 
 class AuditRepairTables(Command):
     COMMAND = 'audit repair'
@@ -30,11 +29,12 @@ class AuditRepairTables(Command):
 
         state, args = self.apply_state(args, state)
 
-        tables = Config().get('audit.tables', 'audit').split(',')
+        tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
         if args:
             tables = args
 
-        self.repair(tables)
+        meta = Audits.get_meta()
+        self.repair(tables, meta)
 
         return state
 
@@ -52,43 +52,20 @@ class AuditRepairTables(Command):
     def auto_repair(self, hours: int):
         self.auto_repaired = True
 
-        state, _, rs = audit_query(f'select * from meta')
-        if state == 'SUCCEEDED':
-            do_repair = True
-            if len(rs) > 1:
-                try:
-                    checked_in = float(rs[1]['Data'][0]['VarCharValue'])
-                    do_repair = checked_in + hours * 60 * 60 < time.time()
-                except:
-                    pass
-
-            if do_repair:
-                tables = Config().get('audit.athena.tables', 'audit').split(',')
-                self.repair(tables, show_sql=True)
-                log2(f'Audit tables have been auto-repaired.')
-
-    def repair(self, tables: list[str], show_sql = False):
+        meta: AuditMeta = Audits.get_meta()
+        if meta.partitions_last_checked + hours * 60 * 60 < time.time():
+            tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
+            self.repair(tables, meta, show_sql=True)
+            log2(f'Audit tables have been auto-repaired.')
+
+    def repair(self, tables: list[str], meta: AuditMeta, show_sql = False):
         with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
             for table in tables:
                 if show_sql:
                     log(f'MSCK REPAIR TABLE {table}')
 
-                executor.submit(run_audit_query, f'MSCK REPAIR TABLE {table}', None,)
-            executor.submit(self.check_in,)
+                executor.submit(Audits.run_audit_query, f'MSCK REPAIR TABLE {table}', None,)
+            executor.submit(Audits.put_meta, Audits.PARTITIONS_ADDED, meta,)
 
     def help(self, _: ReplState):
-        return f"{AuditRepairTables.COMMAND} \t run MSCK REPAIR command for new partition discovery"
-
-    def check_in(self):
-        payload = {
-            'action': 'check-in'
-        }
-        audit_endpoint = Config().get("audit.endpoint", "https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/")
-        try:
-            response = requests.post(audit_endpoint, json=payload, timeout=Config().get("audit.timeout", 10))
-            if response.status_code in [200, 201]:
-                Config().debug(response.text)
-            else:
-                log2(f"Error: {response.status_code} {response.text}")
-        except requests.exceptions.Timeout as e:
-            log2(f"Timeout occurred: {e}")
+        return f"{AuditRepairTables.COMMAND} \t run MSCK REPAIR command for new partition discovery"

kaqing-2.0.105/adam/commands/audit/audit_run.py

@@ -0,0 +1,51 @@
+import concurrent
+from adam.commands.command import Command
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_audits import AuditMeta, Audits
+
+class AuditRun(Command):
+    COMMAND = 'audit run'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(AuditRun, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+        self.auto_repaired = False
+
+    def command(self):
+        return AuditRun.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        state, args = self.apply_state(args, state)
+
+        meta: AuditMeta = Audits.get_meta()
+        clusters = Audits.find_new_clusters(meta.cluster_last_checked)
+        Audits.put_meta(Audits.ADD_CLUSTERS, meta, clusters=clusters)
+        if clusters:
+            log2(f'Added {len(clusters)} new clusters.')
+            tables = Config().get('audit.athena.repair-cluster-tables', 'cluster').split(',')
+            with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
+                for table in tables:
+                    Audits.run_audit_query(f'MSCK REPAIR TABLE {table}')
+        else:
+            log2(f'No new clusters were found.')
+
+        return state
+
+    def completion(self, state: ReplState):
+        if state.device == ReplState.L:
+            return super().completion(state)
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f"{AuditRun.COMMAND} \t run"

kaqing-2.0.105/adam/commands/bash.py

@@ -0,0 +1,150 @@
+from adam.commands.command import Command
+from adam.utils_k8s.cassandra_clusters import CassandraClusters
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.pod_exec_result import PodExecResult
+from adam.repl_state import BashSession, ReplState, RequiredState
+from adam.utils_repl.automata_completer import AutomataCompleter
+from adam.utils_repl.state_machine import StateMachine
+
+class Bash(Command):
+    COMMAND = 'bash'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Bash, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return Bash.COMMAND
+
+    def required(self):
+        return RequiredState.CLUSTER_OR_POD
+
+    def run(self, cmd: str, s0: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, s0)
+
+        state, args = self.apply_state(args, s0, args_to_check=2)
+        if not self.validate_state(state):
+            return state
+
+        if state.in_repl:
+            if s0.sts != state.sts or s0.pod != state.pod:
+                r = self.exec_with_dir(state, args)
+            else:
+                r = self.exec_with_dir(s0, args)
+
+            if not r:
+                state.exit_bash()
+
+                return 'inconsistent pwd'
+
+            return r
+        else:
+            command = ' '.join(args)
+
+            if state.pod:
+                CassandraNodes.exec(state.pod, state.namespace, command, show_out=True)
+            elif state.sts:
+                CassandraClusters.exec(state.sts, state.namespace, command, action='bash', show_out=True)
+
+            return state
+
+    def exec_with_dir(self, state: ReplState, args: list[str]) -> list[PodExecResult]:
+        session_just_created = False
+        if not args:
+            session_just_created = True
+            session = BashSession(state.device)
+            state.enter_bash(session)
+
+        if state.bash_session:
+            if args != ['pwd']:
+                if args:
+                    args.append('&&')
+                args.extend(['pwd', '>', f'/tmp/.qing-{state.bash_session.session_id}'])
+
+            if not session_just_created:
+                if pwd := state.bash_session.pwd(state):
+                    args = ['cd', pwd, '&&'] + args
+
+        command = ' '.join(args)
+
+        rs = []
+
+        if state.pod:
+            rs = [CassandraNodes.exec(state.pod, state.namespace, command,
+                                      show_out=not session_just_created, shell='bash')]
+        elif state.sts:
+            rs = CassandraClusters.exec(state.sts, state.namespace, command, action='bash',
+                                        show_out=not session_just_created, shell='bash')
+
+        return rs
+
+    def completion(self, state: ReplState):
+        if state.pod or state.sts:
+            return {Bash.COMMAND: AutomataCompleter(BashStateMachine())}
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{Bash.COMMAND} [bash-commands]\t run bash on the Cassandra nodes'
+
+BASH_SPEC = [
+    # <command> ::= <simple_command> | <pipeline> | <conditional_command>
+    # <simple_command> ::= <word> <argument>* <redirection>*
+    # <pipeline> ::= <command> '|' <command>
+    # <conditional_command> ::= <command> '&&' <command> | <command> '||' <command>
+    # <word> ::= <letter> <letter_or_digit>*
+    # <argument> ::= <word>
+    # <redirection> ::= '>' <filename> | '<' <filename>
+    # <filename> ::= <word>
+    # <letter> ::= 'a' | 'b' | ... | 'z' | 'A' | 'B' | ... | 'Z'
+    # <digit> ::= '0' | '1' | ... | '9'
+    # <letter_or_digit> ::= <letter> | <digit>
+
+    '                                > word           > word',
+    'word                            > word           > word                                ^ |,>,2>,<,&,&&,||',
+    '-                               > pipe           > word_pipe',
+    '-                               > _rdr0_         > word_rdr0',
+    '-                               > _rdr1_         > word_rdr1',
+    '-                               > _rdr2_         > word_rdr2',
+    '-                               > &              > word_bg                             ^ |,>,2>,<,&,&&,||',
+    '-                               > &&|_or_        > word',
+    'word_a                          > word           > word',
+    'word_pipe                       > word           > word',
+    'word_rdr0                       > word           > word_rdr0_f',
+    'word_rdr1                       > word           > word_rdr1_f',
+    'word_rdr2                       > word           > word_rdr2_f',
+    'word_rdr1_f                     > pipe           > word_pipe                           ^ |,2>,<,&,&&,||',
+    '-                               > _rdr2_         > word_rdr2',
+    '-                               > _rdr0_         > word_rdr0',
+    'word_rdr2_f                     > pipe           > word_pipe                           ^ |,<,&,&&,||',
+    '-                               > _rdr0_         > word_rdr0',
+    '-                               > &              > word_bg                             ^ |,>,2>,<,&,&&,||',
+    '-                               > &&|_or_        > word',
+    'word_rdr0_f                     > pipe           > word_pipe                           ^ |,&,&&,||',
+    '-                               > &              > word_bg                             ^ |,>,2>,<,&,&&,||',
+    '-                               > &&|_or_        > word',
+    'word_bg                         > &&|_or_        >                                     ^ &&,||',
+]
+
+BASH_KEYWORDS = [
+    '&',
+    '&&',
+    '|',
+    '||',
+    '>',
+    '2>',
+    '>>',
+    '<'
+]
+class BashStateMachine(StateMachine[str]):
+    def spec(self) -> str:
+        return BASH_SPEC
+
+    def keywords(self) -> list[str]:
+        return BASH_KEYWORDS

{kaqing-2.0.103 → kaqing-2.0.105}/adam/commands/ls.py

@@ -14,7 +14,7 @@ from adam.pod_exec_result import PodExecResult
 from adam.repl_state import ReplState
 from adam.utils import lines_to_tabular, log, log2
 from adam.apps import Apps
-from adam.utils_athena import audit_table_names
+from adam.utils_audits import Audits
 
 class Ls(Command):
     COMMAND = 'ls'
@@ -131,7 +131,7 @@ class Ls(Command):
         log(lines_to_tabular(pg_table_names(pg.namespace, pg.path()), 'NAME', separator=','))
 
     def show_audit_log_tables(self):
-        log(lines_to_tabular(audit_table_names(), 'NAME', separator=','))
+        log(lines_to_tabular(Audits.audit_table_names(), 'NAME', separator=','))
 
     def completion(self, state: ReplState):
         if state.pod:

{kaqing-2.0.103 → kaqing-2.0.105}/adam/commands/preview_table.py

@@ -4,7 +4,7 @@ from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import lines_to_tabular, log, log2
-from adam.utils_athena import audit_table_names, run_audit_query
+from adam.utils_audits import Audits
 
 class PreviewTable(Command):
     COMMAND = 'preview'
@@ -43,7 +43,7 @@ class PreviewTable(Command):
                     lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresContext.default_schema()]
                     log(lines_to_tabular(lines, separator=','))
                 elif state.device == ReplState.L:
-                    log(lines_to_tabular(audit_table_names(), separator=','))
+                    log(lines_to_tabular(Audits.audit_table_names(), separator=','))
                 else:
                     log(lines_to_tabular(cassandra_table_names(state), separator=','))
 
@@ -66,7 +66,7 @@ class PreviewTable(Command):
         if state.device == ReplState.P:
             PostgresContext.apply(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
         elif state.device == ReplState.L:
-            run_audit_query(f'select * from {table} limit {rows}')
+            Audits.run_audit_query(f'select * from {table} limit {rows}')
         else:
             run_cql(state, f'select * from {table} limit {rows}', show_out=True, use_single_quotes=True, on_any=True)
 

kaqing-2.0.105/adam/embedded_params.py

@@ -0,0 +1,2 @@
+def config():
+    return {'app': {'console-endpoint': 'https://{host}/{env}/{app}/static/console/index.html', 'cr': {'cluster-regex': '(.*?-.*?)-.*', 'group': 'ops.c3.ai', 'v': 'v2', 'plural': 'c3cassandras'}, 'label': 'c3__app_id-0', 'login': {'admin-group': '{host}/C3.ClusterAdmin', 'ingress': '{app_id}-k8singr-appleader-001', 'timeout': 5, 'session-check-url': 'https://{host}/{env}/{app}/api/8/C3/userSessionToken', 'cache-creds': True, 'cache-username': True, 'url': 'https://{host}/{env}/{app}', 'another': "You're logged in to {has}. However, for this app, you need to log in to {need}.", 'token-server-url': 'http://localhost:{port}', 'password-max-length': 128}, 'strip': '0'}, 'audit': {'endpoint': 'https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/', 'workers': 3, 'timeout': 10, 'log-audit-queries': False, 'athena': {'auto-repair': {'elapsed_hours': 12}, 'region': 'us-west-2', 'catalog': 'AwsDataCatalog', 'database': 'audit', 'repair-partition-tables': 'audit', 'output': 's3://s3.ops--audit/ddl/results', 'repair-cluster-tables': 'audit'}}, 'bash': {'workers': 32}, 'cassandra': {'service-name': 'all-pods-service'}, 'cql': {'workers': 32, 'samples': 3, 'secret': {'cluster-regex': '(.*?-.*?)-.*', 'name': '{cluster}-superuser', 'password-item': 'password'}, 'alter-tables': {'excludes': 'system_auth,system_traces,reaper_db,system_distributed,system_views,system,system_schema,system_virtual_schema', 'gc-grace-periods': '3600,86400,864000,7776000', 'batching': True}}, 'checks': {'compactions-threshold': 250, 'cpu-busy-threshold': 98.0, 'cpu-threshold': 0.0, 'cassandra-data-path': '/c3/cassandra', 'root-disk-threshold': 50, 'cassandra-disk-threshold': 50, 'snapshot-size-cmd': "ls /c3/cassandra/data/data/*/*/snapshots | grep snapshots | sed 's/:$//g' | xargs -I {} du -sk {} | awk '{print $1}' | awk '{s+=$1} END {print s}'", 'snapshot-size-threshold': '40G', 'table-sizes-cmd': "ls -Al /c3/cassandra/data/data/ | awk '{print $9}' | sed 's/\\^r//g' | xargs -I {} du -sk /c3/cassandra/data/data/{}"}, 'get-host-id': {'workers': 32}, 'idps': {'ad': {'email-pattern': '.*@c3.ai', 'uri': 'https://login.microsoftonline.com/53ad779a-93e7-485c-ba20-ac8290d7252b/oauth2/v2.0/authorize?response_type=id_token&response_mode=form_post&client_id=00ff94a8-6b0a-4715-98e0-95490012d818&scope=openid+email+profile&redirect_uri=https%3A%2F%2Fplat.c3ci.cloud%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY', 'jwks-uri': 'https://login.microsoftonline.com/common/discovery/keys', 'contact': 'Please contact ted.tran@c3.ai.', 'whitelist-file': '/kaqing/members'}, 'okta': {'default': True, 'email-pattern': '.*@c3iot.com', 'uri': 'https://c3energy.okta.com/oauth2/v1/authorize?response_type=id_token&response_mode=form_post&client_id={client_id}&scope=openid+email+profile+groups&redirect_uri=https%3A%2F%2F{host}%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY', 'jwks-uri': 'https://c3energy.okta.com/oauth2/v1/keys'}}, 'issues': {'workers': 32}, 'logs': {'path': '/c3/cassandra/logs/system.log'}, 'medusa': {'restore-auto-complete': False}, 'nodetool': {'workers': 32, 'samples': 3, 'commands_in_line': 40}, 'pg': {'name-pattern': '^{namespace}.*-k8spg-.*', 'excludes': '.helm., -admin-secret', 'agent': {'name': 'ops-pg-agent', 'just-in-time': False, 'timeout': 86400, 'image': 'seanahnsf/kaqing'}, 'default-db': 'postgres', 'default-schema': 'postgres', 'secret': {'endpoint-key': 'postgres-db-endpoint', 'port-key': 'postgres-db-port', 'username-key': 'postgres-admin-username', 'password-key': 'postgres-admin-password'}}, 'pod': {'name': 'ops', 'image': 'seanahnsf/kaqing-cloud', 'sa': {'name': 'ops', 'proto': 'c3', 'additional-cluster-roles': 'c3aiops-k8ssandra-operator'}, 'label-selector': 'run=ops'}, 'preview': {'rows': 10}, 'processes': {'columns': 'pod,cpu,mem', 'header': 'POD_NAME,CPU,MEM/LIMIT'}, 'reaper': {'service-name': 'reaper-service', 'port-forward': {'timeout': 86400, 'local-port': 9001}, 'abort-runs-batch': 10, 'show-runs-batch': 100, 'pod': {'cluster-regex': '(.*?-.*?-.*?-.*?)-.*', 'label-selector': 'k8ssandra.io/reaper={cluster}-reaper'}, 'secret': {'cluster-regex': '(.*?-.*?)-.*', 'name': '{cluster}-reaper-ui', 'password-item': 'password'}}, 'repair': {'log-path': '/home/cassrepair/logs/', 'image': 'ci-registry.c3iot.io/cloudops/cassrepair:2.0.14', 'secret': 'ciregistryc3iotio', 'env': {'interval': 24, 'timeout': 60, 'pr': False, 'runs': 1}}, 'repl': {'start-drive': 'a', 'a': {'auto-enter': 'c3/c3'}, 'c': {'auto-enter': 'cluster'}, 'history': {'push-cat-remote-log-file': True}, 'background-process': {'auto-nohup': True}}, 'status': {'columns': 'status,address,load,tokens,owns,host_id,gossip,compactions', 'header': '--,Address,Load,Tokens,Owns,Host ID,GOSSIP,COMPACTIONS'}, 'storage': {'columns': 'pod,volume_root,volume_cassandra,snapshots,data,compactions', 'header': 'POD_NAME,VOLUME /,VOLUME CASS,SNAPSHOTS,DATA,COMPACTIONS'}, 'watch': {'auto': 'rollout', 'timeout': 3600, 'interval': 10}, 'debug': False, 'debugs': {'timings': False, 'exit-on-error': False, 'show-parallelism': False}}

{kaqing-2.0.103 → kaqing-2.0.105}/adam/repl_commands.py

@@ -3,6 +3,7 @@ from adam.commands.app import App
 from adam.commands.app_ping import AppPing
 from adam.commands.audit.audit import Audit
 from adam.commands.audit.audit_repair_tables import AuditRepairTables
+from adam.commands.audit.audit_run import AuditRun
 from adam.commands.deploy.code_start import CodeStart
 from adam.commands.deploy.code_stop import CodeStop
 from adam.commands.deploy.deploy import Deploy
@@ -57,7 +58,7 @@ class ReplCommands:
         cmds: list[Command] = ReplCommands.navigation() + ReplCommands.cassandra_check() + ReplCommands.cassandra_ops() + \
             ReplCommands.tools() + ReplCommands.app() + ReplCommands.exit()
 
-        intermediate_cmds: list[Command] = [App(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
+        intermediate_cmds: list[Command] = [App(), Audit(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
         ic = [c.command() for c in intermediate_cmds]
         # 1. dedup commands
         deduped = []
@@ -85,7 +86,7 @@ class ReplCommands:
 
     def tools() -> list[Command]:
         return [Cqlsh(), Postgres(), Bash(), Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(),
-                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), AuditRepairTables(), Audit()]
+                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), Audit()] + Audit.cmd_list()
 
     def app() -> list[Command]:
         return [ShowAppActions(), ShowAppId(), ShowAppQueues(), AppPing(), App()]

{kaqing-2.0.103 → kaqing-2.0.105}/adam/sql/sql_completer.py

@@ -1,16 +1,33 @@
 from typing import Callable
 
-from adam.sql.automata_completer import AutomataCompleter
+import sqlparse
+from sqlparse.sql import Statement, 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
 
 __all__ = [
     "SqlCompleter",
 ]
 
-def default_columns(tables: list[str]):
+def default_columns(_: list[str]):
     return 'id,x.,y.,z.'.split(',')
 
-class SqlCompleter(AutomataCompleter):
+class SqlCompleter(AutomataCompleter[Token]):
+    def tokens(self, text: str) -> list[Token]:
+        tokens = []
+
+        stmts = sqlparse.parse(text)
+        if not stmts:
+            tokens = []
+        else:
+            statement: Statement = stmts[0]
+            tokens = statement.tokens
+
+        return tokens
+
     def __init__(self,
                  tables: Callable[[], list[str]],
                  dml: str = None,
@@ -33,7 +50,17 @@ class SqlCompleter(AutomataCompleter):
         self.variant = variant
         self.debug = debug
 
-    def terms(self, state: str, word: str) -> list[str]:
+    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 == 'tables':

{kaqing-2.0.103 → kaqing-2.0.105}/adam/sql/sql_state_machine.py

@@ -1,4 +1,8 @@
-from adam.sql.state_machine import StateMachine
+from typing import Callable
+from sqlparse.sql import Token
+from sqlparse import tokens as TOKEN
+
+from adam.utils_repl.state_machine import StateMachine, State
 
 __all__ = [
     'SqlStateMachine', 'CqlStateMachine', 'AthenaStateMachine'
@@ -70,7 +74,7 @@ SQL_SPEC = [
     'select_from_lp_                 > select           > select',
     'select_from_x                   > ,                > select_from_x_comma_                    ^ (select,tables',
     'select_from_sq_                 > as               > select_from_x_as                        ^ as',
-    'select_from_x_comma_            > name             > select_from_x                           ^ tables',
+    'select_from_x_comma_            > name|audit       > select_from_x                           ^ tables',
     'select_from_x_                                                                               ^ as,where,inner join,left outer join,right outer join,full outer join,group by,order by,limit',
     'select_from_x_as_x_             > ,                > select_from_x_comma_                    ^ where,inner join,left outer join,right outer join,full outer join,group by,order by,limit',
     '-                               > as               > select_from_x_as',
@@ -92,7 +96,7 @@ SQL_SPEC = [
     '-                               > full outer join  > select_join',
     'select_from_x_as_               > name             > select_from_x_as_x                      ^ x,y,z',
     'select_from_x_as_x              > ,                > select_from_x_as_x_comma_',
-    'select_from_x_as_x_comma_       > name             > select_from_x                           ^ tables',
+    'select_from_x_as_x_comma_       > name|audit       > select_from_x                           ^ tables',
     'select_where_                   > name             > select_where_a                          ^ columns',
     'select_where_a                  > name             > select_where_a                          ^ columns,=,<,<=,>,>=,<>',
     '-                               > comparison       > select_where_a_op',
@@ -134,7 +138,7 @@ SQL_SPEC = [
     'select_where_sc_limit_          > num              > select_where_sc_limit_num               ^ 1',
     'select_where_sc_limit_num_rp__  > as               > select_from_x_as                        ^ as',
     'select_where_x_inner_           > join             > select_join',
-    'select_join_                    > name             > select_x_join_y                         ^ tables',
+    'select_join_                    > name|audit       > select_x_join_y                         ^ tables',
     'select_from_x_left_             > join             > select_join                             ^ outer join',
     '-                               > outer            > select_from_x_left_outer',
     'select_from_x_left_outer_       > join             > select_join                             ^ join',
@@ -171,7 +175,7 @@ SQL_SPEC = [
     # <query_expression> ::= SELECT <select_list> FROM <table_reference_list> [ WHERE <search_condition> ] [ GROUP BY <grouping_column_list> ] [ HAVING <search_condition> ] [ ORDER BY <sort_specification_list> ]
     '                                > insert           > insert',
     'insert_                         > into             > insert_into                             ^ into',
-    'insert_into_                    > name             > insert_into_x                           ^ tables',
+    'insert_into_                    > name|audit       > insert_into_x                           ^ tables',
     'insert_into_x                   > (                > insert_into_x_lp_',
     'insert_into_x_                  > (                > insert_into_x_lp_                       ^ (,values(',
     '-                               > values           > insert_values',
@@ -198,7 +202,7 @@ SQL_SPEC = [
 
     # <search_condition> ::= <boolean_expression>
     '                                > update           > update',
-    'update_                         > name             > update_x                                ^ tables',
+    'update_                         > name|audit       > update_x                                ^ tables',
     'update_x_                       > set              > update_set                              ^ set',
     'update_set_                     > name             > update_set_a                            ^ id',
     'update_set_a                    > comparison       > update_set_a_op',
@@ -267,7 +271,7 @@ SQL_SPEC = [
     # <expression_list> ::= <expression> { , <expression> }...
     '                                > delete           > delete',
     'delete_                         > from             > delete_from                             ^ from',
-    'delete_from_                    > name             > delete_from_x                           ^ tables',
+    'delete_from_                    > name|audit       > delete_from_x                           ^ tables',
     'delete_from_x_                  > where            > update_where                            ^ where',
 
     # <alter table action> ::=
@@ -381,10 +385,90 @@ ATHENA_KEYWORDS = SQL_KEYWORDS + [
     'partition'
 ]
 
-class SqlStateMachine(StateMachine):
+class SqlStateMachine(StateMachine[Token]):
     def __init__(self, indent=0, push_level = 0, debug = False):
         super().__init__(indent, push_level, debug)
 
+    def traverse_tokens(self, tokens: list[Token], state: State = State('')):
+        def handle_push():
+            if f'{state.state} > {it}' in self.states:
+                state_test = self.states[f'{state.state} > {it}']
+                if state_test.comeback_token:
+                    self.comebacks[self.push_level] = state_test.comeback_state
+
+        def handle_pop():
+            if self.push_level in self.comebacks:
+                try:
+                    return State(self.comebacks[self.push_level])
+                finally:
+                    del self.comebacks[self.push_level]
+
+            return None
+
+        for token in tokens:
+            if self.debug:
+                if token.ttype == TOKEN.Whitespace:
+                    print('_ ', end='')
+                elif token.ttype in [TOKEN.DML, TOKEN.Wildcard, TOKEN.Punctuation, TOKEN.CTE]:
+                    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("  " * self.indent + f"Token: {token.value}, Type: {token.ttype}@{token.ttype.__class__}")
+
+            last_name = None
+            if token.is_group:
+                state = self.traverse_tokens(token.tokens, state)
+            else:
+                comeback_state = None
+
+                it = ''
+                if (t := token.value.lower()) in self.keywords():
+                    it = t
+                elif token.ttype == TOKEN.Text.Whitespace:
+                    it = '_'
+                elif token.ttype == TOKEN.Name:
+                    it = 'name'
+                    last_name = token.value
+                elif token.ttype == TOKEN.Literal.String.Single:
+                    it = 'single'
+                elif token.ttype in [TOKEN.Literal.Number.Integer, TOKEN.Literal.Number.Float]:
+                    it = 'num'
+                elif token.ttype == TOKEN.Wildcard:
+                    it = '*'
+                elif token.ttype == TOKEN.Punctuation:
+                    it = token.value
+
+                    if it == '(':
+                        handle_push()
+                        self.push_level += 1
+                    elif it == ')':
+                        self.push_level -= 1
+                        comeback_state = handle_pop()
+
+                elif token.ttype == TOKEN.Operator.Comparison:
+                    it = 'comparison'
+
+                try:
+                    # print(f'\n{state.to_s} > {it} > ', end='')
+                    if comeback_state:
+                        state = comeback_state
+                    else:
+                        context = state.context
+                        state = self.states[f'{state.state} > {it}']
+                        state.context = context
+
+                    if last_name:
+                        state.context['last_name'] = last_name
+                except:
+                    pass
+
+        return state
+
     def spec(self):
         return SQL_SPEC