kaqing 2.0.174__py3-none-any.whl → 2.0.188__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 (138) hide show
  1. adam/app_session.py +2 -2
  2. adam/apps.py +18 -4
  3. adam/batch.py +1 -1
  4. adam/checks/check_utils.py +3 -1
  5. adam/commands/__init__.py +4 -2
  6. adam/commands/app/__init__.py +0 -0
  7. adam/commands/app/app.py +38 -0
  8. adam/commands/app/app_ping.py +38 -0
  9. adam/commands/app/show_app_actions.py +49 -0
  10. adam/commands/app/show_app_id.py +44 -0
  11. adam/commands/app/show_app_queues.py +38 -0
  12. adam/commands/app/utils_app.py +106 -0
  13. adam/commands/audit/audit.py +9 -27
  14. adam/commands/audit/audit_repair_tables.py +5 -7
  15. adam/commands/audit/audit_run.py +1 -1
  16. adam/commands/audit/completions_l.py +15 -0
  17. adam/commands/audit/show_last10.py +2 -14
  18. adam/commands/audit/show_slow10.py +2 -13
  19. adam/commands/audit/show_top10.py +2 -11
  20. adam/commands/audit/utils_show_top10.py +14 -1
  21. adam/commands/bash/bash.py +1 -1
  22. adam/commands/cat.py +3 -7
  23. adam/commands/check.py +2 -2
  24. adam/commands/cli_commands.py +6 -1
  25. adam/commands/{cp.py → clipboard_copy.py} +18 -12
  26. adam/commands/code.py +2 -2
  27. adam/commands/command.py +61 -11
  28. adam/commands/commands_utils.py +19 -12
  29. adam/commands/cql/completions_c.py +28 -0
  30. adam/commands/cql/cqlsh.py +3 -7
  31. adam/commands/cql/utils_cql.py +22 -58
  32. adam/commands/deploy/deploy_pg_agent.py +2 -2
  33. adam/commands/deploy/undeploy_pg_agent.py +2 -2
  34. adam/commands/devices/device.py +39 -8
  35. adam/commands/devices/device_app.py +18 -28
  36. adam/commands/devices/device_auit_log.py +3 -3
  37. adam/commands/devices/device_cass.py +16 -22
  38. adam/commands/devices/device_export.py +6 -3
  39. adam/commands/devices/device_postgres.py +79 -63
  40. adam/commands/download_file.py +47 -0
  41. adam/commands/export/clean_up_all_export_sessions.py +3 -3
  42. adam/commands/export/clean_up_export_sessions.py +5 -10
  43. adam/commands/export/completions_x.py +11 -0
  44. adam/commands/export/download_export_session.py +40 -0
  45. adam/commands/export/export.py +0 -16
  46. adam/commands/export/export_databases.py +26 -9
  47. adam/commands/export/export_select.py +9 -58
  48. adam/commands/export/export_sessions.py +90 -5
  49. adam/commands/export/export_use.py +13 -10
  50. adam/commands/export/export_x_select.py +48 -0
  51. adam/commands/export/exporter.py +60 -22
  52. adam/commands/export/import_files.py +44 -0
  53. adam/commands/export/import_session.py +8 -4
  54. adam/commands/export/importer.py +7 -0
  55. adam/commands/export/importer_athena.py +101 -34
  56. adam/commands/export/importer_sqlite.py +30 -5
  57. adam/commands/export/show_column_counts.py +11 -11
  58. adam/commands/export/show_export_databases.py +5 -3
  59. adam/commands/export/show_export_session.py +5 -6
  60. adam/commands/export/show_export_sessions.py +4 -11
  61. adam/commands/export/utils_export.py +42 -14
  62. adam/commands/find_files.py +51 -0
  63. adam/commands/find_processes.py +76 -0
  64. adam/commands/head.py +36 -0
  65. adam/commands/help.py +2 -2
  66. adam/commands/intermediate_command.py +6 -3
  67. adam/commands/ls.py +1 -1
  68. adam/commands/medusa/medusa_backup.py +12 -14
  69. adam/commands/medusa/medusa_restore.py +20 -15
  70. adam/commands/medusa/medusa_show_backupjobs.py +6 -4
  71. adam/commands/medusa/medusa_show_restorejobs.py +5 -3
  72. adam/commands/medusa/utils_medusa.py +15 -0
  73. adam/commands/nodetool.py +3 -8
  74. adam/commands/param_get.py +2 -3
  75. adam/commands/param_set.py +1 -1
  76. adam/commands/postgres/completions_p.py +22 -0
  77. adam/commands/postgres/postgres.py +14 -21
  78. adam/commands/postgres/postgres_databases.py +270 -0
  79. adam/commands/postgres/utils_postgres.py +29 -20
  80. adam/commands/preview_table.py +3 -1
  81. adam/commands/pwd.py +3 -3
  82. adam/commands/reaper/reaper_forward.py +2 -2
  83. adam/commands/reaper/reaper_runs.py +3 -3
  84. adam/commands/reaper/reaper_schedule_activate.py +6 -2
  85. adam/commands/reaper/reaper_schedule_start.py +1 -2
  86. adam/commands/reaper/reaper_schedule_stop.py +1 -2
  87. adam/commands/reaper/utils_reaper.py +13 -6
  88. adam/commands/repair/repair_scan.py +0 -2
  89. adam/commands/repair/repair_stop.py +0 -1
  90. adam/commands/shell.py +7 -5
  91. adam/commands/show/show.py +1 -1
  92. adam/commands/show/show_adam.py +3 -3
  93. adam/commands/show/show_cassandra_repairs.py +5 -3
  94. adam/commands/show/show_cassandra_status.py +27 -20
  95. adam/commands/show/{show_commands.py → show_cli_commands.py} +2 -2
  96. adam/commands/show/show_login.py +2 -2
  97. adam/commands/show/show_params.py +2 -5
  98. adam/commands/show/show_processes.py +15 -14
  99. adam/commands/show/show_storage.py +9 -8
  100. adam/config.py +1 -0
  101. adam/embedded_params.py +1 -1
  102. adam/repl.py +16 -9
  103. adam/repl_commands.py +16 -9
  104. adam/repl_session.py +8 -1
  105. adam/repl_state.py +33 -10
  106. adam/sql/lark_completer.py +284 -0
  107. adam/sql/lark_parser.py +604 -0
  108. adam/sql/sql_state_machine.py +8 -2
  109. adam/utils.py +116 -29
  110. adam/utils_athena.py +7 -8
  111. adam/utils_issues.py +2 -2
  112. adam/utils_k8s/app_clusters.py +2 -2
  113. adam/utils_k8s/app_pods.py +5 -2
  114. adam/utils_k8s/cassandra_clusters.py +11 -3
  115. adam/utils_k8s/cassandra_nodes.py +2 -2
  116. adam/utils_k8s/k8s.py +9 -0
  117. adam/utils_k8s/kube_context.py +2 -2
  118. adam/utils_k8s/pods.py +23 -5
  119. adam/utils_k8s/statefulsets.py +5 -2
  120. adam/utils_local.py +4 -0
  121. adam/utils_repl/appendable_completer.py +6 -0
  122. adam/utils_repl/repl_completer.py +128 -2
  123. adam/utils_sqlite.py +2 -2
  124. adam/version.py +1 -1
  125. {kaqing-2.0.174.dist-info → kaqing-2.0.188.dist-info}/METADATA +1 -1
  126. kaqing-2.0.188.dist-info/RECORD +253 -0
  127. kaqing-2.0.188.dist-info/top_level.txt +2 -0
  128. teddy/__init__.py +0 -0
  129. teddy/lark_parser.py +436 -0
  130. teddy/lark_parser2.py +618 -0
  131. adam/commands/cql/cql_completions.py +0 -32
  132. adam/commands/export/export_select_x.py +0 -54
  133. adam/commands/postgres/postgres_context.py +0 -272
  134. adam/commands/postgres/psql_completions.py +0 -10
  135. kaqing-2.0.174.dist-info/RECORD +0 -230
  136. kaqing-2.0.174.dist-info/top_level.txt +0 -1
  137. {kaqing-2.0.174.dist-info → kaqing-2.0.188.dist-info}/WHEEL +0 -0
  138. {kaqing-2.0.174.dist-info → kaqing-2.0.188.dist-info}/entry_points.txt +0 -0
