kaqing 2.0.110__py3-none-any.whl → 2.0.214__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 (251) hide show
  1. adam/__init__.py +0 -2
  2. adam/app_session.py +9 -12
  3. adam/apps.py +18 -4
  4. adam/batch.py +19 -19
  5. adam/checks/check_utils.py +16 -46
  6. adam/checks/cpu.py +7 -1
  7. adam/checks/cpu_metrics.py +52 -0
  8. adam/checks/disk.py +2 -3
  9. adam/columns/columns.py +3 -1
  10. adam/columns/cpu.py +3 -1
  11. adam/columns/cpu_metrics.py +22 -0
  12. adam/columns/memory.py +3 -4
  13. adam/commands/__init__.py +24 -0
  14. adam/commands/app/app.py +38 -0
  15. adam/commands/{app_ping.py → app/app_ping.py} +7 -13
  16. adam/commands/{login.py → app/login.py} +22 -24
  17. adam/commands/app/show_app_actions.py +49 -0
  18. adam/commands/{show → app}/show_app_id.py +8 -11
  19. adam/commands/{show → app}/show_app_queues.py +7 -14
  20. adam/commands/app/show_login.py +56 -0
  21. adam/commands/app/utils_app.py +106 -0
  22. adam/commands/audit/audit.py +22 -40
  23. adam/commands/audit/audit_repair_tables.py +15 -19
  24. adam/commands/audit/audit_run.py +15 -22
  25. adam/commands/audit/completions_l.py +15 -0
  26. adam/commands/audit/show_last10.py +4 -18
  27. adam/commands/audit/show_slow10.py +4 -17
  28. adam/commands/audit/show_top10.py +4 -16
  29. adam/commands/audit/utils_show_top10.py +15 -3
  30. adam/commands/bash/__init__.py +5 -0
  31. adam/commands/bash/bash.py +36 -0
  32. adam/commands/bash/bash_completer.py +93 -0
  33. adam/commands/bash/utils_bash.py +16 -0
  34. adam/commands/cassandra/__init__.py +0 -0
  35. adam/commands/cassandra/download_cassandra_log.py +45 -0
  36. adam/commands/{restart.py → cassandra/restart_cluster.py} +12 -26
  37. adam/commands/cassandra/restart_node.py +51 -0
  38. adam/commands/cassandra/restart_nodes.py +47 -0
  39. adam/commands/{rollout.py → cassandra/rollout.py} +20 -25
  40. adam/commands/cassandra/show_cassandra_repairs.py +37 -0
  41. adam/commands/cassandra/show_cassandra_status.py +117 -0
  42. adam/commands/{show → cassandra}/show_cassandra_version.py +5 -18
  43. adam/commands/cassandra/show_processes.py +50 -0
  44. adam/commands/cassandra/show_storage.py +44 -0
  45. adam/commands/{watch.py → cassandra/watch.py} +26 -29
  46. adam/commands/cli/__init__.py +0 -0
  47. adam/commands/{cli_commands.py → cli/cli_commands.py} +8 -4
  48. adam/commands/cli/clipboard_copy.py +86 -0
  49. adam/commands/cli/show_cli_commands.py +56 -0
  50. adam/commands/code.py +57 -0
  51. adam/commands/command.py +211 -40
  52. adam/commands/commands_utils.py +20 -27
  53. adam/commands/config/__init__.py +0 -0
  54. adam/commands/{param_get.py → config/param_get.py} +11 -14
  55. adam/commands/{param_set.py → config/param_set.py} +8 -12
  56. adam/commands/{show → config}/show_params.py +2 -5
  57. adam/commands/cql/alter_tables.py +66 -0
  58. adam/commands/cql/completions_c.py +29 -0
  59. adam/commands/cql/cqlsh.py +10 -32
  60. adam/commands/cql/utils_cql.py +306 -0
  61. adam/commands/debug/__init__.py +0 -0
  62. adam/commands/debug/debug.py +22 -0
  63. adam/commands/debug/debug_completes.py +35 -0
  64. adam/commands/debug/debug_timings.py +35 -0
  65. adam/commands/debug/show_offloaded_completes.py +45 -0
  66. adam/commands/deploy/code_start.py +7 -10
  67. adam/commands/deploy/code_stop.py +4 -21
  68. adam/commands/deploy/code_utils.py +3 -3
  69. adam/commands/deploy/deploy.py +4 -27
  70. adam/commands/deploy/deploy_frontend.py +14 -17
  71. adam/commands/deploy/deploy_pg_agent.py +3 -6
  72. adam/commands/deploy/deploy_pod.py +65 -73
  73. adam/commands/deploy/deploy_utils.py +14 -24
  74. adam/commands/deploy/undeploy.py +4 -27
  75. adam/commands/deploy/undeploy_frontend.py +4 -7
  76. adam/commands/deploy/undeploy_pg_agent.py +6 -8
  77. adam/commands/deploy/undeploy_pod.py +11 -12
  78. adam/commands/devices/__init__.py +0 -0
  79. adam/commands/devices/device.py +149 -0
  80. adam/commands/devices/device_app.py +163 -0
  81. adam/commands/devices/device_auit_log.py +49 -0
  82. adam/commands/devices/device_cass.py +179 -0
  83. adam/commands/devices/device_export.py +87 -0
  84. adam/commands/devices/device_postgres.py +160 -0
  85. adam/commands/devices/devices.py +25 -0
  86. adam/commands/diag/__init__.py +0 -0
  87. adam/commands/{check.py → diag/check.py} +16 -25
  88. adam/commands/diag/generate_report.py +52 -0
  89. adam/commands/diag/issues.py +43 -0
  90. adam/commands/exit.py +1 -4
  91. adam/commands/export/__init__.py +0 -0
  92. adam/commands/export/clean_up_all_export_sessions.py +37 -0
  93. adam/commands/export/clean_up_export_sessions.py +39 -0
  94. adam/commands/export/completions_x.py +11 -0
  95. adam/commands/export/download_export_session.py +40 -0
  96. adam/commands/export/drop_export_database.py +39 -0
  97. adam/commands/export/drop_export_databases.py +37 -0
  98. adam/commands/export/export.py +37 -0
  99. adam/commands/export/export_databases.py +251 -0
  100. adam/commands/export/export_select.py +34 -0
  101. adam/commands/export/export_sessions.py +210 -0
  102. adam/commands/export/export_use.py +49 -0
  103. adam/commands/export/export_x_select.py +48 -0
  104. adam/commands/export/exporter.py +419 -0
  105. adam/commands/export/import_files.py +44 -0
  106. adam/commands/export/import_session.py +40 -0
  107. adam/commands/export/importer.py +81 -0
  108. adam/commands/export/importer_athena.py +157 -0
  109. adam/commands/export/importer_sqlite.py +78 -0
  110. adam/commands/export/show_column_counts.py +45 -0
  111. adam/commands/export/show_export_databases.py +39 -0
  112. adam/commands/export/show_export_session.py +39 -0
  113. adam/commands/export/show_export_sessions.py +37 -0
  114. adam/commands/export/utils_export.py +366 -0
  115. adam/commands/fs/__init__.py +0 -0
  116. adam/commands/fs/cat.py +36 -0
  117. adam/commands/fs/cat_local.py +42 -0
  118. adam/commands/fs/cd.py +41 -0
  119. adam/commands/fs/download_file.py +47 -0
  120. adam/commands/fs/find_files.py +51 -0
  121. adam/commands/fs/find_processes.py +76 -0
  122. adam/commands/fs/head.py +36 -0
  123. adam/commands/fs/ls.py +41 -0
  124. adam/commands/fs/ls_local.py +40 -0
  125. adam/commands/fs/pwd.py +45 -0
  126. adam/commands/fs/rm.py +18 -0
  127. adam/commands/fs/rm_downloads.py +39 -0
  128. adam/commands/fs/rm_logs.py +38 -0
  129. adam/commands/{shell.py → fs/shell.py} +12 -4
  130. adam/commands/{show → fs}/show_adam.py +3 -3
  131. adam/commands/{show → fs}/show_host.py +1 -1
  132. adam/commands/help.py +5 -3
  133. adam/commands/intermediate_command.py +52 -0
  134. adam/commands/kubectl.py +38 -0
  135. adam/commands/medusa/medusa.py +4 -22
  136. adam/commands/medusa/medusa_backup.py +20 -27
  137. adam/commands/medusa/medusa_restore.py +35 -48
  138. adam/commands/medusa/medusa_show_backupjobs.py +16 -18
  139. adam/commands/medusa/medusa_show_restorejobs.py +13 -18
  140. adam/commands/medusa/utils_medusa.py +15 -0
  141. adam/commands/nodetool/__init__.py +0 -0
  142. adam/commands/{nodetool.py → nodetool/nodetool.py} +9 -20
  143. adam/commands/postgres/completions_p.py +22 -0
  144. adam/commands/postgres/postgres.py +47 -55
  145. adam/commands/postgres/postgres_databases.py +269 -0
  146. adam/commands/postgres/postgres_ls.py +5 -9
  147. adam/commands/postgres/postgres_preview.py +5 -9
  148. adam/commands/postgres/utils_postgres.py +80 -0
  149. adam/commands/preview_table.py +8 -44
  150. adam/commands/reaper/reaper.py +4 -27
  151. adam/commands/reaper/reaper_forward.py +49 -56
  152. adam/commands/reaper/reaper_forward_session.py +6 -0
  153. adam/commands/reaper/reaper_forward_stop.py +10 -16
  154. adam/commands/reaper/reaper_restart.py +7 -14
  155. adam/commands/reaper/reaper_run_abort.py +8 -33
  156. adam/commands/reaper/reaper_runs.py +43 -58
  157. adam/commands/reaper/reaper_runs_abort.py +29 -49
  158. adam/commands/reaper/reaper_schedule_activate.py +14 -33
  159. adam/commands/reaper/reaper_schedule_start.py +9 -33
  160. adam/commands/reaper/reaper_schedule_stop.py +9 -33
  161. adam/commands/reaper/reaper_schedules.py +4 -14
  162. adam/commands/reaper/reaper_status.py +8 -16
  163. adam/commands/reaper/utils_reaper.py +203 -0
  164. adam/commands/repair/repair.py +4 -22
  165. adam/commands/repair/repair_log.py +5 -11
  166. adam/commands/repair/repair_run.py +27 -34
  167. adam/commands/repair/repair_scan.py +32 -40
  168. adam/commands/repair/repair_stop.py +5 -12
  169. adam/commands/show.py +40 -0
  170. adam/config.py +5 -15
  171. adam/embedded_params.py +1 -1
  172. adam/log.py +4 -4
  173. adam/repl.py +83 -116
  174. adam/repl_commands.py +86 -45
  175. adam/repl_session.py +9 -1
  176. adam/repl_state.py +176 -40
  177. adam/sql/async_executor.py +62 -0
  178. adam/sql/lark_completer.py +286 -0
  179. adam/sql/lark_parser.py +604 -0
  180. adam/sql/qingl.lark +1076 -0
  181. adam/sql/sql_completer.py +52 -27
  182. adam/sql/sql_state_machine.py +131 -19
  183. adam/sso/authn_ad.py +6 -8
  184. adam/sso/authn_okta.py +4 -6
  185. adam/sso/cred_cache.py +4 -9
  186. adam/sso/idp.py +9 -12
  187. adam/utils.py +670 -31
  188. adam/utils_athena.py +145 -0
  189. adam/utils_audits.py +12 -103
  190. adam/utils_issues.py +32 -0
  191. adam/utils_k8s/app_clusters.py +35 -0
  192. adam/utils_k8s/app_pods.py +41 -0
  193. adam/utils_k8s/cassandra_clusters.py +35 -20
  194. adam/utils_k8s/cassandra_nodes.py +15 -6
  195. adam/utils_k8s/custom_resources.py +16 -17
  196. adam/utils_k8s/ingresses.py +2 -2
  197. adam/utils_k8s/jobs.py +7 -11
  198. adam/utils_k8s/k8s.py +96 -0
  199. adam/utils_k8s/kube_context.py +3 -6
  200. adam/{pod_exec_result.py → utils_k8s/pod_exec_result.py} +13 -4
  201. adam/utils_k8s/pods.py +159 -89
  202. adam/utils_k8s/secrets.py +4 -4
  203. adam/utils_k8s/service_accounts.py +5 -4
  204. adam/utils_k8s/services.py +2 -2
  205. adam/utils_k8s/statefulsets.py +6 -14
  206. adam/utils_local.py +80 -0
  207. adam/utils_net.py +4 -4
  208. adam/utils_repl/__init__.py +0 -0
  209. adam/utils_repl/appendable_completer.py +6 -0
  210. adam/utils_repl/automata_completer.py +48 -0
  211. adam/utils_repl/repl_completer.py +93 -0
  212. adam/utils_repl/state_machine.py +173 -0
  213. adam/utils_sqlite.py +132 -0
  214. adam/version.py +1 -1
  215. {kaqing-2.0.110.dist-info → kaqing-2.0.214.dist-info}/METADATA +1 -1
  216. kaqing-2.0.214.dist-info/RECORD +272 -0
  217. kaqing-2.0.214.dist-info/top_level.txt +2 -0
  218. teddy/__init__.py +0 -0
  219. teddy/lark_parser.py +436 -0
  220. teddy/lark_parser2.py +618 -0
  221. adam/commands/alter_tables.py +0 -81
  222. adam/commands/app.py +0 -67
  223. adam/commands/bash.py +0 -150
  224. adam/commands/cd.py +0 -125
  225. adam/commands/cp.py +0 -95
  226. adam/commands/cql/cql_completions.py +0 -15
  227. adam/commands/cql/cql_utils.py +0 -112
  228. adam/commands/devices.py +0 -118
  229. adam/commands/issues.py +0 -75
  230. adam/commands/logs.py +0 -40
  231. adam/commands/ls.py +0 -146
  232. adam/commands/postgres/postgres_context.py +0 -239
  233. adam/commands/postgres/postgres_utils.py +0 -31
  234. adam/commands/postgres/psql_completions.py +0 -10
  235. adam/commands/pwd.py +0 -77
  236. adam/commands/reaper/reaper_session.py +0 -159
  237. adam/commands/report.py +0 -63
  238. adam/commands/show/show.py +0 -54
  239. adam/commands/show/show_app_actions.py +0 -56
  240. adam/commands/show/show_cassandra_status.py +0 -128
  241. adam/commands/show/show_commands.py +0 -61
  242. adam/commands/show/show_login.py +0 -63
  243. adam/commands/show/show_processes.py +0 -53
  244. adam/commands/show/show_repairs.py +0 -47
  245. adam/commands/show/show_storage.py +0 -52
  246. kaqing-2.0.110.dist-info/RECORD +0 -187
  247. kaqing-2.0.110.dist-info/top_level.txt +0 -1
  248. /adam/commands/{show → app}/__init__.py +0 -0
  249. /adam/commands/{nodetool_commands.py → nodetool/nodetool_commands.py} +0 -0
  250. {kaqing-2.0.110.dist-info → kaqing-2.0.214.dist-info}/WHEEL +0 -0
  251. {kaqing-2.0.110.dist-info → kaqing-2.0.214.dist-info}/entry_points.txt +0 -0
