kaqing 2.0.145__py3-none-any.whl → 2.0.174__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 (174) hide show
  1. adam/__init__.py +0 -2
  2. adam/app_session.py +8 -11
  3. adam/batch.py +3 -3
  4. adam/checks/check_utils.py +14 -46
  5. adam/checks/cpu.py +7 -1
  6. adam/checks/cpu_metrics.py +52 -0
  7. adam/checks/disk.py +2 -3
  8. adam/columns/columns.py +3 -1
  9. adam/columns/cpu.py +3 -1
  10. adam/columns/cpu_metrics.py +22 -0
  11. adam/columns/memory.py +3 -4
  12. adam/commands/__init__.py +22 -0
  13. adam/commands/alter_tables.py +33 -48
  14. adam/commands/audit/audit.py +22 -23
  15. adam/commands/audit/audit_repair_tables.py +14 -17
  16. adam/commands/audit/audit_run.py +15 -23
  17. adam/commands/audit/show_last10.py +10 -13
  18. adam/commands/audit/show_slow10.py +10 -13
  19. adam/commands/audit/show_top10.py +10 -13
  20. adam/commands/audit/utils_show_top10.py +2 -3
  21. adam/commands/bash/__init__.py +5 -0
  22. adam/commands/bash/bash.py +7 -104
  23. adam/commands/bash/utils_bash.py +16 -0
  24. adam/commands/cat.py +7 -23
  25. adam/commands/cd.py +7 -11
  26. adam/commands/check.py +14 -23
  27. adam/commands/cli_commands.py +2 -3
  28. adam/commands/code.py +20 -23
  29. adam/commands/command.py +152 -37
  30. adam/commands/commands_utils.py +8 -17
  31. adam/commands/cp.py +18 -32
  32. adam/commands/cql/cql_completions.py +11 -7
  33. adam/commands/cql/cqlsh.py +10 -30
  34. adam/commands/cql/{cql_utils.py → utils_cql.py} +147 -15
  35. adam/commands/deploy/code_start.py +7 -10
  36. adam/commands/deploy/code_stop.py +4 -21
  37. adam/commands/deploy/code_utils.py +3 -3
  38. adam/commands/deploy/deploy.py +4 -27
  39. adam/commands/deploy/deploy_frontend.py +14 -17
  40. adam/commands/deploy/deploy_pg_agent.py +2 -5
  41. adam/commands/deploy/deploy_pod.py +64 -68
  42. adam/commands/deploy/undeploy.py +4 -27
  43. adam/commands/deploy/undeploy_frontend.py +4 -7
  44. adam/commands/deploy/undeploy_pg_agent.py +4 -7
  45. adam/commands/deploy/undeploy_pod.py +9 -12
  46. adam/commands/devices/device.py +93 -2
  47. adam/commands/devices/device_app.py +37 -10
  48. adam/commands/devices/device_auit_log.py +8 -2
  49. adam/commands/devices/device_cass.py +47 -7
  50. adam/commands/devices/device_export.py +9 -11
  51. adam/commands/devices/device_postgres.py +41 -6
  52. adam/commands/exit.py +1 -4
  53. adam/commands/export/clean_up_all_export_sessions.py +37 -0
  54. adam/commands/export/clean_up_export_sessions.py +12 -8
  55. adam/commands/export/drop_export_database.py +7 -26
  56. adam/commands/export/drop_export_databases.py +5 -14
  57. adam/commands/export/export.py +8 -38
  58. adam/commands/export/export_databases.py +86 -27
  59. adam/commands/export/export_select.py +25 -27
  60. adam/commands/export/export_select_x.py +3 -3
  61. adam/commands/export/export_sessions.py +124 -0
  62. adam/commands/export/export_use.py +8 -17
  63. adam/commands/export/exporter.py +88 -158
  64. adam/commands/export/import_session.py +7 -35
  65. adam/commands/export/importer.py +12 -5
  66. adam/commands/export/importer_athena.py +21 -20
  67. adam/commands/export/importer_sqlite.py +16 -21
  68. adam/commands/export/show_column_counts.py +7 -25
  69. adam/commands/export/show_export_databases.py +4 -6
  70. adam/commands/export/show_export_session.py +7 -18
  71. adam/commands/export/show_export_sessions.py +9 -12
  72. adam/commands/export/utils_export.py +26 -1
  73. adam/commands/intermediate_command.py +49 -0
  74. adam/commands/issues.py +11 -43
  75. adam/commands/kubectl.py +3 -6
  76. adam/commands/login.py +22 -24
  77. adam/commands/logs.py +3 -6
  78. adam/commands/ls.py +8 -9
  79. adam/commands/medusa/medusa.py +4 -22
  80. adam/commands/medusa/medusa_backup.py +20 -25
  81. adam/commands/medusa/medusa_restore.py +34 -36
  82. adam/commands/medusa/medusa_show_backupjobs.py +14 -18
  83. adam/commands/medusa/medusa_show_restorejobs.py +11 -18
  84. adam/commands/nodetool.py +6 -15
  85. adam/commands/param_get.py +11 -13
  86. adam/commands/param_set.py +8 -12
  87. adam/commands/postgres/postgres.py +22 -38
  88. adam/commands/postgres/postgres_context.py +47 -23
  89. adam/commands/postgres/postgres_ls.py +4 -8
  90. adam/commands/postgres/postgres_preview.py +5 -9
  91. adam/commands/postgres/psql_completions.py +1 -1
  92. adam/commands/postgres/utils_postgres.py +70 -0
  93. adam/commands/preview_table.py +6 -45
  94. adam/commands/pwd.py +13 -16
  95. adam/commands/reaper/reaper.py +4 -27
  96. adam/commands/reaper/reaper_forward.py +48 -55
  97. adam/commands/reaper/reaper_forward_session.py +6 -0
  98. adam/commands/reaper/reaper_forward_stop.py +10 -16
  99. adam/commands/reaper/reaper_restart.py +7 -14
  100. adam/commands/reaper/reaper_run_abort.py +8 -33
  101. adam/commands/reaper/reaper_runs.py +42 -57
  102. adam/commands/reaper/reaper_runs_abort.py +29 -49
  103. adam/commands/reaper/reaper_schedule_activate.py +9 -32
  104. adam/commands/reaper/reaper_schedule_start.py +9 -32
  105. adam/commands/reaper/reaper_schedule_stop.py +9 -32
  106. adam/commands/reaper/reaper_schedules.py +4 -14
  107. adam/commands/reaper/reaper_status.py +8 -16
  108. adam/commands/reaper/utils_reaper.py +196 -0
  109. adam/commands/repair/repair.py +4 -22
  110. adam/commands/repair/repair_log.py +5 -11
  111. adam/commands/repair/repair_run.py +27 -34
  112. adam/commands/repair/repair_scan.py +32 -38
  113. adam/commands/repair/repair_stop.py +5 -11
  114. adam/commands/report.py +27 -29
  115. adam/commands/restart.py +25 -26
  116. adam/commands/rollout.py +19 -24
  117. adam/commands/shell.py +10 -4
  118. adam/commands/show/show.py +10 -26
  119. adam/commands/show/show_cassandra_repairs.py +35 -0
  120. adam/commands/show/show_cassandra_status.py +32 -43
  121. adam/commands/show/show_cassandra_version.py +5 -18
  122. adam/commands/show/show_commands.py +19 -24
  123. adam/commands/show/show_host.py +1 -1
  124. adam/commands/show/show_login.py +20 -27
  125. adam/commands/show/show_processes.py +15 -19
  126. adam/commands/show/show_storage.py +10 -20
  127. adam/commands/watch.py +26 -29
  128. adam/config.py +4 -16
  129. adam/embedded_params.py +1 -1
  130. adam/log.py +4 -4
  131. adam/pod_exec_result.py +3 -3
  132. adam/repl.py +31 -32
  133. adam/repl_commands.py +11 -11
  134. adam/repl_state.py +52 -26
  135. adam/sql/sql_completer.py +4 -6
  136. adam/sql/sql_state_machine.py +21 -14
  137. adam/sso/authn_ad.py +6 -8
  138. adam/sso/authn_okta.py +4 -6
  139. adam/sso/cred_cache.py +3 -5
  140. adam/sso/idp.py +9 -12
  141. adam/utils.py +393 -33
  142. adam/utils_athena.py +14 -13
  143. adam/utils_audits.py +12 -12
  144. adam/utils_issues.py +32 -0
  145. adam/utils_k8s/app_clusters.py +13 -18
  146. adam/utils_k8s/app_pods.py +2 -0
  147. adam/utils_k8s/cassandra_clusters.py +21 -18
  148. adam/utils_k8s/custom_resources.py +16 -17
  149. adam/utils_k8s/ingresses.py +2 -2
  150. adam/utils_k8s/jobs.py +7 -11
  151. adam/utils_k8s/k8s.py +87 -0
  152. adam/utils_k8s/pods.py +14 -76
  153. adam/utils_k8s/secrets.py +4 -4
  154. adam/utils_k8s/service_accounts.py +5 -4
  155. adam/utils_k8s/services.py +2 -2
  156. adam/utils_k8s/statefulsets.py +1 -12
  157. adam/utils_repl/state_machine.py +3 -3
  158. adam/utils_sqlite.py +78 -42
  159. adam/version.py +1 -1
  160. {kaqing-2.0.145.dist-info → kaqing-2.0.174.dist-info}/METADATA +1 -1
  161. kaqing-2.0.174.dist-info/RECORD +230 -0
  162. adam/commands/app.py +0 -67
  163. adam/commands/app_ping.py +0 -44
  164. adam/commands/export/clean_up_export_session.py +0 -53
  165. adam/commands/postgres/postgres_utils.py +0 -31
  166. adam/commands/reaper/reaper_session.py +0 -159
  167. adam/commands/show/show_app_actions.py +0 -56
  168. adam/commands/show/show_app_id.py +0 -47
  169. adam/commands/show/show_app_queues.py +0 -45
  170. adam/commands/show/show_repairs.py +0 -47
  171. kaqing-2.0.145.dist-info/RECORD +0 -227
  172. {kaqing-2.0.145.dist-info → kaqing-2.0.174.dist-info}/WHEEL +0 -0
  173. {kaqing-2.0.145.dist-info → kaqing-2.0.174.dist-info}/entry_points.txt +0 -0
  174. {kaqing-2.0.145.dist-info → kaqing-2.0.174.dist-info}/top_level.txt +0 -0
