kaqing 2.0.172__py3-none-any.whl → 2.0.186__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 (140) 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 +8 -2
  6. adam/commands/alter_tables.py +24 -35
  7. adam/commands/app/__init__.py +0 -0
  8. adam/commands/app/app.py +38 -0
  9. adam/commands/app/app_ping.py +38 -0
  10. adam/commands/app/show_app_actions.py +49 -0
  11. adam/commands/app/show_app_id.py +44 -0
  12. adam/commands/app/show_app_queues.py +38 -0
  13. adam/commands/app/utils_app.py +106 -0
  14. adam/commands/audit/audit.py +9 -27
  15. adam/commands/audit/audit_repair_tables.py +5 -7
  16. adam/commands/audit/audit_run.py +1 -1
  17. adam/commands/audit/completions_l.py +15 -0
  18. adam/commands/audit/show_last10.py +2 -14
  19. adam/commands/audit/show_slow10.py +2 -13
  20. adam/commands/audit/show_top10.py +2 -11
  21. adam/commands/audit/utils_show_top10.py +14 -1
  22. adam/commands/bash/bash.py +1 -1
  23. adam/commands/cat.py +5 -19
  24. adam/commands/cd.py +6 -8
  25. adam/commands/check.py +10 -18
  26. adam/commands/cli_commands.py +6 -1
  27. adam/commands/{cp.py → clipboard_copy.py} +34 -36
  28. adam/commands/code.py +2 -2
  29. adam/commands/command.py +94 -10
  30. adam/commands/commands_utils.py +19 -12
  31. adam/commands/cql/completions_c.py +28 -0
  32. adam/commands/cql/cqlsh.py +3 -7
  33. adam/commands/cql/utils_cql.py +22 -60
  34. adam/commands/deploy/deploy_pg_agent.py +2 -2
  35. adam/commands/deploy/undeploy_pg_agent.py +2 -2
  36. adam/commands/devices/device.py +39 -8
  37. adam/commands/devices/device_app.py +19 -29
  38. adam/commands/devices/device_auit_log.py +3 -3
  39. adam/commands/devices/device_cass.py +17 -23
  40. adam/commands/devices/device_export.py +12 -11
  41. adam/commands/devices/device_postgres.py +79 -63
  42. adam/commands/download_file.py +47 -0
  43. adam/commands/export/clean_up_all_export_sessions.py +3 -3
  44. adam/commands/export/clean_up_export_sessions.py +7 -19
  45. adam/commands/export/completions_x.py +11 -0
  46. adam/commands/export/download_export_session.py +40 -0
  47. adam/commands/export/drop_export_database.py +6 -22
  48. adam/commands/export/drop_export_databases.py +3 -9
  49. adam/commands/export/export.py +1 -17
  50. adam/commands/export/export_databases.py +93 -21
  51. adam/commands/export/export_select.py +8 -68
  52. adam/commands/export/export_sessions.py +209 -0
  53. adam/commands/export/export_use.py +13 -16
  54. adam/commands/export/export_x_select.py +48 -0
  55. adam/commands/export/exporter.py +108 -129
  56. adam/commands/export/import_files.py +44 -0
  57. adam/commands/export/import_session.py +10 -6
  58. adam/commands/export/importer.py +19 -5
  59. adam/commands/export/importer_athena.py +112 -41
  60. adam/commands/export/importer_sqlite.py +47 -19
  61. adam/commands/export/show_column_counts.py +11 -20
  62. adam/commands/export/show_export_databases.py +5 -2
  63. adam/commands/export/show_export_session.py +6 -15
  64. adam/commands/export/show_export_sessions.py +4 -11
  65. adam/commands/export/utils_export.py +46 -16
  66. adam/commands/find_files.py +51 -0
  67. adam/commands/find_processes.py +76 -0
  68. adam/commands/head.py +36 -0
  69. adam/commands/help.py +2 -2
  70. adam/commands/intermediate_command.py +6 -3
  71. adam/commands/ls.py +1 -1
  72. adam/commands/medusa/medusa_backup.py +13 -16
  73. adam/commands/medusa/medusa_restore.py +39 -32
  74. adam/commands/medusa/medusa_show_backupjobs.py +6 -4
  75. adam/commands/medusa/medusa_show_restorejobs.py +5 -3
  76. adam/commands/medusa/utils_medusa.py +15 -0
  77. adam/commands/nodetool.py +3 -8
  78. adam/commands/param_get.py +10 -12
  79. adam/commands/param_set.py +7 -10
  80. adam/commands/postgres/completions_p.py +22 -0
  81. adam/commands/postgres/postgres.py +25 -40
  82. adam/commands/postgres/postgres_databases.py +270 -0
  83. adam/commands/postgres/utils_postgres.py +33 -20
  84. adam/commands/preview_table.py +4 -2
  85. adam/commands/pwd.py +3 -3
  86. adam/commands/reaper/reaper_forward.py +2 -2
  87. adam/commands/reaper/reaper_run_abort.py +4 -10
  88. adam/commands/reaper/reaper_runs.py +3 -3
  89. adam/commands/reaper/reaper_schedule_activate.py +12 -12
  90. adam/commands/reaper/reaper_schedule_start.py +7 -12
  91. adam/commands/reaper/reaper_schedule_stop.py +7 -12
  92. adam/commands/reaper/utils_reaper.py +13 -6
  93. adam/commands/repair/repair_scan.py +0 -2
  94. adam/commands/repair/repair_stop.py +0 -1
  95. adam/commands/shell.py +7 -5
  96. adam/commands/show/show.py +1 -1
  97. adam/commands/show/show_adam.py +3 -3
  98. adam/commands/show/show_cassandra_repairs.py +5 -3
  99. adam/commands/show/show_cassandra_status.py +27 -20
  100. adam/commands/show/{show_commands.py → show_cli_commands.py} +2 -2
  101. adam/commands/show/show_login.py +2 -2
  102. adam/commands/show/show_params.py +2 -5
  103. adam/commands/show/show_processes.py +15 -14
  104. adam/commands/show/show_storage.py +9 -8
  105. adam/config.py +1 -0
  106. adam/embedded_params.py +1 -1
  107. adam/repl.py +20 -11
  108. adam/repl_commands.py +16 -9
  109. adam/repl_session.py +8 -1
  110. adam/repl_state.py +33 -10
  111. adam/sql/lark_completer.py +280 -0
  112. adam/sql/lark_parser.py +604 -0
  113. adam/sql/sql_state_machine.py +8 -2
  114. adam/utils.py +116 -29
  115. adam/utils_athena.py +7 -8
  116. adam/utils_issues.py +2 -2
  117. adam/utils_k8s/app_clusters.py +2 -2
  118. adam/utils_k8s/app_pods.py +5 -2
  119. adam/utils_k8s/cassandra_clusters.py +11 -3
  120. adam/utils_k8s/cassandra_nodes.py +2 -2
  121. adam/utils_k8s/k8s.py +14 -5
  122. adam/utils_k8s/kube_context.py +2 -2
  123. adam/utils_k8s/pods.py +23 -5
  124. adam/utils_k8s/statefulsets.py +5 -2
  125. adam/utils_local.py +4 -0
  126. adam/utils_repl/appendable_completer.py +6 -0
  127. adam/utils_repl/repl_completer.py +128 -2
  128. adam/utils_sqlite.py +14 -14
  129. adam/version.py +1 -1
  130. {kaqing-2.0.172.dist-info → kaqing-2.0.186.dist-info}/METADATA +1 -1
  131. kaqing-2.0.186.dist-info/RECORD +250 -0
  132. adam/commands/cql/cql_completions.py +0 -33
  133. adam/commands/export/export_handlers.py +0 -71
  134. adam/commands/export/export_select_x.py +0 -54
  135. adam/commands/postgres/postgres_context.py +0 -272
  136. adam/commands/postgres/psql_completions.py +0 -10
  137. kaqing-2.0.172.dist-info/RECORD +0 -230
  138. {kaqing-2.0.172.dist-info → kaqing-2.0.186.dist-info}/WHEEL +0 -0
  139. {kaqing-2.0.172.dist-info → kaqing-2.0.186.dist-info}/entry_points.txt +0 -0
  140. {kaqing-2.0.172.dist-info → kaqing-2.0.186.dist-info}/top_level.txt +0 -0