@@ -3,7 +3,7 @@ from abc import abstractmethod
3
3
  from adam.commands.command import Command
4
4
  from adam.commands.command_helpers import ClusterCommandHelper
5
5
  from adam.repl_state import ReplState
6
- from adam.utils import lines_to_tabular, log, log2
6
+ from adam.utils import tabulize, log, log2
7
7
 
8
8
  class IntermediateCommand(Command):
9
9
  def run(self, cmd: str, state: ReplState):
@@ -12,6 +12,9 @@ class IntermediateCommand(Command):
12
12
 
13
13
  return self.intermediate_run(cmd, state, args, self.cmd_list())
14
14
 
15
+ def completion(self, state: ReplState):
16
+ return {}
17
+
15
18
  @abstractmethod
16
19
  def cmd_list(self):
17
20
  pass
@@ -21,7 +24,7 @@ class IntermediateCommand(Command):
21
24
 
22
25
  if state.in_repl:
23
26
  if display_help:
24
- log(lines_to_tabular([c.help(state) for c in cmds], separator=separator))
27
+ tabulize(cmds, lambda c: c.help(state), separator=separator)
25
28
 
26
29
  return 'command-missing'
27
30
  else:
@@ -43,7 +46,7 @@ class IntermediateCommand(Command):
43
46
  log()