@@ -1,10 +1,12 @@
1
- from adam.commands.command import Command
1
+ from datetime import datetime
2
+ from adam.commands import extract_trailing_options, validate_args
3
+ from adam.commands.command import Command, InvalidArgumentsException
2
4
  from adam.commands.export.export_databases import ExportDatabases
5
+ from adam.config import Config
3
6
  from adam.repl_state import ReplState, RequiredState
4
7
  from adam.sql.sql_completer import SqlCompleter, SqlVariant
5
8
  from adam.utils import log2
6
9
  from adam.utils_athena import Athena
7
- from adam.utils_sqlite import SQLite
8
10
 
9
11
  class ExportSelect(Command):
10
12
  COMMAND = '.select'
@@ -28,38 +30,34 @@ class ExportSelect(Command):
28
30
  if not(args := self.args(cmd)):
29
31
  return super().run(cmd, state)
30
32
 
31
- state, args = self.apply_state(args, state)
32
- if not self.validate_state(state):
33
- return state
33
+ with self.validate(args, state) as (args, state):
34
+ with extract_trailing_options(args, '&') as (args, backgrounded):
35
+ if not state.export_session:
36
+ if state.in_repl:
37
+ if state.device == ReplState.C:
38
+ log2("Select an export database first with 'use' command.")
39
+ else:
40
+ log2('cd to an export database first.')
41
+ else:
42
+ log2('* export database is missing.')
34
43
 
