kaqing 2.0.101__py3-none-any.whl → 2.0.104__py3-none-any.whl

adam/commands/postgres/postgres_utils.py

@@ -1,6 +1,6 @@
 import functools
 
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 
 TestPG = [False]
@@ -12,8 +12,8 @@ def pg_database_names(ns: str, pg_path: str):
 
     Config().wait_log('Inspecting Postgres Databases...')
 
-    pg = PostgresSession(ns, pg_path)
-    return [db['name'] for db in pg.databases() if db['owner'] == PostgresSession.default_owner()]
+    pg = PostgresContext.apply(ns, pg_path)
+    return [db['name'] for db in pg.databases() if db['owner'] == PostgresContext.default_owner()]
 
 @functools.lru_cache()
 def pg_table_names(ns: str, pg_path: str):
@@ -21,10 +21,10 @@ def pg_table_names(ns: str, pg_path: str):
         return ['C3_2_XYZ1']
 
     Config().wait_log('Inspecting Postgres Database...')
-    return [table['name'] for table in pg_tables(ns, pg_path) if table['schema'] == PostgresSession.default_schema()]
+    return [table['name'] for table in pg_tables(ns, pg_path) if table['schema'] == PostgresContext.default_schema()]
 
 def pg_tables(ns: str, pg_path: str):
-    pg = PostgresSession(ns, pg_path)
+    pg = PostgresContext.apply(ns, pg_path)
     if pg.db:
         return pg.tables()
 

adam/commands/postgres/psql_completions.py

@@ -7,4 +7,4 @@ def psql_completions(ns: str, pg_path: str):
         '\d': None,
         '\dt': None,
         '\du': None
-    } | SqlCompleter.completions(lambda: pg_table_names(ns, pg_path))
+    } | SqlCompleter(lambda: pg_table_names(ns, pg_path)).completions_for_nesting()

adam/commands/preview_table.py

@@ -1,11 +1,6 @@
-import functools
-
-from adam.commands.audit.audit_table_completer import AuditTableNameCompleter
 from adam.commands.command import Command
-from adam.commands.cql.cql_table_completer import CqlTableNameCompleter
-from adam.commands.cql.cql_utils import run_cql, table_names, tables
-from adam.commands.postgres.postgres_session import PostgresSession
-from adam.commands.postgres.psql_table_completer import PsqlTableNameCompleter
+from adam.commands.cql.cql_utils import cassandra_table_names, run_cql
+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
@@ -44,13 +39,13 @@ class PreviewTable(Command):
         if not args:
             def show_tables():
                 if state.device == ReplState.P:
-                    pg = PostgresSession(state.namespace, state.pg_path)
-                    lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresSession.default_schema()]
+                    pg = PostgresContext.apply(state.namespace, state.pg_path)
+                    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=','))
                 else:
-                    run_cql(state, f'describe tables', show_out=True, on_any=True)
+                    log(lines_to_tabular(cassandra_table_names(state), separator=','))
 
             if state.in_repl:
                 log2('Table is required.')
@@ -69,7 +64,7 @@ class PreviewTable(Command):
 
         rows = Config().get('preview.rows', 10)
         if state.device == ReplState.P:
-            PostgresSession(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
+            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}')
         else:
@@ -77,22 +72,8 @@ class PreviewTable(Command):
 
         return state
 
-    def completion(self, state: ReplState):
-        if state.device == ReplState.P:
-            return {PreviewTable.COMMAND: PsqlTableNameCompleter(state.namespace, state.pg_path)}
-        elif state.device == ReplState.L:
-            return {PreviewTable.COMMAND: AuditTableNameCompleter()}
-        elif state.sts:
-            return {PreviewTable.COMMAND: CqlTableNameCompleter(table_names(state))}
-
+    def completion(self, _: ReplState):
         return {}
 
     def help(self, _: ReplState):
-        return f'{PreviewTable.COMMAND} TABLE\t preview table'
-
-    @functools.lru_cache()
-    def cql_tables(state: ReplState):
-        if state.pod:
-            return tables(state)
-
-        return tables(state, on_any=True)
+        return f'{PreviewTable.COMMAND} TABLE\t preview table'