44
47
  log('Sub-Commands:')
45
48
 
46
- log(lines_to_tabular([c.help(ReplState()).replace(f'{cmd} ', ' ', 1) for c in cmd_list], separator=separator))
49
+ tabulize(cmd_list, lambda c: c.help(ReplState()).replace(f'{cmd} ', ' ', 1), separator=separator)
47
50
  if show_cluster_help:
48
51
  log()
49
52
  ClusterCommandHelper.cluster_help()
adam/commands/ls.py CHANGED
@@ -35,7 +35,7 @@ class Ls(Command):
35
35
  return state
36
36
 
37
37
  def completion(self, state: ReplState):
38
- return Devices.device(state).ls_completion(Ls.COMMAND, state, default = super().completion(state))
38
+ return super().completion(state, {'&': None}, pods=Devices.device(state).pods(state, '-'))
39
39
 
40
40
  def help(self, _: ReplState):
41
41
  return f'{Ls.COMMAND} [device:]\t list apps, envs, clusters, nodes, pg hosts/databases or export databases'
@@ -1,5 +1,6 @@
1
1
  from datetime import datetime
2
2
 
3
+ from adam.commands import validate_args
3
4
  from adam.commands.command import Command
4
5
  from adam.utils_k8s.statefulsets import StatefulSets
5
6
  from adam.repl_state import ReplState, RequiredState
@@ -32,23 +33,20 @@ class MedusaBackup(Command):
32
33
  ns = state.namespace
33
34
  sts = state.sts
34
35
  now_dtformat = datetime.now().strftime("%Y-%m-%d.%H.%M.%S")
35
- bkname = 'medusa-' + now_dtformat + 'full-backup-' + sts
36
- if len(args) == 1:
37
- bkname = str(args[0])
36
+ with validate_args(args, state, default='medusa-' + now_dtformat + 'full-backup-' + sts) as bkname:
37
+ dc = StatefulSets.get_datacenter(state.sts, ns)
38
+ if not dc:
39
+ return state
40
+
41
+ try:
42
+ CustomResources.create_medusa_backupjob(bkname, dc, ns)
43
+ except Exception as e:
44
+ log2("Exception: MedusaBackup failed: %s\n" % e)
45
+ finally:
46
+ CustomResources.clear_caches()
38
47
 
39
- dc = StatefulSets.get_datacenter(state.sts, ns)
40
- if not dc:
41
48
  return state
42
49
 
43
- try:
44
- CustomResources.create_medusa_backupjob(bkname, dc, ns)
45
- except Exception as e:
46
- log2("Exception: MedusaBackup failed: %s\n" % e)
47
- finally:
48
- CustomResources.clear_caches()
49
-
50
- return state
51
-
52
50
  def completion(self, state: ReplState):
53
51
  return super().completion(state)
54
52
 
@@ -4,11 +4,12 @@ from functools import partial
4
4
 
5
5
  from adam.commands import validate_args
6
6
  from adam.commands.command import Command, InvalidArgumentsException
7
+ from adam.commands.medusa.utils_medusa import medusa_backup_names
7
8
  from adam.utils_k8s.statefulsets import StatefulSets
8
9
  from adam.repl_state import ReplState, RequiredState
9
10
  from adam.utils_k8s.custom_resources import CustomResources
10
11
  from adam.config import Config
11
- from adam.utils import lines_to_tabular, log2, log_exc
12
+ from adam.utils import tabulize, log2, log_exc
12
13
 
13
14
  class MedusaRestore(Command):
14
15
  COMMAND = 'restore'
@@ -45,8 +46,11 @@ class MedusaRestore(Command):
45
46
  else:
46
47
  log2('\n* Backup job name is not valid.')
47
48
 
48
- bklist = [f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '')}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]
49
- log2(lines_to_tabular(bklist, 'NAME\tCREATED\tFINISHED', separator='\t'))
49
+ tabulize(CustomResources.medusa_show_backupjobs(dc, ns),
50
+ lambda x: f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '')}",
51
+ header='NAME\tCREATED\tFINISHED',
52
+ separator='\t',
53
+ to=2)
50
54
 
51
55
  with validate_args(args, state, msg=partial(msg, True)) as bkname:
52
56
  if not (job := CustomResources.medusa_get_backupjob(dc, ns, bkname)):
@@ -64,18 +68,19 @@ class MedusaRestore(Command):
64
68
  return state
