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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of kaqing might be problematic. Click here for more details.

Files changed (46) hide show
  1. adam/batch.py +0 -14
  2. adam/commands/audit/audit.py +4 -8
  3. adam/commands/audit/audit_repair_tables.py +13 -35
  4. adam/commands/audit/audit_run.py +49 -0
  5. adam/commands/bash.py +60 -2
  6. adam/commands/cd.py +8 -8
  7. adam/commands/commands_utils.py +1 -2
  8. adam/commands/cql/cql_completions.py +4 -4
  9. adam/commands/cql/cql_utils.py +1 -1
  10. adam/commands/cql/cqlsh.py +1 -1
  11. adam/commands/deploy/deploy_pg_agent.py +2 -2
  12. adam/commands/deploy/undeploy_pg_agent.py +2 -2
  13. adam/commands/ls.py +12 -12
  14. adam/commands/nodetool.py +1 -1
  15. adam/commands/postgres/postgres.py +3 -3
  16. adam/commands/postgres/{postgres_session.py → postgres_context.py} +26 -27
  17. adam/commands/postgres/postgres_utils.py +5 -5
  18. adam/commands/postgres/psql_completions.py +1 -1
  19. adam/commands/preview_table.py +8 -27
  20. adam/commands/pwd.py +2 -2
  21. adam/embedded_params.py +1 -1
  22. adam/repl.py +3 -3
  23. adam/repl_commands.py +3 -6
  24. adam/repl_state.py +2 -2
  25. adam/sql/sql_completer.py +67 -76
  26. adam/sql/sql_state_machine.py +518 -0
  27. adam/sql/term_completer.py +3 -0
  28. adam/utils_athena.py +68 -1
  29. adam/utils_k8s/pods.py +2 -2
  30. adam/version.py +1 -1
  31. {kaqing-2.0.102.dist-info → kaqing-2.0.104.dist-info}/METADATA +1 -1
  32. {kaqing-2.0.102.dist-info → kaqing-2.0.104.dist-info}/RECORD +35 -44
  33. adam/commands/audit/audit_table_completer.py +0 -9
  34. adam/commands/cql/cql_table_completer.py +0 -8
  35. adam/commands/describe/__init__.py +0 -0
  36. adam/commands/describe/describe.py +0 -61
  37. adam/commands/describe/describe_keyspace.py +0 -58
  38. adam/commands/describe/describe_keyspaces.py +0 -46
  39. adam/commands/describe/describe_schema.py +0 -46
  40. adam/commands/describe/describe_table.py +0 -57
  41. adam/commands/describe/describe_tables.py +0 -46
  42. adam/commands/postgres/psql_table_completer.py +0 -11
  43. adam/sql/state_machine.py +0 -576
  44. {kaqing-2.0.102.dist-info → kaqing-2.0.104.dist-info}/WHEEL +0 -0
  45. {kaqing-2.0.102.dist-info → kaqing-2.0.104.dist-info}/entry_points.txt +0 -0
  46. {kaqing-2.0.102.dist-info → kaqing-2.0.104.dist-info}/top_level.txt +0 -0
@@ -1,11 +1,6 @@
1
- import functools
2
-
3
- from adam.commands.audit.audit_table_completer import AuditTableNameCompleter
4
1
  from adam.commands.command import Command
5
- from adam.commands.cql.cql_table_completer import CqlTableNameCompleter
6
- from adam.commands.cql.cql_utils import run_cql, table_names, tables
7
- from adam.commands.postgres.postgres_session import PostgresSession
8
- from adam.commands.postgres.psql_table_completer import PsqlTableNameCompleter
2
+ from adam.commands.cql.cql_utils import cassandra_table_names, run_cql
3
+ from adam.commands.postgres.postgres_context import PostgresContext
9
4
  from adam.config import Config
10
5
  from adam.repl_state import ReplState, RequiredState
11
6
  from adam.utils import lines_to_tabular, log, log2
