kaqing 2.0.171__py3-none-any.whl → 2.0.203__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.
Files changed (194) hide show
  1. adam/app_session.py +5 -10
  2. adam/apps.py +18 -4
  3. adam/batch.py +7 -7
  4. adam/checks/check_utils.py +3 -1
  5. adam/checks/disk.py +2 -3
  6. adam/columns/memory.py +3 -4
  7. adam/commands/__init__.py +15 -6
  8. adam/commands/alter_tables.py +26 -41
  9. adam/commands/app/__init__.py +0 -0
  10. adam/commands/{app_cmd.py → app/app.py} +2 -2
  11. adam/commands/{show → app}/show_app_actions.py +7 -15
  12. adam/commands/{show → app}/show_app_queues.py +1 -4
  13. adam/{utils_app.py → commands/app/utils_app.py} +9 -1
  14. adam/commands/audit/audit.py +9 -26
  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 +15 -3
  22. adam/commands/bash/bash.py +1 -1
  23. adam/commands/bash/utils_bash.py +1 -1
  24. adam/commands/cassandra/__init__.py +0 -0
  25. adam/commands/cassandra/download_cassandra_log.py +45 -0
  26. adam/commands/cassandra/nodetool.py +64 -0
  27. adam/commands/cassandra/nodetool_commands.py +120 -0
  28. adam/commands/cassandra/restart_cluster.py +47 -0
  29. adam/commands/cassandra/restart_node.py +51 -0
  30. adam/commands/cassandra/restart_nodes.py +47 -0
  31. adam/commands/cassandra/rollout.py +88 -0
  32. adam/commands/cat.py +5 -19
  33. adam/commands/cd.py +7 -9
  34. adam/commands/check.py +10 -18
  35. adam/commands/cli_commands.py +6 -1
  36. adam/commands/{cp.py → clipboard_copy.py} +34 -36
  37. adam/commands/code.py +2 -2
  38. adam/commands/command.py +139 -22
  39. adam/commands/commands_utils.py +14 -12
  40. adam/commands/cql/alter_tables.py +66 -0
  41. adam/commands/cql/completions_c.py +29 -0
  42. adam/commands/cql/cqlsh.py +3 -7
  43. adam/commands/cql/utils_cql.py +23 -61
  44. adam/commands/debug/__init__.py +0 -0
  45. adam/commands/debug/debug.py +22 -0
  46. adam/commands/debug/debug_completes.py +35 -0
  47. adam/commands/debug/debug_timings.py +35 -0
  48. adam/commands/deploy/deploy_pg_agent.py +2 -2
  49. adam/commands/deploy/deploy_pod.py +2 -4
  50. adam/commands/deploy/undeploy_pg_agent.py +2 -2
  51. adam/commands/devices/device.py +40 -9
  52. adam/commands/devices/device_app.py +19 -29
  53. adam/commands/devices/device_auit_log.py +3 -3
  54. adam/commands/devices/device_cass.py +17 -23
  55. adam/commands/devices/device_export.py +12 -11
  56. adam/commands/devices/device_postgres.py +79 -63
  57. adam/commands/devices/devices.py +1 -1
  58. adam/commands/download_cassandra_log.py +45 -0
  59. adam/commands/download_file.py +47 -0
  60. adam/commands/export/clean_up_all_export_sessions.py +3 -3
  61. adam/commands/export/clean_up_export_sessions.py +7 -19
  62. adam/commands/export/completions_x.py +11 -0
  63. adam/commands/export/download_export_session.py +40 -0
  64. adam/commands/export/drop_export_database.py +6 -22
  65. adam/commands/export/drop_export_databases.py +3 -9
  66. adam/commands/export/export.py +1 -17
  67. adam/commands/export/export_databases.py +109 -32
  68. adam/commands/export/export_select.py +8 -55
  69. adam/commands/export/export_sessions.py +211 -0
  70. adam/commands/export/export_use.py +13 -16
  71. adam/commands/export/export_x_select.py +48 -0
  72. adam/commands/export/exporter.py +176 -167
  73. adam/commands/export/import_files.py +44 -0
  74. adam/commands/export/import_session.py +10 -6
  75. adam/commands/export/importer.py +24 -9
  76. adam/commands/export/importer_athena.py +114 -44
  77. adam/commands/export/importer_sqlite.py +45 -23
  78. adam/commands/export/show_column_counts.py +11 -20
  79. adam/commands/export/show_export_databases.py +5 -2
  80. adam/commands/export/show_export_session.py +6 -15
  81. adam/commands/export/show_export_sessions.py +4 -11
  82. adam/commands/export/utils_export.py +79 -27
  83. adam/commands/find_files.py +51 -0
  84. adam/commands/find_processes.py +76 -0
  85. adam/commands/generate_report.py +52 -0
  86. adam/commands/head.py +36 -0
  87. adam/commands/help.py +2 -2
  88. adam/commands/intermediate_command.py +6 -3
  89. adam/commands/login.py +3 -6
  90. adam/commands/ls.py +2 -2
  91. adam/commands/medusa/medusa_backup.py +13 -16
  92. adam/commands/medusa/medusa_restore.py +26 -37
  93. adam/commands/medusa/medusa_show_backupjobs.py +7 -7
  94. adam/commands/medusa/medusa_show_restorejobs.py +6 -6
  95. adam/commands/medusa/utils_medusa.py +15 -0
  96. adam/commands/nodetool.py +3 -8
  97. adam/commands/os/__init__.py +0 -0
  98. adam/commands/os/cat.py +36 -0
  99. adam/commands/os/download_file.py +47 -0
  100. adam/commands/os/find_files.py +51 -0
  101. adam/commands/os/find_processes.py +76 -0
  102. adam/commands/os/head.py +36 -0
  103. adam/commands/os/shell.py +41 -0
  104. adam/commands/param_get.py +10 -12
  105. adam/commands/param_set.py +7 -10
  106. adam/commands/postgres/completions_p.py +22 -0
  107. adam/commands/postgres/postgres.py +25 -40
  108. adam/commands/postgres/postgres_databases.py +269 -0
  109. adam/commands/postgres/utils_postgres.py +33 -20
  110. adam/commands/preview_table.py +4 -2
  111. adam/commands/pwd.py +4 -6
  112. adam/commands/reaper/reaper_forward.py +2 -2
  113. adam/commands/reaper/reaper_run_abort.py +4 -10
  114. adam/commands/reaper/reaper_runs.py +3 -3
  115. adam/commands/reaper/reaper_schedule_activate.py +12 -12
  116. adam/commands/reaper/reaper_schedule_start.py +7 -12
  117. adam/commands/reaper/reaper_schedule_stop.py +7 -12
  118. adam/commands/reaper/utils_reaper.py +13 -6
  119. adam/commands/repair/repair_log.py +1 -4
  120. adam/commands/repair/repair_run.py +3 -8
  121. adam/commands/repair/repair_scan.py +1 -6
  122. adam/commands/repair/repair_stop.py +1 -5
  123. adam/commands/restart_cluster.py +47 -0
  124. adam/commands/restart_node.py +51 -0
  125. adam/commands/restart_nodes.py +47 -0
  126. adam/commands/shell.py +9 -2
  127. adam/commands/show/show.py +4 -4
  128. adam/commands/show/show_adam.py +3 -3
  129. adam/commands/show/show_cassandra_repairs.py +5 -6
  130. adam/commands/show/show_cassandra_status.py +29 -29
  131. adam/commands/show/show_cassandra_version.py +1 -4
  132. adam/commands/show/{show_commands.py → show_cli_commands.py} +3 -6
  133. adam/commands/show/show_login.py +3 -9
  134. adam/commands/show/show_params.py +2 -5
  135. adam/commands/show/show_processes.py +15 -16
  136. adam/commands/show/show_storage.py +9 -8
  137. adam/config.py +4 -5
  138. adam/embedded_params.py +1 -1
  139. adam/log.py +4 -4
  140. adam/repl.py +26 -18
  141. adam/repl_commands.py +32 -20
  142. adam/repl_session.py +9 -1
  143. adam/repl_state.py +39 -10
  144. adam/sql/async_executor.py +44 -0
  145. adam/sql/lark_completer.py +286 -0
  146. adam/sql/lark_parser.py +604 -0
  147. adam/sql/qingl.lark +1076 -0
  148. adam/sql/sql_completer.py +4 -6
  149. adam/sql/sql_state_machine.py +25 -13
  150. adam/sso/authn_ad.py +2 -5
  151. adam/sso/authn_okta.py +2 -4
  152. adam/sso/cred_cache.py +2 -5
  153. adam/sso/idp.py +8 -11
  154. adam/utils.py +299 -105
  155. adam/utils_athena.py +18 -18
  156. adam/utils_audits.py +3 -7
  157. adam/utils_issues.py +2 -2
  158. adam/utils_k8s/app_clusters.py +4 -4
  159. adam/utils_k8s/app_pods.py +8 -6
  160. adam/utils_k8s/cassandra_clusters.py +16 -5
  161. adam/utils_k8s/cassandra_nodes.py +9 -6
  162. adam/utils_k8s/custom_resources.py +11 -17
  163. adam/utils_k8s/jobs.py +7 -11
  164. adam/utils_k8s/k8s.py +14 -5
  165. adam/utils_k8s/kube_context.py +3 -6
  166. adam/{pod_exec_result.py → utils_k8s/pod_exec_result.py} +4 -4
  167. adam/utils_k8s/pods.py +85 -23
  168. adam/utils_k8s/statefulsets.py +5 -2
  169. adam/utils_local.py +42 -0
  170. adam/utils_repl/appendable_completer.py +6 -0
  171. adam/utils_repl/repl_completer.py +45 -2
  172. adam/utils_repl/state_machine.py +3 -3
  173. adam/utils_sqlite.py +58 -30
  174. adam/version.py +1 -1
  175. {kaqing-2.0.171.dist-info → kaqing-2.0.203.dist-info}/METADATA +1 -1
  176. kaqing-2.0.203.dist-info/RECORD +277 -0
  177. kaqing-2.0.203.dist-info/top_level.txt +2 -0
  178. teddy/__init__.py +0 -0
  179. teddy/lark_parser.py +436 -0
  180. teddy/lark_parser2.py +618 -0
  181. adam/commands/cql/cql_completions.py +0 -33
  182. adam/commands/export/export_handlers.py +0 -71
  183. adam/commands/export/export_select_x.py +0 -54
  184. adam/commands/logs.py +0 -37
  185. adam/commands/postgres/postgres_context.py +0 -274
  186. adam/commands/postgres/psql_completions.py +0 -10
  187. adam/commands/report.py +0 -61
  188. adam/commands/restart.py +0 -60
  189. kaqing-2.0.171.dist-info/RECORD +0 -236
  190. kaqing-2.0.171.dist-info/top_level.txt +0 -1
  191. /adam/commands/{app_ping.py → app/app_ping.py} +0 -0
  192. /adam/commands/{show → app}/show_app_id.py +0 -0
  193. {kaqing-2.0.171.dist-info → kaqing-2.0.203.dist-info}/WHEEL +0 -0
  194. {kaqing-2.0.171.dist-info → kaqing-2.0.203.dist-info}/entry_points.txt +0 -0