65
69
 
66
70
  def completion(self, state: ReplState):
67
- if sc := super().completion(state):
68
- ns = state.namespace
69
- dc: str = StatefulSets.get_datacenter(state.sts, ns)
70
- if not dc:
71
- return {}
72
-
73
- if Config().get('medusa.restore-auto-complete', False):
74
- leaf = {id: None for id in [f"{x['metadata']['name']}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]}
75
-
76
- return super().completion(state, leaf)
77
- else:
78
- return sc
71
+ return super().completion(state, lambda: {id: None for id in medusa_backup_names(state)}, auto_key='medusa.backups')
72
+ # if sc := super().completion(state):
73
+ # ns = state.namespace
74
+ # dc: str = StatefulSets.get_datacenter(state.sts, ns)
75
+ # if not dc:
76
+ # return {}
77
+
78
+ # if Config().get('medusa.restore-auto-complete', False):
79
+ # leaf = {id: None for id in [f"{x['metadata']['name']}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]}
80
+
81
+ # return super().completion(state, leaf)
82
+ # else:
83
+ # return sc
79
84
 
80
85
  return {}
81
86
 
@@ -2,7 +2,7 @@ from adam.commands.command import Command
2
2
  from adam.utils_k8s.statefulsets import StatefulSets
3
3
  from adam.repl_state import ReplState, RequiredState
4
4
  from adam.utils_k8s.custom_resources import CustomResources
5
- from adam.utils import lines_to_tabular, log2, log_exc
5
+ from adam.utils import tabulize, log_exc
6
6
 
7
7
 
8
8
  class MedusaShowBackupJobs(Command):
@@ -33,12 +33,14 @@ class MedusaShowBackupJobs(Command):
33
33
  if not dc:
34
34
  return state
35
35
 
36
- # try:
37
36
  with log_exc(lambda e: "Exception: MedusaShowBackupJobs failed: %s\n" % e):
38
37
  CustomResources.clear_caches()
39
38
 
40
- bklist = [f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '') if 'status' in x else 'unknown'}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]
41
- log2(lines_to_tabular(bklist, 'NAME\tCREATED\tFINISHED', separator='\t'))
39
+ tabulize(CustomResources.medusa_show_backupjobs(dc, ns),
40
+ lambda x: f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '') if 'status' in x else 'unknown'}",
41
+ header='NAME\tCREATED\tFINISHED',
42
+ separator='\t',
43
+ to=2)
42
44
 
43
45
  return state
44
46
 
@@ -2,7 +2,7 @@ from adam.commands.command import Command
2
2
  from adam.utils_k8s.statefulsets import StatefulSets
3
3
  from adam.repl_state import ReplState, RequiredState
4
4
  from adam.utils_k8s.custom_resources import CustomResources
5
- from adam.utils import lines_to_tabular, log2, log_exc
5
+ from adam.utils import tabulize, log_exc
6
6
 
7
7
  class MedusaShowRestoreJobs(Command):
8
8
  COMMAND = 'show restores'
@@ -33,8 +33,10 @@ class MedusaShowRestoreJobs(Command):
33
33
  return state
34
34
 
35
35
  with log_exc(lambda e: "Exception: MedusaShowRestoreJobs failed: %s\n" % e):
36
- rtlist = CustomResources.medusa_show_restorejobs(dc, ns)
37
- log2(lines_to_tabular(rtlist, 'NAME\tCREATED\tFINISHED', separator='\t'))
36
+ tabulize(CustomResources.medusa_show_restorejobs(dc, ns),
37
+ header='NAME\tCREATED\tFINISHED',
38
+ separator='\t',
39
+ to=2)
38
40
 
39
41
  return state
40
42
 