@@ -44,13 +39,13 @@ class PreviewTable(Command):
44
39
  if not args:
45
40
  def show_tables():
46
41
  if state.device == ReplState.P:
47
- pg = PostgresSession(state.namespace, state.pg_path)
48
- lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresSession.default_schema()]
42
+ pg = PostgresContext.apply(state.namespace, state.pg_path)
43
+ lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresContext.default_schema()]
49
44
  log(lines_to_tabular(lines, separator=','))
50
45
  elif state.device == ReplState.L:
51
46
  log(lines_to_tabular(audit_table_names(), separator=','))
52
47
  else:
53
- run_cql(state, f'describe tables', show_out=True, on_any=True)
48
+ log(lines_to_tabular(cassandra_table_names(state), separator=','))
54
49
 
55
50
  if state.in_repl:
56
51
  log2('Table is required.')
@@ -69,7 +64,7 @@ class PreviewTable(Command):
69
64
 
70
65
  rows = Config().get('preview.rows', 10)
71
66
  if state.device == ReplState.P:
72
- PostgresSession(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
67
+ PostgresContext.apply(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
73
68
  elif state.device == ReplState.L:
74
69
  run_audit_query(f'select * from {table} limit {rows}')
75
70
  else:
@@ -77,22 +72,8 @@ class PreviewTable(Command):
77
72
 
78
73
  return state
79
74
 
80
- def completion(self, state: ReplState):
81
- if state.device == ReplState.P:
82
- return {PreviewTable.COMMAND: PsqlTableNameCompleter(state.namespace, state.pg_path)}
83
- elif state.device == ReplState.L:
84
- return {PreviewTable.COMMAND: AuditTableNameCompleter()}
85
- elif state.sts:
86
- return {PreviewTable.COMMAND: CqlTableNameCompleter(table_names(state))}
87
-
75
+ def completion(self, _: ReplState):
88
76
  return {}
89
77
 
90
78
  def help(self, _: ReplState):
91
- return f'{PreviewTable.COMMAND} TABLE\t preview table'
92
-
93
- @functools.lru_cache()
94
- def cql_tables(state: ReplState):
95
- if state.pod:
96
- return tables(state)
97
-
98
- return tables(state, on_any=True)
79
+ return f'{PreviewTable.COMMAND} TABLE\t preview table'
adam/commands/pwd.py CHANGED
@@ -1,6 +1,6 @@
1
1
  from adam.app_session import AppSession
2
2
  from adam.commands.command import Command
3
- from adam.commands.postgres.postgres_session import PostgresSession
3
+ from adam.commands.postgres.postgres_context import PostgresContext
4
4
  from adam.repl_state import ReplState
5
5
  from adam.utils import lines_to_tabular, log
6
6
 
@@ -29,7 +29,7 @@ class Pwd(Command):
29
29
  words = []
30
30
 
31
31
  if device == ReplState.P:
32
- pg = PostgresSession(state.namespace, state.pg_path)
32
+ pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
33
33
 
34
34
  if pg.host:
35
35
  words.append(f'host/{pg.host}')
adam/embedded_params.py CHANGED
@@ -1,2 +1,2 @@
1
1
  def config():
2
- 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}}
2
+ 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 CHANGED
@@ -13,7 +13,7 @@ from adam.cli_group import cli
13
13
  from adam.commands.command import Command
14
14
  from adam.commands.command_helpers import ClusterCommandHelper
15
15
  from adam.commands.help import Help
16
- from adam.commands.postgres.postgres_session import PostgresSession
16
+ from adam.commands.postgres.postgres_context import PostgresContext
17
17
  from adam.config import Config
18
18
  from adam.utils_k8s.kube_context import KubeContext
19
19
  from adam.utils_k8s.statefulsets import StatefulSets
@@ -40,7 +40,7 @@ def enter_repl(state: ReplState):
40
40
  msg = ''
41
41
  if state.device == ReplState.P:
42
42
  msg = f'{ReplState.P}:'
43
- pg = PostgresSession(state.namespace, state.pg_path) if state.pg_path else None
43
+ pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path) if state.pg_path else None
44
44
  if pg and pg.db:
45
45
  msg += pg.db
46
46
  elif pg and pg.host:
@@ -189,7 +189,7 @@ def enter_repl(state: ReplState):
189
189
  def try_device_default_action(state: ReplState, cmds: Command, cmd_list: list[Command], cmd: str):
190
190
  c_sql_tried = False
191
191
  if state.device == ReplState.P:
192
- pg = PostgresSession(state.namespace, state.pg_path)
192
+ pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
193
193
  if pg.db:
194
194
  c_sql_tried = True
195
195
  cmd = f'pg {cmd}'
adam/repl_commands.py CHANGED
@@ -3,6 +3,7 @@ from adam.commands.app import App
3
3
  from adam.commands.app_ping import AppPing
4
4
  from adam.commands.audit.audit import Audit
5
5
  from adam.commands.audit.audit_repair_tables import AuditRepairTables
6
+ from adam.commands.audit.audit_run import AuditRun
6
7
  from adam.commands.deploy.code_start import CodeStart
7
8
  from adam.commands.deploy.code_stop import CodeStop
8
9
  from adam.commands.deploy.deploy import Deploy
@@ -13,7 +14,6 @@ from adam.commands.deploy.undeploy import Undeploy
13
14
  from adam.commands.deploy.undeploy_frontend import UndeployFrontend
14
15
  from adam.commands.deploy.undeploy_pg_agent import UndeployPgAgent
15
16
  from adam.commands.deploy.undeploy_pod import UndeployPod
16
- from adam.commands.describe.describe import Describe
17
17
  from adam.commands.shell import Shell
18
18
  from adam.commands.show.show_app_queues import ShowAppQueues
19
19
  from adam.commands.cp import ClipboardCopy
@@ -59,7 +59,6 @@ class ReplCommands:
59
59
  ReplCommands.tools() + ReplCommands.app() + ReplCommands.exit()
60
60
 
61
61
  intermediate_cmds: list[Command] = [App(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
62
- # intermediate_cmds: list[Command] = [App(), Reaper(), Repair(), Deploy(), Describe(), Show(), Undeploy()]
63
62
  ic = [c.command() for c in intermediate_cmds]
64
63
  # 1. dedup commands
65
64
  deduped = []
@@ -80,16 +79,14 @@ class ReplCommands:
80
79
  GetParam(), SetParam(), ShowParams(), ShowKubectlCommands(), ShowLogin(), ShowAdam(), ShowHost()]
81
80
 
82
81
  def cassandra_check() -> list[Command]:
83
- # return Describe.cmd_list() + [ShowCassandraStatus(),
84
- return [ShowCassandraStatus(),
85
- ShowCassandraVersion(), ShowRepairs(), ShowStorage(), ShowProcesses(), Check(), Issues(), NodeTool(), Report()]
82
+ return [ShowCassandraStatus(), ShowCassandraVersion(), ShowRepairs(), ShowStorage(), ShowProcesses(), Check(), Issues(), NodeTool(), Report()]
86
83
 
87
84
  def cassandra_ops() -> list[Command]:
88
85
  return [AlterTables()] + Medusa.cmd_list() + [Restart(), RollOut(), Watch()] + Reaper.cmd_list() + Repair.cmd_list()
89
86
 
90
87
  def tools() -> list[Command]:
91
88
  return [Cqlsh(), Postgres(), Bash(), Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(),
92
- DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), AuditRepairTables(), Audit()]
89
+ DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), AuditRepairTables(), AuditRun(), Audit()]
93
90
 
94
91
  def app() -> list[Command]:
95
92
  return [ShowAppActions(), ShowAppId(), ShowAppQueues(), AppPing(), App()]