@@ -1,11 +1,15 @@
1
+ from collections.abc import Callable
1
2
  from datetime import datetime
3
+ from functools import partial
2
4
 
3
- from adam.commands.command import Command
5
+ from adam.commands import validate_args
6
+ from adam.commands.command import Command, InvalidArgumentsException
7
+ from adam.commands.medusa.utils_medusa import medusa_backup_names
4
8
  from adam.utils_k8s.statefulsets import StatefulSets
5
9
  from adam.repl_state import ReplState, RequiredState
6
10
  from adam.utils_k8s.custom_resources import CustomResources
7
11
  from adam.config import Config
8
- from adam.utils import lines_to_tabular, log2, log_exc
12
+ from adam.utils import tabulize, log2, log_exc
9
13
 
10
14
  class MedusaRestore(Command):
11
15
  COMMAND = 'restore'
@@ -35,45 +39,48 @@ class MedusaRestore(Command):
35
39
  if not dc:
36
40
  return state
37
41
 
38
- if len(args) == 1:
39
- bkname = args[0]
40
- job = CustomResources.medusa_get_backupjob(dc, ns, bkname)
41
- if not job:
42
+ def msg(missing: bool):
43
+ if missing:
44
+ log2('\n* Missing Backup Name')
45
+ log2('Usage: qing restore <backup> <sts@name_space>\n')
46
+ else:
42
47
  log2('\n* Backup job name is not valid.')