adam/repl_commands.py CHANGED
@@ -1,9 +1,16 @@
1
1
  from adam.commands.alter_tables import AlterTables
2
- from adam.commands.app_cmd import App
3
- from adam.commands.app_ping import AppPing
2
+ from adam.commands.app.app import App
3
+ from adam.commands.app.app_ping import AppPing
4
+ from adam.commands.app.show_app_actions import ShowAppActions
5
+ from adam.commands.app.show_app_id import ShowAppId
6
+ from adam.commands.app.show_app_queues import ShowAppQueues
4
7
  from adam.commands.audit.audit import Audit
5
8
  from adam.commands.cat import Cat
6
9
  from adam.commands.code import Code
10
+ from adam.commands.debug.debug import Debug
11
+ from adam.commands.debug.debug_timings import DebugTimings
12
+ from adam.commands.download_cassandra_log import DownloadCassandraLog
13
+ from adam.commands.download_file import DownloadFile
7
14
  from adam.commands.deploy.code_start import CodeStart
8
15
  from adam.commands.deploy.code_stop import CodeStop
9
16
  from adam.commands.deploy.deploy import Deploy
@@ -19,23 +26,29 @@ from adam.commands.devices.device_auit_log import DeviceAuditLog
19
26
  from adam.commands.devices.device_cass import DeviceCass