adam/repl_state.py CHANGED
@@ -2,7 +2,7 @@ import copy
2
2
  from enum import Enum
3
3
  import re
4
4
 
5
- from adam.commands.postgres.postgres_session import PostgresSession
5
+ from adam.commands.postgres.postgres_context import PostgresContext
6
6
  from adam.utils_k8s.cassandra_clusters import CassandraClusters
7
7
  from adam.utils_k8s.cassandra_nodes import CassandraNodes
8
8
  from adam.utils_k8s.kube_context import KubeContext
@@ -210,7 +210,7 @@ class ReplState:
210
210
  return False
211
211
 
212
212
  if pg_required == RequiredState.PG_DATABASE:
213
- pg = PostgresSession(self.namespace, self.pg_path)
213
+ pg: PostgresContext = PostgresContext.apply(self.namespace, self.pg_path)
214
214
  if not pg.db:
215
215
  if self.in_repl:
216
216
  log2('cd to a database first.')
adam/sql/sql_completer.py CHANGED
@@ -1,22 +1,33 @@
1
- from typing import Callable, Iterable
2
- from prompt_toolkit.completion import CompleteEvent, Completer, Completion
3
- from prompt_toolkit.document import Document
1
+ from typing import Callable
2
+
4
3
  import sqlparse
5
- from sqlparse.sql import Statement
4
+ from sqlparse.sql import Statement, Token
6
5
 
7
- from adam.sql.state_machine import StateMachine, StateTo
8
6
  from adam.sql.term_completer import TermCompleter
7
+ from adam.utils_repl.automata_completer import AutomataCompleter
8
+ from adam.sql.sql_state_machine import AthenaStateMachine, CqlStateMachine, SqlStateMachine
9
+ from adam.utils_repl.state_machine import State
9
10
 
10
11
  __all__ = [
11
12
  "SqlCompleter",
12
13
  ]
13
14
 
14
- DML_COMPLETER = TermCompleter(['select', 'insert', 'delete', 'update', 'alter', 'describe'])
15
-
16
- def default_columns(tables: list[str]):
15
+ def default_columns(_: list[str]):
17
16
  return 'id,x.,y.,z.'.split(',')
18
17
 
19
- class SqlCompleter(Completer):
18
+ class SqlCompleter(AutomataCompleter[Token]):
19
+ def tokens(self, text: str) -> list[Token]:
20
+ tokens = []
21
+
22
+ stmts = sqlparse.parse(text)
23
+ if not stmts:
24
+ tokens = []
25
+ else:
26
+ statement: Statement = stmts[0]
27
+ tokens = statement.tokens
28
+
29
+ return tokens
30
+
20
31
  def __init__(self,
21
32
  tables: Callable[[], list[str]],
22
33
  dml: str = None,
@@ -25,78 +36,58 @@ class SqlCompleter(Completer):
25
36
  table_props: Callable[[], dict[str,list[str]]] = lambda: [],
26
37
  variant = 'sql',
27
38
  debug = False):
28
- super().__init__()
29
- self.dml = dml
39
+ machine = SqlStateMachine(debug=debug)
40
+ if variant == 'cql':
41
+ machine = CqlStateMachine(debug=debug)
42
+ elif variant == 'athena':
43
+ machine = AthenaStateMachine(debug=debug)
44
+ super().__init__(machine, dml, debug)
45
+
30
46
  self.tables = tables
31
47
  self.columns = columns
32
48
  self.partition_columns = partition_columns
33
49
  self.table_props = table_props
50
+ self.variant = variant
34
51
  self.debug = debug