35
- if not state.export_session:
36
- if state.in_repl:
37
- if state.device == ReplState.C:
38
- log2("Select an export database first with 'use' command.")
39
- else:
40
- log2('cd to an export database first.')
41
- else:
42
- log2('* export database is missing.')
43
-
44
- Command.display_help()
45
-
46
- return 'command-missing'
44
+ Command.display_help()
47
45
 
48
- if not args:
49
- if state.in_repl:
50
- log2('Use a SQL statement.')
51
- else:
52
- log2('* SQL statement is missing.')
46
+ raise InvalidArgumentsException()
53
47
 
54
- Command.display_help()
48
+ with validate_args(args, state, name='SQL statement') as query:
49
+ def output(out: str):
50
+ log_prefix = Config().get('export.log-prefix', '/tmp/qing')
51
+ log_file = f'{log_prefix}-{datetime.now().strftime("%d%H%M%S")}-sqlite.log'
55
52
 
56
- return 'command-missing'
53
+ with open(log_file, 'w') as f:
54
+ f.write(out)
57
55
 
58
- query = ' '.join(args)
56
+ return log_file
59
57
 
60
- ExportDatabases.run_query(f'select {query}', database=state.export_session)
58
+ ExportDatabases.run_query(f'select {query}', database=state.export_session, output=output if backgrounded else None)
61
59
 
62
- return state
60
+ return state
63
61
 
64
62
  def completion(self, state: ReplState):
65
63
  if not state.export_session:
@@ -38,15 +38,15 @@ class ExportSelectX(Command):
38
38
  )}
39
39
 
40
40
  if state.export_session:
41
- completions |= {'select': SqlCompleter(
41
+ completions |= {dml: SqlCompleter(
42
42
  lambda: ExportDatabases.table_names(state.export_session),
43
- dml='select',
43
+ dml=dml,
44
44
  expandables={
45
45
  'export-dbs': lambda: ExportDatabases.database_names(),
46
46
  'columns':lambda _: Athena.column_names(database=state.export_session, function='export'),
47
47
  },
48
48
  variant=SqlVariant.ATHENA
49
- )}
49
+ ) for dml in ['select', 'preview']}
50
50
 
51
51
  return completions
52
52
 