@@ -0,0 +1,15 @@
1
+ from adam.config import Config
2
+ from adam.repl_state import ReplState
3
+ from adam.utils_k8s.custom_resources import CustomResources
4
+ from adam.utils_k8s.statefulsets import StatefulSets
5
+
6
+ def medusa_backup_names(state: ReplState, warm=False):
7
+ if warm and (auto := Config().get('medusa.restore-auto-complete', 'off')) in ['off', 'jit', 'lazy']:
8
+ return {}
9
+
10
+ ns = state.namespace
11
+ dc: str = StatefulSets.get_datacenter(state.sts, ns)
12
+ if not dc:
13
+ return {}
14
+
15
+ return {id: None for id in [f"{x['metadata']['name']}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]}
adam/commands/nodetool.py CHANGED
@@ -3,11 +3,11 @@ import click
3
3
  from adam.commands.command import Command
4
4
  from adam.commands.command_helpers import ClusterOrPodCommandHelper
5
5
  from adam.commands.cql.utils_cql import cassandra
6
+ from adam.commands.devices.devices import Devices
6
7
  from adam.commands.nodetool_commands import NODETOOL_COMMANDS
7
8
  from adam.config import Config
8
9
  from adam.repl_state import ReplState, RequiredState
9
10
  from adam.utils import log
10
- from adam.utils_k8s.statefulsets import StatefulSets
11
11
 
12
12
  class NodeTool(Command):
13
13
  COMMAND = 'nodetool'
@@ -33,17 +33,12 @@ class NodeTool(Command):
33
33
 
34
34
  with self.validate(args, state) as (args, state):
35
35
  with cassandra(state) as pods:
36
- pods.nodetool(' '.join(args))
36
+ pods.nodetool(' '.join(args), status=(args[0] == 'status'))
37
37
 
38
38
  return state
39
39
 
40
40
  def completion(self, state: ReplState):
41
- if super().completion(state):
42
- d = {c: {'&': None} for c in NODETOOL_COMMANDS}
43
- return {NodeTool.COMMAND: {'help': None} | d} | \
44
- {f'@{p}': {NodeTool.COMMAND: d} for p in StatefulSets.pod_names(state.sts, state.namespace)}
45
-
46
- return {}
41
+ return super().completion(state, {c: {'&': None} for c in NODETOOL_COMMANDS}, pods=Devices.device(state).pods(state, '-'))
47
42
 
48
43
  def help(self, _: ReplState):
49
44
  return f'{NodeTool.COMMAND} <sub-command> [&]\t run nodetool with arguments'
@@ -2,7 +2,7 @@ from adam.commands import validate_args
2
2
  from adam.commands.command import Command
3
3
  from adam.config import Config
4
4
  from adam.repl_state import ReplState
5
- from adam.utils import lines_to_tabular, log, log2
5
+ from adam.utils import tabulize, log, log2
6
6
 
7
7
  class GetParam(Command):
8
8
  COMMAND = 'get'
@@ -25,8 +25,7 @@ class GetParam(Command):
25
25
 
26
26
  with self.validate(args, state) as (args, state):
27
27
  def msg():
28
- lines = [f'{key}\t{Config().get(key, None)}' for key in Config().keys()]
29
- log(lines_to_tabular(lines, separator='\t'))
28
+ tabulize(Config().keys(), lambda key: f'{key}\t{Config().get(key, None)}', separator='\t')
30
29
 
31
30
  with validate_args(args, state, msg=msg) as key:
32
31
  if v := Config().get(key, None):
@@ -24,7 +24,7 @@ class SetParam(Command):
24
24
  return super().run(cmd, state)
25
25
 
26
26
  with self.validate(args, state) as (args, state):
27
- with validate_args(args, state, count_at_least=2, msg=lambda: log2('set <key> <value>')):
27
+ with validate_args(args, state, exactly=2, msg=lambda: log2('set <key> <value>')):
28
28
  key = args[0]
29
29
  value = args[1]
30
30
  Config().set(key, value)
@@ -0,0 +1,22 @@
1
+ from adam.commands.postgres.postgres_databases import PostgresDatabases
2
+ from adam.commands.postgres.utils_postgres import pg_table_names
3
+ from adam.repl_state import ReplState
4
+ from adam.sql.lark_completer import LarkCompleter
5
+
6
+ def completions_p(state: ReplState):
7
+ return {
8
+ '\h': None,
9
+ '\d': None,
10
+ '\dt': None,
11
+ '\du': None
12
+ } | LarkCompleter(expandables={
13
+ 'tables': lambda x: pg_table_names(state),
14
+ 'columns': lambda x: ['id'],
15
+ 'hosts': lambda x: ['@' + PostgresDatabases.pod_and_container(state.namespace)[0]],
16
+ }, variant=ReplState.P).completions_for_nesting()
17
+
18
+ def psql0_completions(state: ReplState):
19
+ return {
20
+ '\h': None,
21
+ '\l': None,
22
+ }
@@ -3,13 +3,13 @@ import click
3
3
  from adam.commands import extract_trailing_options, validate_args
4
4
  from adam.commands.command import Command
5
5
  from adam.commands.intermediate_command import IntermediateCommand
6
- from adam.commands.postgres.psql_completions import psql_completions
6
+ from adam.commands.postgres.postgres_databases import pg_path
7
+ from adam.commands.postgres.completions_p import psql0_completions, completions_p
7
8
  from adam.commands.postgres.utils_postgres import pg_table_names, postgres
8
9
  from .postgres_ls import PostgresLs
9
10
  from .postgres_preview import PostgresPreview
10
- from .postgres_context import PostgresContext
11
11
  from adam.repl_state import ReplState
12
- from adam.utils import log, log2
12
+ from adam.utils import log, log2, log_timing
13
13
 
14
14
  class Postgres(IntermediateCommand):
15
15
  COMMAND = 'pg'
@@ -43,10 +43,10 @@ class Postgres(IntermediateCommand):
43
43
 
44
44
  if state.in_repl:
45
45
  with postgres(state) as pod:
46
- pod.sql(args, background=backgrounded)
46
+ pod.sql(args, backgrounded=backgrounded)
47
47
  elif not self.run_subcommand(cmd, state):
48
48
  with postgres(state) as pod:
49
- pod.sql(args, background=backgrounded)
49
+ pod.sql(args, backgrounded=backgrounded)
50
50
 
51
51
  return state
52
52
 
@@ -55,24 +55,17 @@ class Postgres(IntermediateCommand):
55
55
 
56
56
  def completion(self, state: ReplState):
57
57
  if state.device != state.P:
58
- # conflicts with cql completions
59
58
  return {}
60
59
 
61
- leaf = {}
62
- session = PostgresContext.apply(state.namespace, state.pg_path)
63
- if session.db:
64
- if pg_table_names(state.namespace, state.pg_path):
65
- leaf = psql_completions(state.namespace, state.pg_path)
66
- elif state.pg_path:
67
- leaf = {
68
- '\h': None,
69
- '\l': None,
70
- }
71
-
72
- if state.pg_path:
73
- return super().completion(state, leaf) | leaf
74
- else:
75
- return {}
60
+ with pg_path(state) as (host, database):
61
+ if database:
62
+ if pg_table_names(state):
63
+ with log_timing('psql_completions'):
64
+ return completions_p(state)
65
+ elif host:
66
+ return psql0_completions(state)
67
+
68
+ return {}
76
69
 
77
70
  def help(self, _: ReplState):
78
71
  return f'<sql-statements> [&]\t run queries on Postgres databases'
@@ -0,0 +1,270 @@
1
+ import functools
2
+ import re
3
+ import subprocess
4
+
5
+ from adam.config import Config
6
+ from adam.repl_session import ReplSession
7
+ from adam.repl_state import ReplState
8
+ from adam.utils_k8s.kube_context import KubeContext
9
+ from adam.utils_k8s.pods import Pods
10
+ from adam.utils_k8s.secrets import Secrets
11
+ from adam.utils import log2, log_exc
12
+
13
+ class ConnectionDetails:
14
+ def __init__(self, state: ReplState, namespace: str, host: str):
15
+ self.state = state
16
+ self.namespace = namespace
17
+ self.host = host
18
+
19
+ def endpoint(self):
20
+ return PostgresDatabases._connection_property(self.state, 'pg.secret.endpoint-key', 'postgres-db-endpoint', host=self.host)
21
+
22
+ def port(self):
23
+ return PostgresDatabases._connection_property(self.state, 'pg.secret.port-key', 'postgres-db-port', host=self.host)
24
+
25
+ def username(self):
26
+ return PostgresDatabases._connection_property(self.state, 'pg.secret.username-key', 'postgres-admin-username', host=self.host)
27
+
28
+ def password(self):
29
+ return PostgresDatabases._connection_property(self.state, 'pg.secret.password-key', 'postgres-admin-password', host=self.host)
30
+
31
+ class PostgresDatabases:
32
+ def hosts(state: ReplState, namespace: str = None):
33
+ if not namespace:
34
+ namespace = state.namespace
35
+
36
+ return [ConnectionDetails(state, namespace, host) for host in PostgresDatabases.host_names(namespace)]
37
+
38
+ @functools.lru_cache()
39
+ def host_names(namespace: str):
40
+ ss = Secrets.list_secrets(namespace, name_pattern=Config().get('pg.name-pattern', '^{namespace}.*k8spg.*'))
41
+
42
+ def excludes(name: str):
43
+ exs = Config().get('pg.excludes', '.helm., -admin-secret')
44
+ if exs:
45
+ for ex in exs.split(','):
46
+ if ex.strip(' ') in name:
47
+ return True
48
+
49
+ return False
50
+
51
+ return [s for s in ss if not excludes(s)]
52
+
53
+ def databases(state: ReplState, default_owner = False):
54
+ dbs = []
55
+ # List of databases
56
+ # Name | Owner | Encoding | Collate | Ctype | ICU Locale | Locale Provider | Access privileges
57
+ # ---------------------------------------+----------+----------+-------------+-------------+------------+-----------------+-----------------------
58
+ # postgres | postgres | UTF8 | en_US.UTF-8 | en_US.UTF-8 | | libc |
59
+ # stgawsscpsr_c3_c3 | postgres | UTF8 | C | C | | libc |
60
+ # template1 | postgres | UTF8 | en_US.UTF-8 | en_US.UTF-8 | | libc | =c/postgres +
61
+ # | | | | | | | postgres=CTc/postgres
62
+ # (48 rows)
63
+ if r := PostgresDatabases.run_sql(state, '\l', show_out=False):
64
+ s = 0
65
+ for line in r.stdout.split('\n'):
66
+ line: str = line.strip(' \r')
67
+ if s == 0:
68
+ if 'List of databases' in line:
69
+ s = 1
70
+ elif s == 1:
71
+ if 'Name' in line and 'Owner' in line and 'Encoding' in line:
72
+ s = 2
73
+ elif s == 2:
74
+ if line.startswith('---------'):
75
+ s = 3
76
+ elif s == 3:
77
+ groups = re.match(r'^\s*(\S*)\s*\|\s*(\S*)\s*\|.*', line)
78
+ if groups and groups[1] != '|':
79
+ dbs.append({'name': groups[1], 'owner': groups[2]})
80
+
81
+ if default_owner:
82
+ dbs = [db for db in dbs if db['owner'] == PostgresDatabases.default_owner()]
83
+
84
+ return dbs
85
+
86
+ def tables(state: ReplState, default_schema = False):
87
+ dbs = []
88
+ # List of relations
89
+ # Schema | Name | Type | Owner
90
+ # ----------+------------------------------------------------------------+-------+---------------
91
+ # postgres | c3_2_admin_aclpriv | table | postgres
92
+ # postgres | c3_2_admin_aclpriv_a | table | postgres
93
+ if r := PostgresDatabases.run_sql(state, '\dt', show_out=False):
94
+ s = 0
95
+ for line in r.stdout.split('\n'):
96
+ line: str = line.strip(' \r')
97
+ if s == 0:
98
+ if 'List of relations' in line:
99
+ s = 1
100
+ elif s == 1:
101
+ if 'Schema' in line and 'Name' in line and 'Type' in line:
102
+ s = 2
103
+ elif s == 2:
104
+ if line.startswith('---------'):
105
+ s = 3
106
+ elif s == 3:
107
+ groups = re.match(r'^\s*(\S*)\s*\|\s*(\S*)\s*\|.*', line)
108
+ if groups and groups[1] != '|':
109
+ dbs.append({'schema': groups[1], 'name': groups[2]})
110
+
111
+ if default_schema:
112
+ dbs = [db for db in dbs if db["schema"] == PostgresDatabases.default_schema()]
113
+
114
+ return dbs
115
+
116
+ def run_sql(state: ReplState, sql: str, database: str = None, show_out = True, backgrounded = False):
117
+ if not database:
118
+ database = PostgresDatabases.database(state)
119
+ if not database:
120
+ database = PostgresDatabases.default_db()
121
+
122
+ username = PostgresDatabases.username(state)
123
+ password = PostgresDatabases.password(state)
124
+ endpoint = PostgresDatabases.endpoint(state)
125
+
126
+ if KubeContext.in_cluster():
127
+ cmd1 = f'env PGPASSWORD={password} psql -h {endpoint} -p {PostgresDatabases.port()} -U {username} {database} --pset pager=off -c'
128
+ log2(f'{cmd1} "{sql}"')
129
+ # remove double quotes from the sql argument
130
+ cmd = cmd1.split(' ') + [sql]
131
+
132
+ r = subprocess.run(cmd, capture_output=not backgrounded, text=True)
133
+ if show_out:
134
+ log2(r.stdout)
135
+ log2(r.stderr)
136
+
137
+ return r
138
+ else:
139
+ pod_name, container_name = PostgresDatabases.pod_and_container(state.namespace)
140
+ if not pod_name:
141
+ return
142
+
143
+ cmd = f'psql -h {endpoint} -p {PostgresDatabases.port(state)} -U {username} {database} --pset pager=off -c "{sql}"'
144
+ env_prefix = f'PGPASSWORD="{password}"'
145
+
146
+ r = Pods.exec(pod_name, container_name, state.namespace, cmd, show_out=show_out, backgrounded=backgrounded, env_prefix=env_prefix)
147
+ if r and Config().get('repl.history.push-cat-remote-log-file', True):
148
+ if r.log_file and ReplSession().prompt_session:
149
+ ReplSession().prompt_session.history.append_string(f'@{r.pod} cat {r.log_file}')
150
+
151
+ return r
152
+
153
+ @functools.lru_cache()
154
+ def pod_and_container(namespace: str):
155
+ container_name = Config().get('pg.agent.name', 'ops-pg-agent')
156
+
157
+ if Config().get('pg.agent.just-in-time', False):
158
+ if not PostgresDatabases.deploy_pg_agent(container_name, namespace):
159
+ return None
160
+
161
+ pod_name = container_name
162
+ try:
163
+ # try with dedicated pg agent pod name configured
164
+ Pods.get(namespace, container_name)
165
+ except:
166
+ try:
167
+ # try with the ops pod
168
+ container_name = Config().get('pod.name', 'ops')
169
+ pod_name = Pods.get_with_selector(namespace, label_selector = Config().get('pod.label-selector', 'run=ops')).metadata.name
170
+ except:
171
+ log2(f"Could not locate {container_name} pod.")
172
+ return None
173
+
174
+ return pod_name, container_name
175
+
176
+ def deploy_pg_agent(pod_name: str, namespace: str) -> str:
177
+ image = Config().get('pg.agent.image', 'seanahnsf/kaqing')
178
+ timeout = Config().get('pg.agent.timeout', 3600)
179
+ try:
180
+ Pods.create(namespace, pod_name, image, ['sleep', f'{timeout}'], env={'NAMESPACE': namespace}, sa_name='c3')
181
+ except Exception as e:
182
+ if e.status == 409:
183
+ if Pods.completed(namespace, pod_name):
184
+ with log_exc(lambda e2: "Exception when calling BatchV1Api->create_pod: %s\n" % e2):
185
+ Pods.delete(pod_name, namespace)
186
+ Pods.create(namespace, pod_name, image, ['sleep', f'{timeout}'], env={'NAMESPACE': namespace}, sa_name='c3')
187
+
188
+ return
189
+ else:
190
+ log2("Exception when calling BatchV1Api->create_pod: %s\n" % e)
191
+
192
+ return
193
+
194
+ Pods.wait_for_running(namespace, pod_name)
195
+
196
+ return pod_name
197
+
198
+ def undeploy_pg_agent(pod_name: str, namespace: str):
199
+ Pods.delete(pod_name, namespace, grace_period_seconds=0)
200
+
201
+ def endpoint(state: ReplState):
202
+ return PostgresDatabases._connection_property(state, 'pg.secret.endpoint-key', 'postgres-db-endpoint')
203
+
204
+ def port(state: ReplState):
205
+ return PostgresDatabases._connection_property(state, 'pg.secret.port-key', 'postgres-db-port')
206
+
207
+ def username(state: ReplState):
208
+ return PostgresDatabases._connection_property(state, 'pg.secret.username-key', 'postgres-admin-username')
209
+
210
+ def password(state: ReplState):
211
+ return PostgresDatabases._connection_property(state, 'pg.secret.password-key', 'postgres-admin-password')
212
+
213
+ def _connection_property(state: ReplState, config_key: str, default: str, host: str = None, database: str = None):
214
+ with pg_path(state, host=host, database=database) as (host, _):
215
+ if not (conn := PostgresDatabases.conn_details(state.namespace, host)):
216
+ return ''
217
+
218
+ key = Config().get(config_key, default)
219
+ return conn[key] if key in conn else ''
220
+
221
+ def default_db():
222
+ return Config().get('pg.default-db', 'postgres')
223
+
224
+ def default_owner():
225
+ return Config().get('pg.default-owner', 'postgres')
226
+
227
+ def default_schema():
228
+ return Config().get('pg.default-schema', 'postgres')
229
+
230
+ def host(state: ReplState):
231
+ if not state.pg_path:
232
+ return None
233
+
234
+ return state.pg_path.split('/')[0]
235
+
236
+ def database(state: ReplState):
237
+ if not state.pg_path:
238
+ return None
239
+
240
+ tokens = state.pg_path.split('/')
241
+ if len(tokens) > 1:
242
+ return tokens[1]
243
+
244
+ return None
245
+
246
+ @functools.lru_cache()
247
+ def conn_details(namespace: str, host: str):
248
+ return Secrets.get_data(namespace, host)
249
+
250
+ class PostgresPathHandler:
251
+ def __init__(self, state: ReplState, host: str = None, database: str = None):
252
+ self.state = state
253
+ self.host = host
254
+ self.database = database
255
+
256
+ def __enter__(self) -> tuple[str, str]:
257
+ if self.state and self.state.pg_path:
258
+ host_n_db = self.state.pg_path.split('/')
259
+ if not self.host:
260
+ self.host = host_n_db[0]
261
+ if not self.database and len(host_n_db) > 1:
262
+ self.database = host_n_db[1]
263
+
264
+ return self.host, self.database
265
+
266
+ def __exit__(self, exc_type, exc_val, exc_tb):
267
+ return False
268
+
269
+ def pg_path(state: ReplState, host: str = None, database: str = None):
270
+ return PostgresPathHandler(state, host=host, database=database)