35
- self.machine = StateMachine(variant=variant, debug=self.debug)
36
-
37
- def get_completions(
38
- self, document: Document, complete_event: CompleteEvent
39
- ) -> Iterable[Completion]:
40
- text = document.text_before_cursor.lstrip()
41
- state = ''
42
- if self.dml:
43
- state = f'{self.dml}_'
44
- text = f'{self.dml} {text}'
45
-
46
- completer: Completer = None
47
- stmts = sqlparse.parse(text)
48
- if not stmts:
49
- completer = DML_COMPLETER
52
+
53
+ def suggestions_completer(self, state: State, suggestions: str) -> list[str]:
54
+ if not suggestions:
55
+ return None
56
+
57
+ terms = []
58
+ for suggestion in suggestions.split(','):
59
+ terms.extend(self._terms(state, suggestion))
60
+
61
+ return TermCompleter(terms)
62
+
63
+ def _terms(self, state: State, word: str) -> list[str]:
64
+ terms = []
65
+
66
+ if word == 'tables':
67
+ terms.extend(self.tables())
68
+ elif word == '`tables`':
69
+ terms.append('tables')
70
+ elif word == 'columns':
71
+ terms.extend(self.columns([]))
72
+ elif word == 'partition-columns':
73
+ terms.extend(self.partition_columns([]))
74
+ elif word == 'table-props':
75
+ terms.extend(self.table_props().keys())
76
+ elif word == 'table-prop-values':
77
+ if 'last_name' in state.context and state.context['last_name']:
78
+ terms.extend(self.table_props()[state.context['last_name']])
79
+ elif word == 'single':
80
+ terms.append("'")
81
+ elif word == 'comma':
82
+ terms.append(",")
50
83
  else:
51
- statement: Statement = stmts[0]
52
- state: StateTo = self.machine.traverse_tokens(statement.tokens, StateTo(state))
53
- if self.debug:
54
- print('\n =>', state.to_s if isinstance(state, StateTo) else '')
55
-
56
- if not state or not state.to_s:
57
- completer = DML_COMPLETER
58
-
59
- if state and state.to_s in self.machine.suggestions:
60
- terms = []
61
-
62
- for word in self.machine.suggestions[state.to_s].strip(' ').split(','):
63
- if word == 'tables':
64
- terms.extend(self.tables())
65
- elif word == '`tables`':
66
- terms.append('tables')
67
- elif word == 'columns':
68
- terms.extend(self.columns([]))
69
- elif word == 'partition-columns':
70
- terms.extend(self.partition_columns([]))
71
- elif word == 'table-props':
72
- terms.extend(self.table_props().keys())
73
- elif word == 'table-prop-values':
74
- if 'last_name' in state.context and state.context['last_name']:
75
- terms.extend(self.table_props()[state.context['last_name']])
76
- elif word == 'single':
77
- terms.append("'")
78
- elif word == 'comma':
79
- terms.append(",")
80
- else:
81
- terms.append(word)
82
-
83
- if terms:
84
- completer = TermCompleter(terms)
85
-
86
- if completer:
87
- for c in completer.get_completions(document, complete_event):
88
- yield c
89
-
90
- def completions(table_names: Callable[[], list[str]],
91
- columns: Callable[[list[str]], list[str]] = default_columns,
92
- partition_columns: Callable[[list[str]], list[str]] = lambda x: [],
93
- table_props: Callable[[], dict[str, list[str]]] = lambda: [],
94
- variant = 'sql'):
84
+ terms.append(word)
85
+
86
+ return terms
87
+
88
+ def completions_for_nesting(self):
95
89
  return {
96
- 'delete': SqlCompleter(table_names, 'delete', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
97
- 'insert': SqlCompleter(table_names, 'insert', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
98
- 'select': SqlCompleter(table_names, 'select', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
99
- 'update': SqlCompleter(table_names, 'update', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
100
- 'alter': SqlCompleter(table_names, 'alter', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
101
- 'describe': SqlCompleter(table_names, 'describe', columns=columns, partition_columns=partition_columns, table_props=table_props, variant=variant),
90
+ word : SqlCompleter(self.tables, word, columns=self.columns, partition_columns=self.partition_columns,
91
+ table_props=self.table_props, variant=self.variant)
92
+ for word in self.machine.suggestions[''].strip(' ').split(',')
102
93
  }