adam/commands/pwd.py

@@ -1,6 +1,6 @@
 from adam.app_session import AppSession
 from adam.commands.command import Command
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.repl_state import ReplState
 from adam.utils import lines_to_tabular, log
 
@@ -29,7 +29,7 @@ class Pwd(Command):
             words = []
 
             if device == ReplState.P:
-                pg = PostgresSession(state.namespace, state.pg_path)
+                pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
 
                 if pg.host:
                     words.append(f'host/{pg.host}')

adam/embedded_params.py

@@ -1,2 +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', 'tables': 'audit', 'output': 's3://s3.ops--audit/ddl/results'}}, '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}}
+    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}}

adam/repl.py

@@ -13,7 +13,7 @@ from adam.cli_group import cli
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterCommandHelper
 from adam.commands.help import Help
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 from adam.utils_k8s.kube_context import KubeContext
 from adam.utils_k8s.statefulsets import StatefulSets
@@ -40,7 +40,7 @@ def enter_repl(state: ReplState):
         msg = ''
         if state.device == ReplState.P:
             msg = f'{ReplState.P}:'
-            pg = PostgresSession(state.namespace, state.pg_path) if state.pg_path else None
+            pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path) if state.pg_path else None
             if pg and pg.db:
                 msg += pg.db
             elif pg and pg.host:
@@ -141,7 +141,7 @@ def enter_repl(state: ReplState):
                     cmd = f'bash {cmd}'
 
                 if cmd and cmd.strip(' ') and not cmds.run(cmd, state):
-                    try_device_default_action(state, cmds, cmd_list)
+                    try_device_default_action(state, cmds, cmd_list, cmd)
                     # not served by any command in the chain; try SQL query or C3 action
                     # c_sql_tried = False
                     # if state.device == ReplState.P:
@@ -186,10 +186,10 @@ def enter_repl(state: ReplState):
                 if cmd and (state.device != ReplState.L or Config().get('audit.log-audit-queries', False)):
                     executor.submit(audit_log, cmd, state)
 
-def try_device_default_action(state: ReplState, cmds: Command, cmd_list: list[Command]):
+def try_device_default_action(state: ReplState, cmds: Command, cmd_list: list[Command], cmd: str):
     c_sql_tried = False
     if state.device == ReplState.P:
-        pg = PostgresSession(state.namespace, state.pg_path)
+        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
         if pg.db:
             c_sql_tried = True
             cmd = f'pg {cmd}'

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
@@ -13,7 +14,6 @@ from adam.commands.deploy.undeploy import Undeploy
 from adam.commands.deploy.undeploy_frontend import UndeployFrontend
 from adam.commands.deploy.undeploy_pg_agent import UndeployPgAgent
 from adam.commands.deploy.undeploy_pod import UndeployPod
-from adam.commands.describe.describe import Describe
 from adam.commands.shell import Shell
 from adam.commands.show.show_app_queues import ShowAppQueues
 from adam.commands.cp import ClipboardCopy
@@ -58,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(), Describe(), Show(), Undeploy()]
+        intermediate_cmds: list[Command] = [App(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
         ic = [c.command() for c in intermediate_cmds]
         # 1. dedup commands
         deduped = []
@@ -79,15 +79,14 @@ class ReplCommands:
                 GetParam(), SetParam(), ShowParams(), ShowKubectlCommands(), ShowLogin(), ShowAdam(), ShowHost()]
 
     def cassandra_check() -> list[Command]:
-        return Describe.cmd_list() + [ShowCassandraStatus(),
-                ShowCassandraVersion(), ShowRepairs(), ShowStorage(), ShowProcesses(), Check(), Issues(), NodeTool(), Report()]
+        return [ShowCassandraStatus(), ShowCassandraVersion(), ShowRepairs(), ShowStorage(), ShowProcesses(), Check(), Issues(), NodeTool(), Report()]
 
     def cassandra_ops() -> list[Command]:
         return [AlterTables()] + Medusa.cmd_list() + [Restart(), RollOut(), Watch()] + Reaper.cmd_list() + Repair.cmd_list()
 
     def tools() -> list[Command]:
         return [Cqlsh(), Postgres(), Bash(), Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(),
-                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), AuditRepairTables(), Audit()]
+                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), AuditRepairTables(), AuditRun(), Audit()]
 
     def app() -> list[Command]:
         return [ShowAppActions(), ShowAppId(), ShowAppQueues(), AppPing(), App()]