43
- bklist = [f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '')}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]
44
- log2(lines_to_tabular(bklist, 'NAME\tCREATED\tFINISHED', separator='\t'))
45
48
 
46
- return state
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)
54
+
55
+ with validate_args(args, state, msg=partial(msg, True)) as bkname:
56
+ if not (job := CustomResources.medusa_get_backupjob(dc, ns, bkname)):
57
+ msg(False)
58
+ raise InvalidArgumentsException()
47
59
 
48
60
  if not input(f"Restoring from {bkname} created at {job['metadata']['creationTimestamp']}. Please enter Yes to continue: ").lower() in ['y', 'yes']:
49
61
  return state
50
- else:
51
- bklist = [f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '')}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]
52
- log2('\n* Missing Backup Name')
53
- log2('Usage: qing medusa restore <backup> <sts@name_space>\n')
54
- log2(lines_to_tabular(bklist, 'NAME\tCREATED\tFINISHED', separator='\t'))
55
- return state
56
62
 
57
- now_dtformat = datetime.now().strftime("%Y-%m-%d.%H.%M.%S")
58
- rtname = 'medusa-' + now_dtformat + '-restore-from-' + bkname
59
- with log_exc(lambda e: "Exception: MedusaRestore failed: %s\n" % e):
60
- CustomResources.create_medusa_restorejob(rtname, bkname, dc, ns)
63
+ with log_exc(lambda e: "Exception: MedusaRestore failed: %s\n" % e):
64
+ now_dtformat = datetime.now().strftime("%Y-%m-%d.%H.%M.%S")
65
+ rtname = 'medusa-' + now_dtformat + '-restore-from-' + bkname
66
+ CustomResources.create_medusa_restorejob(rtname, bkname, dc, ns)
61
67
 
62
- return state
68
+ return state
63
69
 
64
70
  def completion(self, state: ReplState):
65
- if sc := super().completion(state):
66
- ns = state.namespace
67
- dc: str = StatefulSets.get_datacenter(state.sts, ns)
68
- if not dc:
69
- return {}
70
-
71
- if Config().get('medusa.restore-auto-complete', False):
72
- leaf = {id: None for id in [f"{x['metadata']['name']}" for x in CustomResources.medusa_show_backupjobs(dc, ns)]}
73
-
74
- return super().completion(state, leaf)
75
- else:
76
- 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
77
84
 
78
85
  return {}
79
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'
@@ -1,7 +1,8 @@
1
+ from adam.commands import validate_args
1
2
  from adam.commands.command import Command
2
3
  from adam.config import Config
3
4
  from adam.repl_state import ReplState
4
- from adam.utils import lines_to_tabular, log, log2
5
+ from adam.utils import tabulize, log, log2
5
6
 
6
7
  class GetParam(Command):
7
8
  COMMAND = 'get'
@@ -23,19 +24,16 @@ class GetParam(Command):
23
24
  return super().run(cmd, state)
24
25
 
25
26
  with self.validate(args, state) as (args, state):
26
- if len(args) < 1:
27
- lines = [f'{key}\t{Config().get(key, None)}' for key in Config().keys()]
28
- log(lines_to_tabular(lines, separator='\t'))
27
+ def msg():
28
+ tabulize(Config().keys(), lambda key: f'{key}\t{Config().get(key, None)}', separator='\t')
29
29
 