@@ -0,0 +1,124 @@
1
+ import functools
2
+ import re
3
+
4
+ from adam.commands.export.importer import Importer
5
+ from adam.commands.export.utils_export import ExportTableStatus, csv_dir, find_files
6
+ from adam.config import Config
7
+ from adam.utils import lines_to_tabular, log, parallelize
8
+ from adam.utils_k8s.cassandra_nodes import CassandraNodes
9
+ from adam.utils_k8s.pods import log_prefix
10
+ from adam.utils_k8s.statefulsets import StatefulSets
11
+
12
+ class ExportSessions:
13
+ def clear_export_session_cache():
14
+ ExportSessions.find_export_sessions.cache_clear()
15
+ ExportSessions.export_session_names.cache_clear()
16
+
17
+ @functools.lru_cache()
18
+ def export_session_names(sts: str, pod: str, namespace: str, importer: str = None, export_state = None):
19
+ if not sts or not namespace:
20
+ return []
21
+
22
+ if not pod:
23
+ pod = StatefulSets.pod_names(sts, namespace)[0]
24
+
25
+ if not pod:
26
+ return []
27
+
28
+ return [session for session, state in ExportSessions.find_export_sessions(pod, namespace, importer).items() if not export_state or state == export_state]
29
+
30
+ @functools.lru_cache()
31
+ def find_export_sessions(pod: str, namespace: str, importer: str = None, limit = 100):
32
+ sessions: dict[str, str] = {}
33
+
34
+ prefix = Importer.prefix_from_importer(importer)
35
+
36
+ log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{prefix}*_*.log*')
37
+
38
+ if not log_files:
39
+ return {}
40
+
41
+ for log_file in log_files[:limit]:
42
+ m = re.match(f'{log_prefix()}-(.*?)_.*\.log?(.*)', log_file)
43
+ if m:
44
+ s = m.group(1)
45
+ state = m.group(2) # '', '.pending_import', '.done'
46
+ if state:
47
+ state = state.strip('.')
48
+ else:
49
+ state = 'in_export'
50
+
51
+ if s not in sessions:
52
+ sessions[s] = state
53
+ elif sessions[s] == 'done' and state != 'done':
54
+ sessions[s] = state
55
+
56
+ return sessions
57
+
58
+ def clean_up_all_sessions(sts: str, pod: str, namespace: str):
59
+ if not sts or not namespace:
60
+ return False
61
+
62
+ if not pod:
63
+ pod = StatefulSets.pod_names(sts, namespace)[0]
64
+
65
+ CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/*', show_out=Config().is_debug(), shell='bash')
66
+ CassandraNodes.exec(pod, namespace, f'rm -rf {log_prefix()}-*.log*', show_out=Config().is_debug(), shell='bash')
67
+
68
+ return True
69
+
70
+ def clean_up_sessions(sts: str, pod: str, namespace: str, sessions: list[str], max_workers = 0):
71
+ if not sessions:
72
+ return []
73
+
74
+ if not max_workers:
75
+ max_workers = Config().action_workers('export', 8)
76
+
77
+ with parallelize(sessions, max_workers, msg='Cleaning|Cleaned up {size} export sessions') as exec:
78
+ cnt_tuples = exec.map(lambda session: ExportSessions.clean_up_session(sts, pod, namespace, session, True))
79
+ csv_cnt = 0
80
+ log_cnt = 0
81
+ for (csv, log) in cnt_tuples:
82
+ csv_cnt += csv
83
+ log_cnt += log
84
+
85
+ return csv_cnt, log_cnt
86
+
87
+ def clean_up_session(sts: str, pod: str, namespace: str, session: str, multi_tables = True):
88
+ if not sts or not namespace:
89
+ return 0, 0
90
+
91
+ if not pod:
92
+ pod = StatefulSets.pod_names(sts, namespace)[0]
93
+
94
+ if not pod:
95
+ return 0, 0
96
+
97
+ csv_cnt = 0
98
+ log_cnt = 0
99
+
100
+ log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{session}_*.log*')
101
+
102
+ for log_file in log_files:
103
+ m = re.match(f'{log_prefix()}-{session}_(.*?)\.(.*?)\.log.*', log_file)
104
+ if m:
105
+ table = m.group(2)
106
+
107
+ CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/{session}_{table}', show_out=not multi_tables, shell='bash')
108
+ csv_cnt += 1
109
+
110
+ CassandraNodes.exec(pod, namespace, f'rm -rf {log_file}', show_out=not multi_tables, shell='bash')
111
+ log_cnt += 1
112
+
113
+ return csv_cnt, log_cnt
114
+
115
+ def disply_export_session(sts: str, pod: str, namespace: str, session: str):
116
+ if not pod:
117
+ pod = StatefulSets.pod_names(sts, namespace)[0]
118
+
119
+ if not pod:
120
+ return
121
+
122
+ tables, _ = ExportTableStatus.from_session(sts, pod, namespace, session)
123
+ log()
124
+ log(lines_to_tabular([f'{table.keyspace}\t{table.table}\t{table.target_table}\t{"export_completed_pending_import" if table.status == "pending_import" else table.status}' for table in tables], header='KEYSPACE\tTABLE\tTARGET_TABLE\tSTATUS', separator='\t'))
@@ -2,8 +2,6 @@ from adam.commands.command import Command
2
2
  from adam.commands.export.export_databases import ExportDatabases
3
3
  from adam.repl_state import ReplState
4
4
  from adam.utils import log2
5
- from adam.utils_athena import Athena
6
- from adam.utils_sqlite import SQLite
7
5
 
8
6
  class ExportUse(Command):
9
7
  COMMAND = 'use'
@@ -27,26 +25,19 @@ class ExportUse(Command):
27
25
  if not(args := self.args(cmd)):
28
26
  return super().run(cmd, state)
29
27
 
30
- state, args = self.apply_state(args, state)
31
- if not self.validate_state(state):
32
- return state
33
-
34
- if not args:
35
- state.export_session = None
28
+ with self.validate(args, state) as (args, state):
29
+ if not args:
30
+ state.export_session = None
36
31
 
37
- log2('Export database is unset.')
32
+ log2('Export database is unset.')
38
33
 
39
- return state
34
+ return state
40
35
 
41
- state.export_session = args[0]
42
- if state.export_session.startswith('e'):
43
- Athena.clear_cache()
44
- else:
45
- SQLite.clear_cache()
36
+ state.export_session = args[0]
46
37
 
47
- ExportDatabases.display_export_db(state.export_session)
38
+ ExportDatabases.display_export_db(state.export_session)
48
39
 
49
- return state
40
+ return state
50
41
 
51
42
  def completion(self, state: ReplState):
52
43
  return super().completion(state, {n: None for n in ExportDatabases.database_names()})
@@ -1,23 +1,21 @@
1
- from concurrent.futures import ThreadPoolExecutor, as_completed
2
1
  from datetime import datetime
3
- import functools
4
- import re
5
2
  import time
6
- import traceback
7
3
 
8
- from adam.commands.cql.cql_utils import cassandra_table_names, run_cql, table_spec
4
+ from adam.commands import validate_args
5
+ from adam.commands.command import Command
6
+ from adam.commands.cql.utils_cql import cassandra_table_names, run_cql, table_spec
9
7
  from adam.commands.export.export_databases import ExportDatabases
8
+ from adam.commands.export.export_sessions import ExportSessions
10
9
  from adam.commands.export.importer import Importer
11
10
  from adam.commands.export.importer_athena import AthenaImporter
12
11
  from adam.commands.export.importer_sqlite import SqliteImporter
13
- from adam.commands.export.utils_export import ExportSpec, ExportTableStatus, ExportTableSpec, ImportSpec, csv_dir, find_files
12
+ from adam.commands.export.utils_export import ExportSpec, ExportTableStatus, ExportTableSpec, ImportSpec, csv_dir, find_files, state_with_pod
14
13
  from adam.config import Config
15
14
  from adam.pod_exec_result import PodExecResult
16
15
  from adam.repl_state import ReplState
17
- from adam.utils import elapsed_time, log2, ing
16
+ from adam.utils import debug, log, parallelize, log2, ing, log_exc
18
17
  from adam.utils_k8s.cassandra_nodes import CassandraNodes
19
18
  from adam.utils_k8s.pods import log_prefix
20
- from adam.utils_k8s.statefulsets import StatefulSets
21
19
 
22
20
  class Exporter:
23
21
  def export_tables(args: list[str], state: ReplState, export_only: bool = False, max_workers = 0) -> tuple[list[str], ExportSpec]:
@@ -25,7 +23,7 @@ class Exporter:
25
23
  log2('export-only for testing')
26
24
 
27
25
  spec: ExportSpec = None
28
- try:
26
+ with log_exc(True):
29
27
  spec = Exporter.export_spec(' '.join(args), state)
30
28
 
31
29
  statuses, spec = Exporter._export_tables(spec, state, max_workers=max_workers, export_state='init')
@@ -33,8 +31,6 @@ class Exporter:
33
31
  return statuses, spec
34
32
 
35
33
  return Exporter._export_tables(spec, state, export_only, max_workers, 'pending_export')
36
- except Exception as e:
37
- log2(e)
38
34
 
39
35
  return [], None
40
36
 
@@ -53,12 +49,19 @@ class Exporter:
53
49
  raise Exception(f"You're currently using {importer_from_session} export database. You cannot export tables with {spec.importer} type database.")
54
50
  else:
55
51
  spec.importer = Importer.importer_from_session(session)
52
+
53
+ if spec.importer == 'athena' and not AthenaImporter.ping():
54
+ raise Exception('Credentials for Athena is not present.')
56
55
  else:
57
56
  if not spec.importer:
58
57
  spec.importer = Config().get('export.default-importer', 'sqlite')
59
58
 
60
59
  prefix = Importer.prefix_from_importer(spec.importer)
61
60
  session = f'{prefix}{datetime.now().strftime("%Y%m%d%H%M%S")[3:]}'
61
+
62
+ if spec.importer == 'athena' and not AthenaImporter.ping():
63
+ raise Exception('Credentials for Athena is not present.')
64
+
62
65
  if spec.importer != 'csv':
63
66
  state.export_session = session
64
67
 
@@ -68,7 +71,7 @@ class Exporter:
68
71
 
69
72
  def import_session(args: list[str], state: ReplState, max_workers = 0) -> tuple[list[str], ExportSpec]:
70
73
  import_spec: ImportSpec = None
71
- try:
74
+ with log_exc(True):
72
75
  import_spec = Exporter.import_spec(' '.join(args), state)
73
76
  tables, status_in_whole = ExportTableStatus.from_session(state.sts, state.pod, state.namespace, import_spec.session)
74
77
  if status_in_whole == 'done':
@@ -77,12 +80,7 @@ class Exporter:
77
80
 
78
81
  spec = ExportSpec(None, None, importer=import_spec.importer, tables=[ExportTableSpec.from_status(table) for table in tables], session=import_spec.session)
79
82
 
80
- return Exporter._export_tables(spec, state, max_workers=max_workers)
81
- except Exception as e:
82
- if Config().is_debug():
83
- traceback.print_exception(e)
84
- else:
85
- log2(e)
83
+ return Exporter._export_tables(spec, state, max_workers=max_workers, export_state = 'import')
86
84
 
87
85
  return [], None
88
86
 
@@ -99,16 +97,19 @@ class Exporter:
99
97
  spec.importer = Importer.importer_from_session(state.export_session)
100
98
  if not spec.importer:
101
99
  spec.importer = Config().get('export.default-importer', 'sqlite')
100
+
101
+ if spec.importer == 'athena' and not AthenaImporter.ping():
102
+ raise Exception('Credentials for Athena is not present.')
102
103
  else:
103
- if spec.importer:
104
- if not AthenaImporter.ping():
105
- raise Exception('Credentials for Athena are not present.')
106
- else:
104
+ if not spec.importer:
107
105
  spec.importer = Importer.importer_from_session(spec.session)
108
106
 
109
107
  if spec.importer == 'csv':
110
108
  spec.importer = Config().get('export.default-importer', 'sqlite')
111
109
 
110
+ if spec.importer == 'athena' and not AthenaImporter.ping():
111
+ raise Exception('Credentials for Athena is not present.')
112
+
112
113
  prefix = Importer.prefix_from_importer(spec.importer)
113
114
  session = f'{prefix}{spec.session[1:]}'
114
115
  state.export_session = session
@@ -128,20 +129,14 @@ class Exporter:
128
129
  if export_state == 'init':
129
130
  CassandraNodes.exec(state.pod, state.namespace, f'rm -rf {csv_dir()}/{spec.session}_*', show_out=Config().is_debug(), shell='bash')
130
131
 
131
- if max_workers > 1 and len(spec.tables) > 1:
132
- log2(f'Executing on {len(spec.tables)} Cassandra tables in parallel...')
133
- start_time = time.time()
134
- try:
135
- with ThreadPoolExecutor(max_workers=max_workers) as executor:
136
- futures = [executor.submit(Exporter.export_table, table, state, spec.session, spec.importer, export_only, True, consistency=spec.consistency, export_state=export_state) for table in spec.tables]
137
- if len(futures) == 0:
138
- return [], spec
139
-
140
- return [future.result() for future in as_completed(futures)], spec
141
- finally:
142
- log2(f"{len(spec.tables)} parallel table export elapsed time: {elapsed_time(start_time)} with {max_workers} workers")
143
- else:
144
- return [Exporter.export_table(table, state, spec.session, spec.importer, export_only, multi_tables=len(spec.tables) > 1, consistency=spec.consistency, export_state=export_state) for table in spec.tables], spec
132
+ action = f'[{spec.session}] Exporting|Exported'
133
+ if export_state == 'init':
134
+ action = f'[{spec.session}] Preparing|Prepared'
135
+ elif export_state == 'import':
136
+ action = f'[{spec.session}] Importing|Imported'
137
+
138
+ with parallelize(spec.tables, max_workers, msg=action + ' {size} Cassandra tables') as exec:
139
+ return exec.map(lambda table: Exporter.export_table(table, state, spec.session, spec.importer, export_only, len(spec.tables) > 1, consistency=spec.consistency, export_state=export_state)), spec
145
140
 
146
141
  def export_table(spec: ExportTableSpec, state: ReplState, session: str, importer: str, export_only = False, multi_tables = True, consistency: str = None, export_state=None):
147
142
  s: str = None
@@ -167,8 +162,7 @@ class Exporter:
167
162
  status: ExportTableStatus = ExportTableStatus.from_log_file(state.pod, state.namespace, session, log_file)
168
163
  while status.status != 'done':
169
164
  if status.status == 'export_in_pregress':
170
- if Config().is_debug():
171
- log2('Exporting to CSV is still in progess, sleeping for 1 sec...')
165
+ debug('Exporting to CSV is still in progess, sleeping for 1 sec...')
172
166
  time.sleep(1)
173
167
  elif status.status == 'exported':
174
168
  log_file = Exporter.rename_to_pending_import(spec, state, session, target_table)
@@ -217,114 +211,7 @@ class Exporter:
217
211
 
218
212
  def import_from_csv(spec: ExportTableSpec, state: ReplState, session: str, importer: str, table: str, target_table: str, columns: str, multi_tables = True, create_db = False):
219
213
  im = AthenaImporter() if importer == 'athena' else SqliteImporter()
220
- return im.import_from_csv(state.pod, state.namespace, state.export_session, session if session else state.export_session, spec.keyspace, table, target_table, columns, multi_tables, create_db)
221
-
222
- def clear_export_session_cache():
223
- Exporter.find_export_sessions.cache_clear()
224
- Exporter.export_session_names.cache_clear()
225
-
226
- @functools.lru_cache()
227
- def export_session_names(sts: str, pod: str, namespace: str, importer: str = None, export_state = None):
228
- if not sts or not namespace:
229
- return []
230
-
231
- if not pod:
232
- pod = StatefulSets.pod_names(sts, namespace)[0]
233
-
234
- if not pod:
235
- return []
236
-
237
- return [session for session, state in Exporter.find_export_sessions(pod, namespace, importer).items() if not export_state or state == export_state]
238
-
239
- @functools.lru_cache()
240
- def find_export_sessions(pod: str, namespace: str, importer: str = None, limit = 100):
241
- sessions: dict[str, str] = {}
242
-
243
- prefix = Importer.prefix_from_importer(importer)
244
-
245
- log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{prefix}*_*.log*')
246
-
247
- if not log_files:
248
- return {}
249
-
250
- for log_file in log_files[:limit]:
251
- m = re.match(f'{log_prefix()}-(.*?)_.*\.log?(.*)', log_file)
252
- if m:
253
- s = m.group(1)
254
- state = m.group(2) # '', '.pending_import', '.done'
255
- if state:
256
- state = state.strip('.')
257
- else:
258
- state = 'in_export'
259
-
260
- if s not in sessions:
261
- sessions[s] = state
262
- elif sessions[s] == 'done' and state != 'done':
263
- sessions[s] = state
264
-
265
- return sessions
266
-
267
- def clean_up_all_sessions(sts: str, pod: str, namespace: str):
268
- if not sts or not namespace:
269
- return False
270
-
271
- if not pod:
272
- pod = StatefulSets.pod_names(sts, namespace)[0]
273
-
274
- CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/*', show_out=Config().is_debug(), shell='bash')
275
- CassandraNodes.exec(pod, namespace, f'rm -rf {log_prefix()}-*.log*', show_out=Config().is_debug(), shell='bash')
276
-
277
- return True
278
-
279
- def clean_up_sessions(sts: str, pod: str, namespace: str, sessions: list[str], max_workers = 0):
280
- if not sessions:
281
- return []
282
-
283
- if not max_workers:
284
- max_workers = Config().action_workers('export', 8)
285
-
286
- if max_workers > 1 and len(sessions) > 1:
287
- log2(f'Executing on {len(sessions)} export session clean ups in parallel...')
288
- start_time = time.time()
289
- try:
290
- with ThreadPoolExecutor(max_workers=max_workers) as executor:
291
- futures = [executor.submit(Exporter.clean_up_session, sts, pod, namespace, session, True) for session in sessions]
292
- if len(futures) == 0:
293
- return []
294
-
295
- return [future.result() for future in as_completed(futures)]
296
- finally:
297
- log2(f"{len(sessions)} parallel session clean ups elapsed time: {elapsed_time(start_time)} with {max_workers} workers")
298
- else:
299
- return [Exporter.clean_up_session(sts, pod, namespace, session) for session in sessions]
300
-
301
- def clean_up_session(sts: str, pod: str, namespace: str, session: str, multi_tables = True):
302
- if not sts or not namespace:
303
- return 0, 0
304
-
305
- if not pod:
306
- pod = StatefulSets.pod_names(sts, namespace)[0]
307
-
308
- if not pod:
309
- return 0, 0
310
-
311
- csv_cnt = 0
312
- log_cnt = 0
313
-
314
- log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{session}_*.log*')
315
-
316
- for log_file in log_files:
317
- m = re.match(f'{log_prefix()}-{session}_(.*?)\.(.*?)\.log.*', log_file)
318
- if m:
319
- table = m.group(2)
320
-
321
- CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/{session}_{table}', show_out=not multi_tables, shell='bash')
322
- csv_cnt += 1
323
-
324
- CassandraNodes.exec(pod, namespace, f'rm -rf {log_file}', show_out=not multi_tables, shell='bash')
325
- log_cnt += 1
326
-
327
- return csv_cnt, log_cnt
214
+ return im.import_from_csv(state, session if session else state.export_session, spec.keyspace, table, target_table, columns, multi_tables, create_db)
328
215
 
329
216
  def resove_table_n_columns(spec: ExportTableSpec, state: ReplState, include_ks_in_target = False, importer = 'sqlite'):
330
217
  table = spec.table
@@ -350,15 +237,58 @@ class Exporter:
350
237
 
351
238
  return table, target_table, columns
352
239
 
353
- def drop_databases(sts: str, pod: str, namespace: str, db: str = None):
354
- importer = None
355
- if db:
356
- importer = Importer.importer_from_session(db)
357
-
358
- sessions_done = Exporter.export_session_names(sts, pod, namespace, importer=importer, export_state='done')
359
- sessions = ExportDatabases.sessions_from_dbs(ExportDatabases.drop_export_dbs(db))
360
- if sessions_done and sessions:
361
- intersects = list(set(sessions_done) & set(sessions))
362
- with ing(f'Cleaning up {len(intersects)} completed sessions'):
363
- Exporter.clean_up_sessions(sts, pod, namespace, list(intersects))
364
- Exporter.clear_export_session_cache()
240
+ class ExportService:
241
+ def __init__(self, handler: 'ExporterHandler'):
242
+ self.handler = handler
243
+
244
+ def export(self, args: list[str], export_only=False):
245
+ state = self.handler.state
246
+ export_session = state.export_session
247
+ spec: ExportSpec = None
248
+ try:
249
+ with state_with_pod(state) as state:
250
+ # --export-only for testing only
251
+ statuses, spec = Exporter.export_tables(args, state, export_only=export_only)
252
+ if not statuses:
253
+ return state
254
+
255
+ ExportSessions.clear_export_session_cache()
256
+
257
+ if spec.importer == 'csv' or export_only:
258
+ ExportSessions.disply_export_session(state.sts, state.pod, state.namespace, spec.session)
259
+ else:
260
+ log()
261
+ ExportDatabases.display_export_db(state.export_session)
262
+ finally:
263
+ # if exporting to csv, do not bind the new session id to repl state
264
+ if spec and spec.importer == 'csv':
265
+ state.export_session = export_session
266
+
267
+ return state
268
+
269
+ def import_sesion(self, args: list[str]):
270
+ state = self.handler.state
271
+
272
+ with validate_args(args, state, name='export session') as args_str:
273
+ with state_with_pod(state) as state:
274
+ tables, _ = Exporter.import_session(args_str, state)
275
+ if tables:
276
+ ExportSessions.clear_export_session_cache()
277
+
278
+ log()
279
+ ExportDatabases.display_export_db(state.export_session)
280
+
281
+ return state
282
+
283
+ class ExporterHandler:
284
+ def __init__(self, state: ReplState):
285
+ self.state = state
286
+
287
+ def __enter__(self):
288
+ return ExportService(self)
289
+
290
+ def __exit__(self, exc_type, exc_val, exc_tb):
291
+ return False
292
+
293
+ def export(state: ReplState):
294
+ return ExporterHandler(state)
@@ -1,9 +1,7 @@
1
1
  from adam.commands.command import Command
2
- from adam.commands.export.export_databases import ExportDatabases
3
- from adam.commands.export.exporter import Exporter
2
+ from adam.commands.export.export_sessions import ExportSessions
3
+ from adam.commands.export.exporter import export
4
4
  from adam.repl_state import ReplState, RequiredState
5
- from adam.utils import log, log2
6
- from adam.utils_k8s.statefulsets import StatefulSets
7
5
 
8
6
  class ImportSession(Command):
9
7
  COMMAND = 'import session'
@@ -27,40 +25,14 @@ class ImportSession(Command):
27
25
  if not(args := self.args(cmd)):
28
26
  return super().run(cmd, state)
29
27
 
30
- state, args = self.apply_state(args, state)
31
- if not self.validate_state(state):
32
- return state
33
-
34
- if not args:
35
- if state.in_repl:
36
- log2('Specify export session name.')
37
- else:
38
- log2('* Export session name is missing.')
39
-
40
- Command.display_help()
41
-
42
- return 'command-missing'
43
-
44
- if not state.pod:
45
- state.push()
46
- state.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
47
-
48
- try:
49
- tables, _ = Exporter.import_session(args, state)
50
- if tables:
51
- Exporter.clear_export_session_cache()
52
-
53
- log()
54
- ExportDatabases.display_export_db(state.export_session)
55
- finally:
56
- state.pop()
57
-
58
- return state
28
+ with self.validate(args, state) as (args, state):
29
+ with export(state) as exporter:
30
+ return exporter.import_sesion(args)
59
31
 
60
32
  def completion(self, state: ReplState):
61
33
  # warm up cache
62
- Exporter.export_session_names(state.sts, state.pod, state.namespace)
63
- Exporter.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')
34
+ ExportSessions.export_session_names(state.sts, state.pod, state.namespace)
35
+ ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')
64
36
 
65
37
  return {}
66
38