20
27
  from adam.commands.devices.device_export import DeviceExport
21
28
  from adam.commands.devices.device_postgres import DevicePostgres
29
+ from adam.commands.export.download_export_session import DownloadExportSession
22
30
  from adam.commands.export.drop_export_database import DropExportDatabase
23
31
  from adam.commands.export.export import ExportTables
32
+ from adam.commands.export.import_files import ImportCSVFiles
24
33
  from adam.commands.export.import_session import ImportSession
25
34
  from adam.commands.export.clean_up_export_sessions import CleanUpExportSessions
26
35
  from adam.commands.export.clean_up_all_export_sessions import CleanUpAllExportSessions
27
36
  from adam.commands.export.drop_export_databases import DropExportDatabases
28
- from adam.commands.export.export_select import ExportSelect
37
+ from adam.commands.export.export_x_select import ExportXSelect
29
38
  from adam.commands.export.export_use import ExportUse
30
- from adam.commands.export.export_select_x import ExportSelectX
39
+ from adam.commands.export.export_select import ExportSelect
31
40
  from adam.commands.export.show_column_counts import ShowColumnCounts
32
41
  from adam.commands.export.show_export_databases import ShowExportDatabases
33
42
  from adam.commands.export.show_export_session import ShowExportSession
34
43
  from adam.commands.export.show_export_sessions import ShowExportSessions
44
+ from adam.commands.find_files import FindLocalFiles
45
+ from adam.commands.find_processes import FindProcesses
46
+ from adam.commands.head import Head
35
47
  from adam.commands.kubectl import Kubectl
48
+ from adam.commands.restart_cluster import RestartCluster
49
+ from adam.commands.restart_node import RestartNode
36
50
  from adam.commands.shell import Shell
37
- from adam.commands.show.show_app_queues import ShowAppQueues
38
- from adam.commands.cp import ClipboardCopy
51
+ from adam.commands.clipboard_copy import ClipboardCopy
39
52
  from adam.commands.bash.bash import Bash
40
53
  from adam.commands.cd import Cd
41
54
  from adam.commands.check import Check