30
- return state
30
+ with validate_args(args, state, msg=msg) as key:
31
+ if v := Config().get(key, None):
32
+ log(v)
33
+ else:
34
+ log2(f'{key} is not set.')
31
35
 
32
- key = args[0]
33
- if v := Config().get(key, None):
34
- log(v)
35
- else:
36
- log2(f'{key} is not set.')
37
-
38
- return v if v else state
36
+ return v if v else state
39
37
 
40
38
  def completion(self, _: ReplState):
41
39
  return {GetParam.COMMAND: {key: None for key in Config().keys()}}
@@ -1,3 +1,4 @@
1
+ from adam.commands import validate_args
1
2
  from adam.commands.command import Command
2
3
  from adam.config import Config
3
4
  from adam.repl_state import ReplState
@@ -23,18 +24,14 @@ class SetParam(Command):
23
24
  return super().run(cmd, state)
24
25
 
25
26
  with self.validate(args, state) as (args, state):
26
- if len(args) < 2:
27
- log2('set <key> <value>')
27
+ with validate_args(args, state, exactly=2, msg=lambda: log2('set <key> <value>')):
28
+ key = args[0]
29
+ value = args[1]
30
+ Config().set(key, value)
28
31
 
29
- return 'invalid args'
32
+ log(Config().get(key, None))
30
33
 
31
- key = args[0]
32
- value = args[1]
33
- Config().set(key, value)
34
-
35
- log(Config().get(key, None))
36
-
37
- return value
34
+ return value
38
35
 
39
36
  def completion(self, _: ReplState):
40
37
  return {SetParam.COMMAND: {key: ({'true': None, 'false': None} if Config().get(key, None) in [True, False] else None) for key in Config().keys()}}
@@ -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
+ }
@@ -1,15 +1,15 @@
1
1
  import click
2
2
 
3
- from adam.commands import extract_trailing_options
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'
@@ -32,55 +32,40 @@ class Postgres(IntermediateCommand):
32
32
 
33
33
  with self.validate(args, state) as (args, state):
34
34
  with extract_trailing_options(args, '&') as (args, backgrounded):
35
- if not args:
36
- if state.in_repl:
37
- log2('Please use SQL statement. e.g. pg \l')
38
- else:
39
- log2('* Command or SQL statements is missing.')
40
- Command.display_help()
35
+ with validate_args(args, state, name='SQL statement') as sql:
36
+ if not state.pg_path:
37
+ if state.in_repl:
38
+ log2('Enter "use <pg-name>" first.')
39
+ else:
40
+ log2('* pg-name is missing.')
41
41
 
42
- return 'command-missing'
42
+ return state
43
43
 
44
- if not state.pg_path:
45
44
  if state.in_repl:
46
- log2('Enter "use <pg-name>" first.')
47
- else:
48
- log2('* pg-name is missing.')
45
+ with postgres(state) as pod:
46
+ pod.sql(args, backgrounded=backgrounded)
47
+ elif not self.run_subcommand(cmd, state):
48
+ with postgres(state) as pod:
49
+ pod.sql(args, backgrounded=backgrounded)
49
50
 
50
51
  return state
51
52
 
52
- if state.in_repl:
53
- with postgres(state) as pod:
54
- pod.sql(args, background=backgrounded)
55
- elif not self.run_subcommand(cmd, state):
56
- with postgres(state) as pod:
57
- pod.sql(args, background=backgrounded)
58
-
59
- return state
60
-
61
53
  def cmd_list(self):
62
54
  return [PostgresLs(), PostgresPreview(), PostgresPg()]
63
55
 
64
56
  def completion(self, state: ReplState):
65
57
  if state.device != state.P:
66
- # conflicts with cql completions
67
58
  return {}
68
59
 
69
- leaf = {}
70
- session = PostgresContext.apply(state.namespace, state.pg_path)
71
- if session.db:
72
- if pg_table_names(state.namespace, state.pg_path):
73
- leaf = psql_completions(state.namespace, state.pg_path)
74
- elif state.pg_path:
75
- leaf = {
76
- '\h': None,
77
- '\l': None,
78
- }
79
-
80
- if state.pg_path:
81
- return super().completion(state, leaf) | leaf
82
- else:
83
- 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 {}
84
69
 
85
70
  def help(self, _: ReplState):
86
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)