adam/utils_k8s/jobs.py CHANGED
@@ -1,7 +1,7 @@
1
1
  from kubernetes import client
2
2
  from time import sleep
3
3
  from .pods import Pods
4
- from adam.utils import log2
4
+ from adam.utils import log2, log_exc
5
5
 
6
6
  # utility collection on jobs; methods are all static
7
7
  class Jobs:
@@ -20,11 +20,10 @@ class Jobs:
20
20
  metadata=client.V1ObjectMeta(name=job_name),
21
21
  spec=spec)
22
22
 
23
- try:
23
+ with log_exc(lambda e: "Exception when calling BatchV1Apii->create_namespaced_job: %s\n" % e):
24
24
  client.BatchV1Api().create_namespaced_job(body=job, namespace=namespace)
25
25
  log2(f"Job {job_name} created in {namespace}")
26
- except Exception as e:
27
- log2("Exception when calling BatchV1Apii->create_namespaced_job: %s\n" % e)
26
+
28
27
  return
29
28
 
30
29
  def get_job_pods(job_name: str, namespace: str):
@@ -32,7 +31,7 @@ class Jobs:
32
31
  return pods
33
32
 
34
33
  def delete(job_name: str, namespace: str, wait=True):
35
- try:
34
+ with log_exc(lambda e: "Exception when calling BatchV1Apii->delete_namespaced_job: %s\n" % e):
36
35
  client.BatchV1Api().delete_namespaced_job(name=job_name, namespace=namespace, propagation_policy='Background')