@@ -52,16 +65,14 @@ from adam.commands.preview_table import PreviewTable
52
65
  from adam.commands.pwd import Pwd
53
66
  from adam.commands.reaper.reaper import Reaper
54
67
  from adam.commands.repair.repair import Repair
55
- from adam.commands.report import Report
56
- from adam.commands.restart import Restart
68
+ from adam.commands.generate_report import GenerateReport
69
+ from adam.commands.restart_nodes import RestartNodes
57
70
  from adam.commands.rollout import RollOut
58
71
  from adam.commands.param_set import SetParam
59
72
  from adam.commands.show.show import Show
60
- from adam.commands.show.show_app_actions import ShowAppActions
61
- from adam.commands.show.show_app_id import ShowAppId
62
73
  from adam.commands.show.show_cassandra_status import ShowCassandraStatus
63
74
  from adam.commands.show.show_cassandra_version import ShowCassandraVersion
64
- from adam.commands.show.show_commands import ShowKubectlCommands
75
+ from adam.commands.show.show_cli_commands import ShowKubectlCommands
65
76
  from adam.commands.show.show_host import ShowHost
66
77
  from adam.commands.show.show_login import ShowLogin
67
78
  from adam.commands.show.show_params import ShowParams
@@ -76,7 +87,7 @@ class ReplCommands:
76
87
  cmds: list[Command] = ReplCommands.navigation() + ReplCommands.cassandra_ops() + ReplCommands.postgres_ops() + \
77
88
  ReplCommands.app_ops() + ReplCommands.audit_ops() + ReplCommands.export_ops() + ReplCommands.tools() + ReplCommands.exit()
78
89
 