adam/repl_state.py

@@ -2,7 +2,7 @@ import copy
 from enum import Enum
 import re
 
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.utils_k8s.cassandra_clusters import CassandraClusters
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.kube_context import KubeContext
@@ -210,7 +210,7 @@ class ReplState:
                     return False
 
         if pg_required == RequiredState.PG_DATABASE:
-            pg = PostgresSession(self.namespace, self.pg_path)
+            pg: PostgresContext = PostgresContext.apply(self.namespace, self.pg_path)
             if not pg.db:
                 if self.in_repl:
                     log2('cd to a database first.')

adam/sql/sql_completer.py

@@ -1,22 +1,33 @@
-from typing import Callable, Iterable
-from prompt_toolkit.completion import CompleteEvent, Completer, Completion
-from prompt_toolkit.document import Document
+from typing import Callable
+
 import sqlparse
-from sqlparse.sql import Statement
+from sqlparse.sql import Statement, 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', 'alter', 'describe'])
-
-def default_columns(tables: list[str]):
+def default_columns(_: list[str]):
     return 'id,x.,y.,z.'.split(',')
 
-class SqlCompleter(Completer):
+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,
@@ -25,78 +36,58 @@ class SqlCompleter(Completer):
                  table_props: Callable[[], dict[str,list[str]]] = lambda: [],
                  variant = 'sql',
                  debug = False):
-        super().__init__()
-        self.dml = dml
+        machine = SqlStateMachine(debug=debug)
+        if variant == 'cql':
+            machine = CqlStateMachine(debug=debug)
+        elif variant == 'athena':
+            machine = AthenaStateMachine(debug=debug)
+        super().__init__(machine, dml, debug)
+
         self.tables = tables
         self.columns = columns
         self.partition_columns = partition_columns
         self.table_props = table_props
+        self.variant = variant
         self.debug = debug
-        self.machine = StateMachine(variant=variant, 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 == 'tables':
+            terms.extend(self.tables())
+        elif word == '`tables`':
+            terms.append('tables')
+        elif word == 'columns':
+            terms.extend(self.columns([]))
+        elif word == 'partition-columns':
+            terms.extend(self.partition_columns([]))
+        elif word == 'table-props':
+            terms.extend(self.table_props().keys())
+        elif word == 'table-prop-values':
+            if 'last_name' in state.context and state.context['last_name']:
+                terms.extend(self.table_props()[state.context['last_name']])
+        elif word == 'single':
+            terms.append("'")
+        elif word == 'comma':
+            terms.append(",")
         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 == '`tables`':
-                        terms.append('tables')
-                    elif word == 'columns':
-                        terms.extend(self.columns([]))
-                    elif word == 'partition-columns':
-                        terms.extend(self.partition_columns([]))
-                    elif word == 'table-props':
-                        terms.extend(self.table_props().keys())
-                    elif word == 'table-prop-values':
-                        if 'last_name' in state.context and state.context['last_name']:
-                            terms.extend(self.table_props()[state.context['last_name']])
-                    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,
-                    partition_columns: Callable[[list[str]], list[str]] = lambda x: [],
-                    table_props: Callable[[], dict[str, list[str]]] = lambda: [],
-                    variant = 'sql'):
+            terms.append(word)
+
+        return terms
+
+    def completions_for_nesting(self):
         return {
-            'delete': SqlCompleter(table_names, 'delete', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
-            'insert': SqlCompleter(table_names, 'insert', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
-            'select': SqlCompleter(table_names, 'select', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
-            'update': SqlCompleter(table_names, 'update', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
-            'alter': SqlCompleter(table_names, 'alter', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
-            'describe': SqlCompleter(table_names, 'describe', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
+            word : SqlCompleter(self.tables, word, columns=self.columns, partition_columns=self.partition_columns,
+                                table_props=self.table_props, variant=self.variant)
+            for word in self.machine.suggestions[''].strip(' ').split(',')
         }