37
36
  if wait:
38
37
  while True:
@@ -41,14 +40,11 @@ class Jobs:
41
40
  return
42
41
  sleep(5)
43
42
  log2(f"Job {job_name} in {namespace} deleted.")
44
- except Exception as e:
45
- log2("Exception when calling BatchV1Apii->delete_namespaced_job: %s\n" % e)
43
+
46
44
  return
47
45
 
48
46
  def get_logs(job_name: str, namespace: str):
49
47
  v1 = client.CoreV1Api()
50
- try:
48
+ with log_exc(lambda e: "Exception when calling CorV1Apii->list_namespaced_pod, cannot find job pod: %s\n" % e):
51
49
  pod_name = Jobs.get_job_pods(job_name, namespace).items[0].metadata.name
52
- log2(v1.read_namespaced_pod_log(name=pod_name, namespace=namespace))
53
- except Exception as e:
54
- log2("Exception when calling CorV1Apii->list_namespaced_pod, cannot find job pod: %s\n" % e)
50
+ log2(v1.read_namespaced_pod_log(name=pod_name, namespace=namespace))
adam/utils_k8s/k8s.py ADDED
@@ -0,0 +1,96 @@
1
+ from collections.abc import Callable
2
+ import inspect
3
+ import re
4
+ import portforward
5
+
6
+ from adam.commands.command import InvalidStateException
7
+ from adam.repl_state import ReplState
8
+ from adam.utils import log2
9
+ from adam.utils_k8s.kube_context import KubeContext
10
+
11
+ class PortForwardHandler:
12
+ connections: dict[str, int] = {}
13
+
14
+ def __init__(self, state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
15
+ self.state = state
16
+ self.local_port = local_port
17
+ self.svc_or_pod = svc_or_pod
18
+ self.target_port = target_port
19
+ self.forward_connection = None
20
+ self.pod = None
21
+
22
+ def __enter__(self) -> tuple[str, str]:
23
+ state = self.state
24
+
25
+ if not self.svc_or_pod:
26
+ log2('No service or pod found.')
27
+
28
+ raise InvalidStateException(state)
29
+
30
+ if KubeContext.in_cluster():
31
+ svc_name = self.svc_or_pod(True)
32
+ if not svc_name:
33
+ log2('No service found.')
34
+
35
+ raise InvalidStateException(state)
36
+
37
+ # cs-a526330d23-cs-a526330d23-default-sts-0 ->
38
+ # curl http://cs-a526330d23-cs-a526330d23-reaper-service.stgawsscpsr.svc.cluster.local:8080
39
+ groups = re.match(r'^(.*?-.*?-.*?-.*?-).*', state.sts)
40
+ if groups:
41
+ svc = f'{groups[1]}{svc_name}.{state.namespace}.svc.cluster.local:{self.target_port}'
42
+ return (svc, svc)
43
+ else:
44
+ raise InvalidStateException(state)
45
+ else:
46
+ pod = self.svc_or_pod(False)
47
+ if not pod:
48
+ log2('No pod found.')
49
+
50
+ raise InvalidStateException(state)
51
+
52
+ self.pod = pod
53
+
54
+ # pf = portforward.forward(state.namespace, pod, self.local_port + 1, self.target_port, log_level=portforward.LogLevel.DEBUG)
55
+ # print(inspect.getsource(pf.__enter__))
56
+ # print('test portforward START', state.namespace, pod, self.local_port + 1, self.target_port, pf.__enter__)
57
+ # with pf:
58
+ # print('test portforward BODY')
59
+ # print('test portforward OK')
60
+
61
+ self.forward_connection = portforward.forward(state.namespace, pod, self.local_port, self.target_port)
62
+ if self.inc_connection_cnt() == 1:
63
+ self.forward_connection.__enter__()
64
+
65
+ return (f'localhost:{self.local_port}', f'{pod}:{self.target_port}')
66
+
67
+ def __exit__(self, exc_type, exc_val, exc_tb):
68
+ if self.forward_connection:
69
+ if not self.dec_connection_cnt():
70
+ return self.forward_connection.__exit__(exc_type, exc_val, exc_tb)
71
+
72
+ return False
73
+
74
+ def inc_connection_cnt(self):
75
+ id = self.connection_id(self.pod)
76
+ if id not in PortForwardHandler.connections:
77
+ PortForwardHandler.connections[id] = 1
78
+ else:
79
+ PortForwardHandler.connections[id] += 1
80
+
81
+ return PortForwardHandler.connections[id]
82
+
83
+ def dec_connection_cnt(self):
84
+ id = self.connection_id(self.pod)
85
+ if id not in PortForwardHandler.connections:
86
+ PortForwardHandler.connections[id] = 0
87
+ elif PortForwardHandler.connections[id] > 0:
88
+ PortForwardHandler.connections[id] -= 1
89
+
90
+ return PortForwardHandler.connections[id]
91
+
92
+ def connection_id(self, pod: str):
93
+ return f'{self.local_port}:{pod}:{self.target_port}'
94
+
95
+ def port_forwarding(state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
96
+ return PortForwardHandler(state, local_port, svc_or_pod, target_port)
@@ -3,7 +3,7 @@ import re
3
3
  from kubernetes import config as kconfig
4
4
 
5
5
  from adam.config import Config
6
- from adam.utils import idp_token_from_env, lines_to_tabular, log2
6
+ from adam.utils import idp_token_from_env, log2, tabulize
7
7
 
8
8
  class KubeContext:
9
9
  _in_cluster = False
@@ -56,7 +56,7 @@ class KubeContext:
56
56
  log2('Use -v <key>=<value> format.')
57
57
  log2()
58
58
  lines = [f'{key}\t{Config().get(key, None)}' for key in Config().keys()]
59
- log2(lines_to_tabular(lines, separator='\t'))
59
+ tabulize(lines, separator='\t', to=2)
60
60
 
61
61
  for p in param_ovrs:
62
62
  tokens = p.split('=')
@@ -102,7 +102,4 @@ class KubeContext:
102
102
  return name if re.match(r"^(?!pg-).*-k8spg-.*$", name) else None
103
103
 
104
104
  def show_out(s: bool):
105
- return s or Config().is_debug()
106
-
107
- def show_parallelism():
108
- return Config().get('debugs.show-parallelism', False)
105
+ return s or Config().is_debug()
@@ -1,6 +1,8 @@
1
1
  import yaml
2
2
 
3
- class PodExecResult:
3
+ from adam.utils import ExecResult, log_exc
4
+
5
+ class PodExecResult(ExecResult):
4
6
  # {
5
7
  # 'metadata': {},
6
8
  # 'status': 'Failure',
@@ -27,12 +29,19 @@ class PodExecResult:
27
29
  def exit_code(self) -> int:
28
30
  code = 0
29
31
 
30
- try:
32
+ with log_exc(False):
31
33
  code = self.error['details']['causes'][0]['message']
32
- except:
33
- pass
34
34
 
35
35
  return code
36
36
 
37
+ def cat_log_file_cmd(self):
38
+ if self.pod and self.log_file:
39
+ return f'@{self.pod} cat {self.log_file}'
40
+
41
+ return None
42
+
43
+ def __str__(self):
44
+ return f'{"OK" if self.exit_code() == 0 else self.exit_code()} {self.command}'
45
+
37
46
  def __audit_extra__(self):
38
47
  return self.log_file if self.log_file else None
adam/utils_k8s/pods.py CHANGED
@@ -1,24 +1,32 @@
1
1
  from collections.abc import Callable
2
- from concurrent.futures import ThreadPoolExecutor, as_completed
3
2
  from datetime import datetime
3
+ import os
4
+ import re
5
+ import subprocess
4
6
  import sys
5
7
  import time
6
- from typing import TypeVar, cast
8
+ from typing import TypeVar
7
9
  from kubernetes import client
8
10
  from kubernetes.stream import stream
9
- from kubernetes.stream.ws_client import ERROR_CHANNEL
11
+ from kubernetes.stream.ws_client import ERROR_CHANNEL, WSClient
10
12
 
11
13
  from adam.config import Config
14
+ from adam.repl_session import ReplSession
12
15
  from adam.utils_k8s.volumes import ConfigMapMount
13
- from adam.pod_exec_result import PodExecResult
14
- from adam.utils import elapsed_time, log2
16
+ from adam.utils_k8s.pod_exec_result import PodExecResult
17
+ from adam.utils import GeneratorStream, ParallelMapHandler, log2, debug, log_dir, log_exc
18
+ from adam.utils_local import local_downloads_dir
15
19
  from .kube_context import KubeContext
16
20
 
21
+ from websocket._core import WebSocket
22
+
17
23
  T = TypeVar('T')
18
24
  _TEST_POD_EXEC_OUTS: PodExecResult = None
19
25
 
20
26
  # utility collection on pods; methods are all static
21
27
  class Pods:
28
+ _TEST_POD_CLOSE_SOCKET: bool = False
29
+
22
30
  def set_test_pod_exec_outs(outs: PodExecResult):
23
31
  global _TEST_POD_EXEC_OUTS
24
32
  _TEST_POD_EXEC_OUTS = outs
@@ -26,11 +34,9 @@ class Pods:
26
34
  return _TEST_POD_EXEC_OUTS
27
35
 
28
36
  def delete(pod_name: str, namespace: str, grace_period_seconds: int = None):
29
- try:
37
+ with log_exc(lambda e: "Exception when calling CoreV1Api->delete_namespaced_pod: %s\n" % e):
30
38
  v1 = client.CoreV1Api()
31
39
  v1.delete_namespaced_pod(pod_name, namespace, grace_period_seconds=grace_period_seconds)
32
- except Exception as e:
33
- log2("Exception when calling CoreV1Api->delete_namespaced_pod: %s\n" % e)
34
40
 
35
41
  def delete_with_selector(namespace: str, label_selector: str, grace_period_seconds: int = None):
36
42
  v1 = client.CoreV1Api()
@@ -39,98 +45,85 @@ class Pods:
39
45
  for i in ret.items:
40
46
  v1.delete_namespaced_pod(name=i.metadata.name, namespace=namespace, grace_period_seconds=grace_period_seconds)
41
47
 
42
- def on_pods(pods: list[str],
43
- namespace: str,
44
- body: Callable[[ThreadPoolExecutor, str, str, bool], T],
45
- post: Callable[[T], T] = None,
46
- action: str = 'action',
47
- max_workers=0,
48
- show_out=True,
49
- on_any = False,
50
- background = False) -> list[T]:
51
- show_out = KubeContext.show_out(show_out)
52
-
48
+ def parallelize(collection: list, max_workers: int = 0, samples = sys.maxsize, msg: str = None, action: str = 'action'):
53
49
  if not max_workers:
54
50
  max_workers = Config().action_workers(action, 0)
55
- if not on_any and max_workers > 0:
56
- # if parallel, node sampling is suppressed
57
- if KubeContext.show_parallelism():
58
- log2(f'Executing on all nodes from statefulset in parallel...')
59
- start_time = time.time()
60
- try:
61
- with ThreadPoolExecutor(max_workers=max_workers) as executor:
62
- # disable stdout from the pod_exec, then show the output in a for loop
63
- futures = [body(executor, pod, namespace, show_out) for pod in pods]
64
- if len(futures) == 0:
65
- return cast(list[T], [])
66
-
67
- rs = [future.result() for future in as_completed(futures)]
68
- if post:
69
- rs = [post(r, show_out=show_out) for r in rs]
70
-
71
- return rs
72
- finally:
73
- if KubeContext.show_parallelism():
74
- log2(f"Parallel {action} elapsed time: {elapsed_time(start_time)} with {max_workers} workers")
75
- else:
76
- results: list[T] = []
77
-
78
- samples = 1 if on_any else Config().action_node_samples(action, sys.maxsize)
79
- l = min(len(pods), samples)
80
- adj = 'all'
81
- if l < len(pods):
82
- adj = f'{l} sample'
83
- if show_out:
84
- log2(f'Executing on {adj} nodes from statefulset...')
85
- for pod_name in pods:
86
- try:
87
- # disable stdout from the pod_exec, then show the output in a for loop
88
- result = body(None, pod_name, namespace, False)
89
- if post:
90
- result = post(result, show_out=show_out)
91
- results.append(result)
92
- if result:
93
- l -= 1
94
- if not l:
95
- break
96
- except Exception as e:
97
- log2(e)
98
-
99
- return results
100
-
101
- def exec(pod_name: str, container: str, namespace: str, command: str,
102
- show_out = True, throw_err = False, shell = '/bin/sh',
103
- background = False,
104
- interaction: Callable[[any, list[str]], any] = None):
51
+ if samples == sys.maxsize:
52
+ samples = Config().action_node_samples(action, sys.maxsize)
53
+
54
+ return ParallelMapHandler(collection, max_workers, samples = samples, msg = msg)
55
+
56
+ def exec(pod_name: str,
57
+ container: str,
58
+ namespace: str,
59
+ command: str,
60
+ show_out = True,
61
+ throw_err = False,
62
+ shell = '/bin/sh',
63
+ backgrounded = False,
64
+ log_file = None,
65
+ interaction: Callable[[any, list[str]], any] = None,
66
+ env_prefix: str = None):
105
67
  if _TEST_POD_EXEC_OUTS:
106
68
  return _TEST_POD_EXEC_OUTS
107
69
 
108
70
  show_out = KubeContext.show_out(show_out)
109
71
 
72
+ if backgrounded or command.endswith(' &'):
73
+ command = command.strip(' &')
74
+
75
+ log_all_file = None
76
+ log_pod_file = None
77
+ if log_file:
78
+ log_pod_file = Pods.log_file_from_template(log_file, pod_name=pod_name)
79
+ if (a := Pods.log_file_from_template(log_file, pod_name='all')) != log_file:
80
+ log_all_file = a
81
+ else:
82
+ log_pod_file = Pods.log_file(command, pod_name=pod_name)
83
+
84
+ if env_prefix:
85
+ command = f'{env_prefix} {command}'
86
+
87
+ command = command.replace('"', '\\"')
88
+ cmd = f'nohup kubectl exec {pod_name} -c {container} -- {shell} -c "{command} &" > {log_pod_file} 2>&1 &'
89
+ if log_all_file:
90
+ cmd = f'{cmd} >> {log_all_file}'
91
+
92
+ if show_out:
93
+ log2(cmd)
94
+
95
+ result = subprocess.run(cmd, capture_output=True, text=True, shell=True)
96
+
97
+ return PodExecResult(result.stdout, result.stderr, cmd, None, pod=pod_name, log_file=log_pod_file)
98
+
110
99
  api = client.CoreV1Api()
111
100
 
112
- log_file = None
113
101
  tty = True
114
102
  exec_command = [shell, '-c', command]
115
- if background or command.endswith(' &'):
116
- # should be false for starting a background process
117
- tty = False
118
-
119
- if Config().get('repl.background-process.auto-nohup', True):
120
- command = command.strip(' &')
121
- cmd_name = ''
122
- if command.startswith('nodetool '):
123
- cmd_name = f".{'_'.join(command.split(' ')[5:])}"
124
-
125
- log_file = f'/tmp/qing-{datetime.now().strftime("%d%H%M%S")}{cmd_name}.log'
126
- command = f"nohup {command} > {log_file} 2>&1 &"
127
- exec_command = [shell, '-c', command]
103
+ if env_prefix:
104
+ exec_command = [shell, '-c', f'{env_prefix} {command}']
105
+
106
+ # if backgrounded or command.endswith(' &'):
107
+ # # should be false for starting a background process
108
+ # tty = False
109
+
110
+ # if Config().get('repl.background-process.auto-nohup', True):
111
+ # command = command.strip(' &')
112
+ # cmd_name = ''
113
+ # if command.startswith('nodetool '):
114
+ # cmd_name = f".{'_'.join(command.split(' ')[5:])}"
115
+
116
+ # if not log_file:
117
+ # log_file = f'{log_prefix()}-{datetime.now().strftime("%d%H%M%S")}{cmd_name}.log'
118
+ # command = f"nohup {command} > {log_file} 2>&1 &"
119
+ # if env_prefix:
120
+ # command = f'{env_prefix} {command}'
121
+ # exec_command = [shell, '-c', command]
128
122
 
129
123
  k_command = f'kubectl exec {pod_name} -c {container} -n {namespace} -- {shell} -c "{command}"'
130
- if show_out:
131
- print(k_command)
124
+ debug(k_command)
132
125
 
133
- resp = stream(
126
+ resp: WSClient = stream(
134
127
  api.connect_get_namespaced_pod_exec,
135
128
  pod_name,
136
129
  namespace,
@@ -143,6 +136,7 @@ class Pods:
143
136
  _preload_content=False,
144
137
  )
145
138
 
139
+ s: WebSocket = resp.sock
146
140
  stdout = []
147
141
  stderr = []
148
142
  error_output = None
@@ -161,11 +155,9 @@ class Pods:
161
155
  stderr.append(frag)
162
156
  if show_out: print(frag, end="")
163
157
 
164
- try:
158
+ with log_exc():
165
159
  # get the exit code from server
166
160
  error_output = resp.read_channel(ERROR_CHANNEL)
167
- except Exception as e:
168
- pass
169
161
  except Exception as e:
170
162
  if throw_err:
171
163
  raise e
@@ -173,9 +165,87 @@ class Pods:
173
165
  log2(e)
174
166
  finally:
175
167
  resp.close()
168
+ if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
169
+ with log_exc():
170
+ s.sock.close()
176
171
 
177
172
  return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
178
173
 
174
+ def log_file(command: str, pod_name: str = None, dt: datetime = None):
175
+ cmd_name = ''
176
+ if command.startswith('nodetool '):
177
+ command = command.strip(' &')
178
+ cmd_name = f".{'_'.join(command.split(' ')[5:])}"
179
+
180
+ pod_suffix = '{pod}'
181
+ if pod_name:
182
+ pod_suffix = pod_name
183
+ if groups := re.match(r'.*-(.*)', pod_name):
184
+ pod_suffix = f'-{groups[1]}'
185
+
186
+ return f'{log_dir()}/{Pods.job_id()}{cmd_name}{pod_suffix}.log'
187
+
188
+ def job_id(dt: datetime = None):
189
+ if not dt:
190
+ dt = datetime.now()
191
+
192
+ return dt.strftime("%d%H%M%S")
193
+
194
+ def log_file_from_template(log_file: str, pod_name: str):
195
+ pod_suffix = pod_name
196
+ if pod_name and (groups := re.match(r'.*-(.*)', pod_name)):
197
+ pod_suffix = f'-{groups[1]}'
198
+
199
+ if not pod_suffix.startswith('-'):
200
+ pod_suffix = f'-{pod_suffix}'
201
+
202
+ return log_file.replace('{pod}', pod_suffix)
203
+
204
+ def read_file(pod_name: str, container: str, namespace: str, file_path: str):
205
+ v1 = client.CoreV1Api()
206
+
207
+ resp = stream(
208
+ v1.connect_get_namespaced_pod_exec,
209
+ name=pod_name,
210
+ namespace=namespace,
211
+ container=container,
212
+ command=["cat", file_path],
213
+ stderr=True, stdin=False,
214
+ stdout=True, tty=False,
215
+ _preload_content=False, # Important for streaming
216
+ )
217
+
218
+ s: WebSocket = resp.sock
219
+ try:
220
+ while resp.is_open():
221
+ resp.update(timeout=1)
222
+ if resp.peek_stdout():
223
+ yield resp.read_stdout()
224
+
225
+ with log_exc():
226
+ # get the exit code from server
227
+ error_output = resp.read_channel(ERROR_CHANNEL)
228
+ except Exception as e:
229
+ raise e
230
+ finally:
231
+ resp.close()
232
+ if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
233
+ with log_exc():
234
+ s.sock.close()
235
+
236
+ def download_file(pod_name: str, container: str, namespace: str, from_path: str, to_path: str = None):
237
+ if not to_path:
238
+ to_path = f'{local_downloads_dir()}/{os.path.basename(from_path)}'
239
+
240
+ bytes = Pods.read_file(pod_name, container, namespace, from_path)
241
+ with open(to_path, 'wb') as f:
242
+ for item in GeneratorStream(bytes):
243
+ f.write(item)
244
+
245
+ ReplSession().append_history(f':cat {to_path}')
246
+
247
+ return to_path
248
+
179
249
  def get_container(namespace: str, pod_name: str, container_name: str):
180
250
  pod = Pods.get(namespace, pod_name)
181
251
  if not pod:
adam/utils_k8s/secrets.py CHANGED
@@ -6,13 +6,13 @@ from kubernetes import client
6
6
  from kubernetes.client import V1Secret
7
7
 
8
8
  from adam.config import Config
9
- from adam.utils import log2
9
+ from adam.utils import log2, wait_log
10
10
 
11
11
  # utility collection on secrets; methods are all static
12
12
  class Secrets:
13
13
  @functools.lru_cache()
14
14
  def list_secrets(namespace: str = None, name_pattern: str = None):
15
- Config().wait_log('Inspecting Cassandra Instances...')
15
+ wait_log('Inspecting Cassandra Instances...')
16
16
 
17
17
  secrets_names = []
18
18
 
@@ -39,14 +39,14 @@ class Secrets:
39
39
 
40
40
  return secrets_names
41
41
 
42
- def get_user_pass(ss_name: str, namespace: str, secret_path: str = 'cql.secret'):
42
+ def get_user_pass(sts_or_pod_name: str, namespace: str, secret_path: str = 'cql.secret'):
43
43
  # cs-d0767a536f-cs-d0767a536f-default-sts ->
44
44
  # cs-d0767a536f-superuser
45
45
  # cs-d0767a536f-reaper-ui
46
46
  user = 'superuser'
47
47
  if secret_path == 'reaper.secret':
48
48
  user = 'reaper-ui'
49
- groups = re.match(Config().get(f'{secret_path}.cluster-regex', r'(.*?-.*?)-.*'), ss_name)
49
+ groups = re.match(Config().get(f'{secret_path}.cluster-regex', r'(.*?-.*?)-.*'), sts_or_pod_name)
50
50
  secret_name = Config().get(f'{secret_path}.name', '{cluster}-' + user).replace('{cluster}', groups[1], 1)
51
51
 
52
52
  secret = Secrets.get_data(namespace, secret_name)
@@ -1,6 +1,7 @@
1
1
  from kubernetes import client, config
2
2
 
3
3
  from adam.config import Config
4
+ from adam.utils import debug
4
5
 
5
6
  # utility collection on service accounts; methods are all static
6
7
  class ServiceAccounts:
@@ -37,7 +38,7 @@ class ServiceAccounts:
37
38
  namespace=namespace,
38
39
  body=service_account
39
40
  )
40
- Config().debug(f"Service Account '{api_response.metadata.name}' created in namespace '{namespace}'.")
41
+ debug(f"Service Account '{api_response.metadata.name}' created in namespace '{namespace}'.")
41
42
 
42
43
  def delete_service_account(namespace: str, label_selector: str) -> list:
43
44
  refs = []
@@ -45,7 +46,7 @@ class ServiceAccounts:
45
46
  v1 = client.CoreV1Api()
46
47
  sas = v1.list_namespaced_service_account(namespace=namespace, label_selector=label_selector).items
47
48
  for sa in sas:
48
- Config().debug(f'delete {sa.metadata.name}')
49
+ debug(f'delete {sa.metadata.name}')
49
50
  v1.delete_namespaced_service_account(name=sa.metadata.name, namespace=namespace)
50
51
  refs.append(sa)
51
52
 
@@ -102,7 +103,7 @@ class ServiceAccounts:
102
103
  v1_rbac = client.RbacAuthorizationV1Api()
103
104
  cluster_role_bindings = v1_rbac.list_namespaced_role_binding(namespace=namespace, label_selector=label_selector).items
104
105
  for binding in cluster_role_bindings:
105
- Config().debug(f'delete {binding.metadata.name}')
106
+ debug(f'delete {binding.metadata.name}')
106
107
  v1_rbac.delete_namespaced_role_binding(name=binding.metadata.name, namespace=namespace)
107
108
  refs.append(binding)
108
109
 
@@ -162,7 +163,7 @@ class ServiceAccounts:
162
163
  v1_rbac = client.RbacAuthorizationV1Api()
163
164
  cluster_role_bindings = v1_rbac.list_cluster_role_binding(label_selector=label_selector).items
164
165
  for binding in cluster_role_bindings:
165
- Config().debug(f'delete {binding.metadata.name}')
166
+ debug(f'delete {binding.metadata.name}')
166
167
  v1_rbac.delete_cluster_role_binding(binding.metadata.name)
167
168
  refs.append(binding)
168
169
 
@@ -2,7 +2,7 @@ from typing import List
2
2
  from kubernetes import client
3
3
 
4
4
  from adam.config import Config
5
- from adam.utils import log2
5
+ from adam.utils import debug, log2
6
6
 
7
7
  from .kube_context import KubeContext
8
8
 
@@ -71,7 +71,7 @@ class Services:
71
71
  namespace=namespace,
72
72
  body=delete_options
73
73
  )
74
- Config().debug(f"200 Service '{name}' in namespace '{namespace}' deleted successfully.")
74
+ debug(f"200 Service '{name}' in namespace '{namespace}' deleted successfully.")
75
75
  except client.ApiException as e:
76
76
  log2(f"Error deleting Service '{name}': {e}")
77
77
  except Exception as e:
@@ -1,12 +1,9 @@
1
- from collections.abc import Callable
2
- from concurrent.futures import ThreadPoolExecutor
3
1
  from datetime import datetime
4
2
  import functools
5
3
  import re
6
4
  from typing import List, TypeVar, cast
7
5
  from kubernetes import client
8
6
 
9
- from .pods import Pods
10
7
  from .kube_context import KubeContext
11
8
  from adam.utils import log2
12
9
 
@@ -58,18 +55,12 @@ class StatefulSets:
58
55
 
59
56
  return statefulset_pods
60
57
 
61
- def on_cluster(statefulset: str,
62
- namespace: str,
63
- body: Callable[[ThreadPoolExecutor, str, str, bool], T],
64
- post: Callable[[T], T] = None,
65
- action: str = 'action', max_workers=0, show_out=True, on_any = False, background = False) -> list[T]:
66
- pods = StatefulSets.pod_names(statefulset, namespace)
67
-
68
- return Pods.on_pods(pods, namespace, body, post=post, action=action, max_workers=max_workers, show_out=show_out, on_any=on_any, background=background)
69
-
70
58
  @functools.lru_cache()
71
- def pod_names(ss: str, ns: str):
72
- return [pod.metadata.name for pod in StatefulSets.pods(ss, ns)]
59
+ def pod_names(sts: str, ns: str):
60
+ if not sts:
61
+ return []
62
+
63
+ return [pod.metadata.name for pod in StatefulSets.pods(sts, ns)]
73
64
 
74
65
  def restarted_at(ss: str, ns: str):
75
66
  # returns timestamp and if being rolled out
@@ -92,6 +83,7 @@ class StatefulSets:
92
83
 
93
84
  return restarted, False
94
85
 
86
+ @functools.lru_cache()
95
87
  def get_datacenter(sts: str, ns: str) -> str:
96
88
  v1 = client.AppsV1Api()
97
89
  namespace = ns