79
- intermediate_cmds: list[Command] = [App(), Audit(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
90
+ intermediate_cmds: list[Command] = [App(), Audit(), Reaper(), Repair(), Debug(), Deploy(), Show(), Undeploy()]
80
91
  ic = [c.command() for c in intermediate_cmds]
81
92
  # 1. dedup commands
82
93
  deduped = []
@@ -93,17 +104,18 @@ class ReplCommands:
93
104
  return deduped
94
105
 
95
106
  def navigation() -> list[Command]:
96
- return [Ls(), PreviewTable(), DeviceApp(), DevicePostgres(), DeviceCass(), DeviceAuditLog(), DeviceExport(), Cd(), Cat(), Pwd(), ClipboardCopy(),
107
+ return [Ls(), PreviewTable(), DeviceApp(), DevicePostgres(), DeviceCass(), DeviceAuditLog(), DeviceExport(),
108
+ Cd(), Cat(), Head(), DownloadFile(), FindLocalFiles(), FindProcesses(), Pwd(), ClipboardCopy(),
97
109
  GetParam(), SetParam(), ShowParams(), ShowKubectlCommands(), ShowLogin(), ShowAdam(), ShowHost()]
98
110
 
99
111
  def cassandra_ops() -> list[Command]:
100
- return [Cqlsh(), ShowCassandraStatus(), ShowCassandraVersion(), ShowCassandraRepairs(), ShowStorage(), ShowProcesses(),
101
- Check(), Issues(), NodeTool(), Report(), AlterTables(), Bash(),
102
- ExportTables(), ExportSelect(), ExportUse(), ShowExportDatabases(), ShowColumnCounts(),
112
+ return [Cqlsh(), DownloadCassandraLog(), ShowCassandraStatus(), ShowCassandraVersion(), ShowCassandraRepairs(), ShowStorage(), ShowProcesses(),
113
+ Check(), Issues(), NodeTool(), GenerateReport(), AlterTables(), Bash(),
114
+ ExportTables(), ExportXSelect(), ExportUse(), ShowExportDatabases(), ShowColumnCounts(),
103
115
  DropExportDatabase(), DropExportDatabases(),
104
- ShowExportSessions(), ShowExportSession(),
105
- CleanUpExportSessions(), CleanUpAllExportSessions(), ImportSession()] + \
106
- Medusa().cmd_list() + [Restart(), RollOut(), Watch()] + Reaper().cmd_list() + Repair().cmd_list()
116
+ ShowExportSessions(), ShowExportSession(), DownloadExportSession(),
117
+ CleanUpExportSessions(), CleanUpAllExportSessions(), ImportSession(), ImportCSVFiles()] + \
118
+ Medusa().cmd_list() + [RestartNodes(), RestartNode(), RestartCluster(), RollOut(), Watch()] + Reaper().cmd_list() + Repair().cmd_list() + Debug().cmd_list()
107
119
 
108
120
  def postgres_ops() -> list[Command]:
109
121
  return [Postgres(), DeployPgAgent(), UndeployPgAgent(), PostgresPg()]
@@ -115,7 +127,7 @@ class ReplCommands:
115
127
  return [Audit()] + Audit().cmd_list()
116
128
 
117
129
  def export_ops() -> list[Command]:
118
- return [ExportSelectX(), DropExportDatabase(), DropExportDatabases(), ShowColumnCounts()]
130
+ return [ExportSelect(), DropExportDatabase(), DropExportDatabases(), ShowColumnCounts()]
119
131
 
120
132
  def tools() -> list[Command]:
121
133
  return [Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(), DeployPod(), UndeployPod(), Kubectl(), Code()]
adam/repl_session.py CHANGED
@@ -1,6 +1,9 @@
1
1
  from prompt_toolkit import PromptSession
2
2
  from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
3
3
 
4
+ from adam.config import Config
5
+ from adam.utils import ConfigHolder
6
+
4
7
  class ReplSession:
5
8
  # the singleton pattern
6
9
  def __new__(cls, *args, **kwargs):
@@ -10,4 +13,9 @@ class ReplSession:
10
13
 
11
14
  def __init__(self):
12
15
  if not hasattr(self, 'prompt_session'):
13
- self.prompt_session = PromptSession(auto_suggest=AutoSuggestFromHistory())
16
+ self.prompt_session = PromptSession(auto_suggest=AutoSuggestFromHistory())
17
+ ConfigHolder().append_command_history = self.append_history
18
+
19
+ def append_history(self, entry: str):
20
+ if entry and self.prompt_session and Config().get('repl.history.push-cat-remote-log-file', True):
21
+ self.prompt_session.history.append_string(entry)
adam/repl_state.py CHANGED
@@ -3,7 +3,6 @@ from enum import Enum
3
3
  import re
4
4
  from typing import Callable
5
5
 
6
- from adam.commands.postgres.postgres_context import PostgresContext
7
6
  from adam.utils_k8s.app_clusters import AppClusters
8
7
  from adam.utils_k8s.app_pods import AppPods
9
8
  from adam.utils_k8s.cassandra_clusters import CassandraClusters
@@ -62,6 +61,7 @@ class ReplState:
62
61
  self.bash_session = bash_session
63
62
  self.remote_dir = remote_dir
64
63
  self.original_state: ReplState = None
64
+ self.pod_targetted = False
65
65
 
66
66
  self.export_session: str = None
67
67
 
@@ -82,11 +82,11 @@ class ReplState:
82
82
  msg = ''
83
83
  if self.device == ReplState.P:
84
84
  msg = f'{ReplState.P}:'
85
- pg: PostgresContext = PostgresContext.apply(self.namespace, self.pg_path) if self.pg_path else None
86
- if pg and pg.db:
87
- msg += pg.db
88
- elif pg and pg.host:
89
- msg += pg.host
85
+ host, database = self.pg_host_n_database()
86
+ if database:
87
+ msg += database
88
+ elif host:
89
+ msg += host
90
90
  elif self.device == ReplState.A:
91
91
  msg = f'{ReplState.A}:'
92
92
  if self.app_env:
@@ -314,8 +314,8 @@ class ReplState:
314
314
  if self.device != ReplState.P:
315
315
  return (False, None)
316
316
 
317
- pg: PostgresContext = PostgresContext.apply(self.namespace, self.pg_path)
318
- if not pg.db:
317
+ _, database = self.pg_host_n_database()
318
+ if not database:
319
319
  def error():
320
320
  if self.in_repl:
321
321
  log2('cd to a database first.')
@@ -346,7 +346,10 @@ class ReplState:
346
346
  if not self.export_session:
347
347
  def error():
348
348
  if self.in_repl:
349
- log2("Select an export database first with 'use' command.")
349
+ if self.device == ReplState.C:
350
+ log2("Select an export database first with 'use' command.")
351
+ else:
352
+ log2('cd to an export database first.')
350
353
  else:
351
354
  log2('* export database is missing.')
352
355
  log2()
@@ -390,9 +393,10 @@ class ReplState:
390
393
  self.pop()
391
394
  self.bash_session = None
392
395
 
393
- def push(self):
396
+ def push(self, pod_targetted=False):
394
397
  if not self.original_state:
395
398
  self.original_state = copy(self)
399
+ self.pod_targetted = pod_targetted
396
400
 
397
401
  def pop(self):
398
402
  if o := self.original_state:
@@ -407,6 +411,31 @@ class ReplState:
407
411
  self.namespace = o.namespace
408
412
 
409
413
  self.original_state = None
414
+ self.pod_targetted = False
415
+
416
+ def pg_host_n_database(self):
417
+ host = None
418
+ database = None
419
+
420
+ if self.pg_path:
421
+ host_n_db = self.pg_path.split('/')
422
+ host = host_n_db[0]
423
+ if len(host_n_db) > 1:
424
+ database = host_n_db[1]
425
+
426
+ return host, database
427
+
428
+ def with_no_pod(self):
429
+ state1 = copy(self)
430
+ state1.pod = None
431
+
432
+ return state1
433
+
434
+ def with_pod(self, pod: str):
435
+ state1 = copy(self)
436
+ state1.pod = pod
437
+
438
+ return state1
410
439
 
411
440
  class DevicePodService:
412
441
  def __init__(self, handler: 'DeviceExecHandler'):
@@ -0,0 +1,44 @@
1
+ import asyncio
2
+ from concurrent.futures import ThreadPoolExecutor
3
+ import inspect
4
+ import threading
5
+ import traceback
6
+
7
+ from adam.utils import log2, log_timing
8
+
9
+ class AsyncExecutor:
10
+ # some lib does not handle asyncio loop properly, as sync exec submit does not work, use another async loop
11
+
12
+ lock = threading.Lock()
13
+ in_queue = set()
14
+
15
+ loop: asyncio.AbstractEventLoop = None
16
+ async_exec: ThreadPoolExecutor = None
17
+
18
+ def preload(action: callable, log_key: str = None):
19
+ with AsyncExecutor.lock:
20
+ if not AsyncExecutor.loop:
21
+ AsyncExecutor.loop = asyncio.new_event_loop()
22
+ AsyncExecutor.async_exec = ThreadPoolExecutor(max_workers=6, thread_name_prefix='async')
23
+ AsyncExecutor.loop.set_default_executor(AsyncExecutor.async_exec)
24
+
25
+ async def a():
26
+ try:
27
+ arg_needed = len(action.__code__.co_varnames)
28
+
29
+ if log_key:
30
+ with log_timing(log_key):
31
+ r = action(None) if arg_needed else action()
32
+ else:
33
+ r = action(None) if arg_needed else action()
34
+ if inspect.isawaitable(r):
35
+ await r
36
+
37
+ AsyncExecutor.in_queue.remove(log_key)
38
+ except Exception as e:
39
+ log2('preloading error', e, inspect.getsourcelines(action)[0][0])
40
+ traceback.print_exc()
41
+
42
+ if log_key not in AsyncExecutor.in_queue:
43
+ AsyncExecutor.in_queue.add(log_key)
44
+ AsyncExecutor.async_exec.submit(lambda: AsyncExecutor.loop.run_until_complete(a()))
@@ -0,0 +1,286 @@
1
+ import functools
2
+ import os
3
+ from typing import Iterable
4
+ from prompt_toolkit.completion import CompleteEvent, Completer, Completion, NestedCompleter, WordCompleter
5
+ from prompt_toolkit.document import Document
6
+
7
+ from adam.config import Config
8
+ from adam.sql.async_executor import AsyncExecutor
9
+ from adam.sql.lark_parser import LarkParser
10
+ from adam.utils import debug, log_timing, offload
11
+ from adam.utils_repl.appendable_completer import AppendableCompleter
12
+ from adam.utils_repl.repl_completer import merge_completions
13
+
14
+ __all__ = [
15
+ "LarkCompleter",
16
+ ]
17
+
18
+ def default_columns(x: list[str]):
19
+ return 'id,x.,y.,z.'.split(',')
20
+
21
+ class LarkCompleter(Completer, AppendableCompleter):
22
+ SYSTEM = 'system'
23
+
24
+ def __init__(self,
25
+ dml: str = None,
26
+ expandables: dict = {},
27
+ variant: str = 'c',
28
+
29
+ name: str = None,
30
+ options_lambda: callable = None,
31
+ auto: str = 'lazy',
32
+ debug = False
33
+ ) -> None:
34
+ self.nested: NestedCompleter = None
35
+ self.options_lambda = options_lambda
36
+ if options_lambda and auto == 'lazy':
37
+ AsyncExecutor.preload(options_lambda, log_key=name)
38
+
39
+ self.variant = variant
40
+ self.parser = None
41
+ self.dml = dml
42
+ self.expandables = expandables
43
+
44
+ self.display_dict = {}
45
+ self.meta_dict = {}
46
+ self.WORD = None
47
+ self.sentence = False
48
+ self.match_middle = False
49
+ self.pattern = None
50
+
51
+ self.debug = debug
52
+
53
+ if variant:
54
+ self.parser = LarkCompleter.lark_parser(variant)
55
+ self.preload_lazy_auto_completes()
56
+
57
+ def __repr__(self):
58
+ return f"LarkCompleter.{self.variant}"
59
+
60
+ def preload_lazy_auto_completes(self):
61
+ for key, value in self.expandables.items():
62
+ if callable(value):
63
+ if self.auto_complete(key) == 'lazy':
64
+ AsyncExecutor.preload(value, log_key=key)
65
+
66
+ def from_lambda(name: str, options_lambda: callable, auto: str = 'lazy'):
67
+ return LarkCompleter(name=name, options_lambda=options_lambda, auto=auto, variant=None)
68
+
69
+ @functools.lru_cache()
70
+ def lark_parser(variant: str):
71
+ dir_path = os.path.dirname(os.path.realpath(__file__))
72
+ with open(dir_path + f"/qingl.lark") as f:
73
+ grammar: str = None
74
+ with log_timing(f'lark.{variant}.file-read'):
75
+ grammar = f.read()
76
+
77
+ common_contexts = {
78
+ 'cd_command.file_name': 'direct-dirs',
79
+ }
80
+
81
+ if variant in ['a', 'c0', 'p0', 'x0']:
82
+ grammar = grammar.replace('start: statement_sequence', f'start: qing_{variant}_statement')
83
+ contexts_by_path = {
84
+ } | common_contexts
85
+
86
+ debug(f'* GRAMMAR replaced to start: qing_{variant}_statement')
87
+
88
+ return LarkParser(grammar, contexts_by_path)
89
+ elif variant == 'system':
90
+ grammar = grammar.replace('start: statement_sequence', f'start: qing_{variant}_statement')
91
+ contexts_by_path = {
92
+ } | common_contexts
93
+
94
+ return LarkParser(grammar, contexts_by_path)
95
+
96
+ grammar = grammar.replace('qing_statement: qing_p_statement', f'qing_statement: qing_{variant}_statement')
97
+ debug(f'* GRAMMAR replaced to qing_statement: qing_{variant}_statement')
98
+
99
+ bash_contexts = {
100
+ 'bash_statement.host_name': 'hosts',
101
+ 'bash_statement.bash_command': 'bash-commands',
102
+ }
103
+
104
+ contexts_by_path = {
105
+ 'describe_keyspace.keyspace_name': 'keyspaces',
106
+ 'keyspace_ref.keyspace_path.namespace_ref.identifier_ref': 'tables',
107
+ 'preview_table_statement.path.identifier_ref': 'tables',
108
+
109
+ 'insert_statement.insert_select': 'column-names',
110
+ 'update_statement.set_clause.path.identifier_ref': 'column-names',
111
+ 'update_statement.where_clause.cond.expr.path.identifier_ref': 'column-names',
112
+ 'delete_statement.where_clause.cond.expr.path.identifier_ref': 'column-names',
113
+
114
+ 'select_clause.projection.result_expr.expr.path.identifier_ref': 'columns',
115
+ 'select_from.where_clause.cond.expr.path.identifier_ref': 'columns',
116
+ 'select_from.where_clause.cond.expr.logical_term.and_expr.cond.expr.path.identifier_ref': 'columns',
117
+ 'select_from.group_by_clause.group_term.expr.path.identifier_ref': 'columns',
118
+ 'select_statement.order_by_clause.ordering_term.expr.path.identifier_ref': 'columns',
119
+ 'select_from.from_clause.from_terms.join_clause.ansi_join_clause.ansi_join_predicate.expr.path.identifier_ref': 'columns',
120
+ 'select_from.from_clause.from_terms.join_clause.ansi_join_clause.ansi_join_predicate.expr.comparison_term.relational_expr.expr.path.identifier_ref': 'columns',
121
+
122
+ 'select_from.from_clause.from_terms.from_generic.alias.identifier_ref': 'column-aliases',
123
+
124
+ 'select_statement.limit_clause.expr.literal.nbr.digit': 'limits',
125
+ } | common_contexts
126
+
127
+ if variant == 'p':
128
+ contexts_by_path = bash_contexts | contexts_by_path
129
+ elif variant == 'c':
130
+ contexts_by_path = {
131
+ 'export_table.path.identifier_ref': 'tables',
132
+ 'show_column_counts_command.path.identifier_ref': 'tables',
133
+ 'export_statement.export_tables.keyspace_name': 'keyspaces',
134
+
135
+ 'alter_tables_statement.properties.property.property_name': 'table-props',
136
+ 'alter_cql_table_statement.properties.property.property_name': 'table-props',
137
+ 'alter_tables_statement.properties.property.property_value.literal': 'table-props-value',
138
+ 'alter_cql_table_statement.properties.property.property_value.literal': 'table-props-value',
139
+
140
+ 'select_clause.projection.result_expr.expr.path.identifier_ref': 'columns',
141
+ 'export_statement.export_tables.export_table.column_name_list.column_name': 'columns',
142
+
143
+ 'consistency_statement.consistency': 'consistencies',
144
+ 'export_statement.export_to.export_database_type': 'export-database-types',
145
+ 'drop_export_database.export_database_name': 'export-databases',
146
+ 'use_export_db_statement.export_database_name': 'export-databases',
147
+ 'clean_up_export_session_statement.clean_up_export_sessions.export_session_name': 'export-sessions',
148
+ 'show_export_command.export_session_name': 'export-sessions',
149
+ 'import_statement.import_session.export_session_name': 'export-sessions-incomplete',
150
+ 'download_session_statement.export_session_name': 'export-sessions-incomplete',
151
+ } | bash_contexts | contexts_by_path
152
+ elif variant == 'l':
153
+ contexts_by_path = {
154
+ 'add_partition_action.partition_ref.partition_name': 'partition-columns',
155
+ 'show_topn_statement.topn_count': 'topn-counts',
156
+ 'show_topn_statement.topn_type': 'topn-types',
157
+ 'show_topn_statement.topn_window': 'topn-windows'
158
+ } | contexts_by_path
159
+ elif variant == 'x':
160
+ contexts_by_path = {
161
+ 'show_column_counts_command.path.identifier_ref': 'tables',
162
+ 'drop_export_database.export_database_name': 'export-databases',
163
+ 'use_export_db_statement.export_database_name': 'export-databases',
164
+ } | contexts_by_path
165
+
166
+ grammar = grammar.replace('select_clause: "SELECT"i hint_comment? projection', 'select_clause: ("SELECT"i | "XELECT"i) hint_comment? projection')
167
+
168
+ with offload():
169
+ with open('/tmp/grammar.lark', 'wt') as f:
170
+ f.write(grammar)
171
+
172
+ return LarkParser(grammar, contexts_by_path)
173
+
174
+ def get_completions(
175
+ self, document: Document, complete_event: CompleteEvent
176
+ ) -> Iterable[Completion]:
177
+ if not self.nested and self.options_lambda:
178
+ # for lazy completions
179
+ self.nested = NestedCompleter.from_nested_dict(self.options_lambda())
180
+
181
+ nested_words = set()
182
+
183
+ if self.nested:
184
+ # from NestedCompleter
185
+
186
+ # Split document.
187
+ text = document.text_before_cursor.lstrip()
188
+ stripped_len = len(document.text_before_cursor) - len(text)
189
+
190
+ # If there is a space, check for the first term, and use a
191
+ # subcompleter.
192
+ if " " in text:
193
+ first_term = text.split()[0]
194
+ completer = self.nested.options.get(first_term)
195
+
196
+ # If we have a sub completer, use this for the completions.
197
+ if completer is not None:
198
+ remaining_text = text[len(first_term) :].lstrip()
199
+ move_cursor = len(text) - len(remaining_text) + stripped_len
200
+
201
+ new_document = Document(
202
+ remaining_text,
203
+ cursor_position=document.cursor_position - move_cursor,
204
+ )
205
+
206
+ for c in completer.get_completions(new_document, complete_event):
207
+ nested_words.add(c.text)
208
+ yield c
209
+
210
+ # No space in the input: behave exactly like `WordCompleter`.
211
+ else:
212
+ completer = WordCompleter(
213
+ list(self.nested.options.keys()), ignore_case=self.nested.ignore_case
214
+ )
215
+ for c in completer.get_completions(document, complete_event):
216
+ nested_words.add(c.text)
217
+ yield c
218
+
219
+ if self.parser:
220
+ full = document.text_before_cursor
221
+ if self.dml:
222
+ full = self.dml + ' ' + full
223
+
224
+ words0 = []
225
+ words1 = []
226
+ context = {}
227
+ for word in self.parser.next_terminals(full, context=context):
228
+ if ex := self.expandable(word):
229
+ if ex in self.expandables:
230
+ e = self.expandables[ex]
231
+ if callable(e):
232
+ if self.auto_complete(ex) != 'off':
233
+ ctx = None
234
+ if 'last-id' in context:
235
+ ctx = context['last-id']
236
+ e = e(ctx)
237
+ words0.extend(e)
238
+ else:
239
+ words0.extend(e)
240
+ else:
241
+ words1.append(word)
242
+ words = words0 + words1
243
+
244
+ word_before_cursor = document.get_word_before_cursor(
245
+ WORD=self.WORD, pattern=self.pattern
246
+ )
247
+
248
+ word_before_cursor = word_before_cursor.lower()
249
+
250
+ def word_matches(word: str) -> bool:
251
+ return word.lower().startswith(word_before_cursor)
252
+
253
+ for word in words:
254
+ if word_matches(word) and word not in nested_words:
255
+ display = self.display_dict.get(word, word)
256
+ display_meta = self.meta_dict.get(word, "")
257
+ yield Completion(
258
+ word,
259
+ -len(word_before_cursor),
260
+ display=display,
261
+ display_meta=display_meta,
262
+ )
263
+
264
+ def completions_for_nesting(self, dml: str = None):
265
+ if dml:
266
+ return {dml: LarkCompleter(dml, expandables=self.expandables, variant=self.variant)}
267
+
268
+ return {
269
+ word.text.lower(): LarkCompleter(word.text, expandables=self.expandables, variant=self.variant)
270
+ for word in self.get_completions(Document(''), None)
271
+ }
272
+
273
+ def expandable(self, word: str):
274
+ return word.strip('`') if word.startswith('`') else None
275
+
276
+ def append_completions(self, key: str, value: dict[str, any]):
277
+ if isinstance(value, LarkCompleter) and self.variant == value.variant:
278
+ return
279
+
280
+ if self.nested:
281
+ self.nested = NestedCompleter.from_nested_dict(merge_completions(self.nested.options, value))
282
+ else:
283
+ self.nested = NestedCompleter.from_nested_dict(value)
284
+
285
+ def auto_complete(self, key: str, default = 'lazy'):
286
+ return Config().get(f'auto-complete.{self.variant}.{key}', default=default)