|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +from io import TextIOWrapper |
| 4 | +import logging |
| 5 | +import os |
| 6 | +import threading |
| 7 | +from typing import IO, Literal |
| 8 | + |
| 9 | +from prompt_toolkit.formatted_text import to_formatted_text |
| 10 | +from prompt_toolkit.shortcuts import PromptSession |
| 11 | +import sqlparse |
| 12 | + |
| 13 | +from mycli.app_state import ( |
| 14 | + AppStateMixin, |
| 15 | + configure_prompt_state, |
| 16 | + destructive_keywords_from_config, |
| 17 | + ensure_my_cnf_sections, |
| 18 | + llm_prompt_truncation, |
| 19 | + normalize_ssl_mode, |
| 20 | +) |
| 21 | +from mycli.cli_args import DEFAULT_PROMPT |
| 22 | +from mycli.client_commands import ClientCommandsMixin |
| 23 | +from mycli.client_connection import ClientConnectionMixin |
| 24 | +from mycli.client_query import ClientQueryMixin |
| 25 | +from mycli.main_modes import repl as repl_package |
| 26 | +from mycli.output import OutputMixin |
| 27 | +from mycli.packages import special |
| 28 | +from mycli.packages.special.favoritequeries import FavoriteQueries |
| 29 | +from mycli.packages.tabular_output import sql_format |
| 30 | +from mycli.sqlexecute import SQLExecute |
| 31 | +from mycli.types import Query |
| 32 | + |
| 33 | +sqlparse.engine.grouping.MAX_GROUPING_DEPTH = None # type: ignore[assignment] |
| 34 | +sqlparse.engine.grouping.MAX_GROUPING_TOKENS = None # type: ignore[assignment] |
| 35 | + |
| 36 | + |
| 37 | +class MyCli(AppStateMixin, OutputMixin, ClientCommandsMixin, ClientConnectionMixin, ClientQueryMixin): |
| 38 | + default_prompt = DEFAULT_PROMPT |
| 39 | + default_prompt_splitln = "\\u@\\h\\n(\\t):\\d>" |
| 40 | + max_len_prompt = 45 |
| 41 | + defaults_suffix = None |
| 42 | + prompt_lines: int |
| 43 | + sqlexecute: SQLExecute | None |
| 44 | + numeric_alignment: str |
| 45 | + |
| 46 | + # In order of being loaded. Files lower in list override earlier ones. |
| 47 | + cnf_files: list[str | IO[str]] = [ |
| 48 | + "/etc/my.cnf", |
| 49 | + "/etc/mysql/my.cnf", |
| 50 | + "/usr/local/etc/my.cnf", |
| 51 | + os.path.expanduser("~/.my.cnf"), |
| 52 | + ] |
| 53 | + |
| 54 | + # check XDG_CONFIG_HOME exists and not an empty string |
| 55 | + xdg_config_home = os.environ.get("XDG_CONFIG_HOME", "~/.config") |
| 56 | + system_config_files: list[str | IO[str]] = [ |
| 57 | + "/etc/myclirc", |
| 58 | + os.path.join(os.path.expanduser(xdg_config_home), "mycli", "myclirc"), |
| 59 | + ] |
| 60 | + |
| 61 | + pwd_config_file = os.path.join(os.getcwd(), ".myclirc") |
| 62 | + |
| 63 | + def __init__( |
| 64 | + self, |
| 65 | + sqlexecute: SQLExecute | None = None, |
| 66 | + prompt: str | None = None, |
| 67 | + toolbar_format: str | None = None, |
| 68 | + logfile: TextIOWrapper | Literal[False] | None = None, |
| 69 | + defaults_suffix: str | None = None, |
| 70 | + defaults_file: str | None = None, |
| 71 | + login_path: str | None = None, |
| 72 | + auto_vertical_output: bool = False, |
| 73 | + warn: bool | None = None, |
| 74 | + myclirc: str = "~/.myclirc", |
| 75 | + show_warnings: bool | None = None, |
| 76 | + cli_verbosity: int = 0, |
| 77 | + ) -> None: |
| 78 | + self.sqlexecute = sqlexecute |
| 79 | + self.logfile = logfile |
| 80 | + self.defaults_suffix = defaults_suffix |
| 81 | + self.login_path = login_path |
| 82 | + self.toolbar_error_message: str | None = None |
| 83 | + self.prompt_session: PromptSession | None = None |
| 84 | + self._keepalive_counter = 0 |
| 85 | + self.keepalive_ticks: int | None = 0 |
| 86 | + self.sandbox_mode: bool = False |
| 87 | + |
| 88 | + # self.cnf_files is a class variable that stores the list of mysql |
| 89 | + # config files to read in at launch. |
| 90 | + # If defaults_file is specified then override the class variable with |
| 91 | + # defaults_file. |
| 92 | + if defaults_file: |
| 93 | + self.cnf_files = [defaults_file] |
| 94 | + |
| 95 | + # Load config. |
| 96 | + config_files: list[str | IO[str]] = self.system_config_files + [myclirc] + [self.pwd_config_file] |
| 97 | + from mycli import main as main_module |
| 98 | + |
| 99 | + c = self.config = main_module.read_config_files(config_files) |
| 100 | + # this parallel config exists to |
| 101 | + # * compare with my.cnf |
| 102 | + # * support the --checkup feature |
| 103 | + # todo: after removing my.cnf, create the parallel configs only when --checkup is set |
| 104 | + self.config_without_package_defaults = main_module.read_config_files(config_files, ignore_package_defaults=True) |
| 105 | + # this parallel config exists to compare with my.cnf support the --checkup feature |
| 106 | + self.config_without_user_options = main_module.read_config_files(config_files, ignore_user_options=True) |
| 107 | + self.multi_line = c["main"].as_bool("multi_line") |
| 108 | + self.key_bindings = c["main"]["key_bindings"] |
| 109 | + self.emacs_ttimeoutlen = c['keys'].as_float('emacs_ttimeoutlen') |
| 110 | + self.vi_ttimeoutlen = c['keys'].as_float('vi_ttimeoutlen') |
| 111 | + special.set_timing_enabled(c["main"].as_bool("timing")) |
| 112 | + special.set_show_favorite_query(c["main"].as_bool("show_favorite_query")) |
| 113 | + if show_warnings is not None: |
| 114 | + special.set_show_warnings_enabled(show_warnings) |
| 115 | + else: |
| 116 | + special.set_show_warnings_enabled(c['main'].as_bool('show_warnings')) |
| 117 | + self.beep_after_seconds = float(c["main"]["beep_after_seconds"] or 0) |
| 118 | + self.default_keepalive_ticks = c['connection'].as_int('default_keepalive_ticks') |
| 119 | + |
| 120 | + FavoriteQueries.instance = FavoriteQueries.from_config(self.config) |
| 121 | + |
| 122 | + self.dsn_alias: str | None = None |
| 123 | + self.main_formatter = main_module.TabularOutputFormatter(format_name=c["main"]["table_format"]) |
| 124 | + self.redirect_formatter = main_module.TabularOutputFormatter(format_name=c["main"].get("redirect_format", "csv")) |
| 125 | + sql_format.register_new_formatter(self.main_formatter) |
| 126 | + sql_format.register_new_formatter(self.redirect_formatter) |
| 127 | + self.main_formatter.mycli = self |
| 128 | + self.redirect_formatter.mycli = self |
| 129 | + self.syntax_style = c["main"]["syntax_style"] |
| 130 | + self.verbosity = -1 if c["main"].as_bool("less_chatty") else 0 |
| 131 | + if cli_verbosity: |
| 132 | + self.verbosity = cli_verbosity |
| 133 | + self.cli_style = c["colors"] |
| 134 | + self.ptoolkit_style = main_module.style_factory_ptoolkit(self.syntax_style, self.cli_style) |
| 135 | + self.helpers_style = main_module.style_factory_helpers(self.syntax_style, self.cli_style) |
| 136 | + self.helpers_warnings_style = main_module.style_factory_helpers(self.syntax_style, self.cli_style, warnings=True) |
| 137 | + self.wider_completion_menu = c["main"].as_bool("wider_completion_menu") |
| 138 | + c_dest_warning = c["main"].as_bool("destructive_warning") |
| 139 | + self.destructive_warning = c_dest_warning if warn is None else warn |
| 140 | + self.login_path_as_host = c["main"].as_bool("login_path_as_host") |
| 141 | + self.post_redirect_command = c['main'].get('post_redirect_command') |
| 142 | + self.null_string = c['main'].get('null_string') |
| 143 | + self.numeric_alignment = c['main'].get('numeric_alignment', 'right') or 'right' |
| 144 | + self.binary_display = c['main'].get('binary_display') |
| 145 | + self.llm_prompt_field_truncate, self.llm_prompt_section_truncate = llm_prompt_truncation(c) |
| 146 | + |
| 147 | + self.ssl_mode, ssl_mode_error = normalize_ssl_mode(c) |
| 148 | + if ssl_mode_error: |
| 149 | + self.echo(ssl_mode_error, err=True, fg="red") |
| 150 | + |
| 151 | + # read from cli argument or user config file |
| 152 | + self.auto_vertical_output = auto_vertical_output or c["main"].as_bool("auto_vertical_output") |
| 153 | + |
| 154 | + # Write user config if system config wasn't the last config loaded. |
| 155 | + if c.filename not in self.system_config_files and not os.path.exists(myclirc): |
| 156 | + main_module.write_default_config(myclirc) |
| 157 | + |
| 158 | + # audit log |
| 159 | + if self.logfile is None and "audit_log" in c["main"]: |
| 160 | + try: |
| 161 | + self.logfile = open(os.path.expanduser(c["main"]["audit_log"]), "a") |
| 162 | + except (IOError, OSError): |
| 163 | + self.echo("Error: Unable to open the audit log file. Your queries will not be logged.", err=True, fg="red") |
| 164 | + self.logfile = False |
| 165 | + |
| 166 | + self.completion_refresher = main_module.CompletionRefresher() |
| 167 | + self.prefetch_schemas_mode = c["main"].get("prefetch_schemas_mode", "always") or "always" |
| 168 | + raw_prefetch_list = c["main"].as_list("prefetch_schemas_list") if "prefetch_schemas_list" in c["main"] else [] |
| 169 | + self.prefetch_schemas_list = [s.strip() for s in raw_prefetch_list if s and s.strip()] |
| 170 | + self.schema_prefetcher = main_module.SchemaPrefetcher(self) |
| 171 | + |
| 172 | + self.logger = logging.getLogger(__name__) |
| 173 | + self.initialize_logging() |
| 174 | + |
| 175 | + keyword_casing = c["main"].get("keyword_casing", "auto") |
| 176 | + |
| 177 | + self.highlight_preview = c['search'].as_bool('highlight_preview') |
| 178 | + |
| 179 | + self.query_history: list[Query] = [] |
| 180 | + |
| 181 | + # Initialize completer. |
| 182 | + self.smart_completion = c["main"].as_bool("smart_completion") |
| 183 | + self.completer = main_module.SQLCompleter( |
| 184 | + self.smart_completion, supported_formats=self.main_formatter.supported_formats, keyword_casing=keyword_casing |
| 185 | + ) |
| 186 | + self._completer_lock = threading.Lock() |
| 187 | + |
| 188 | + self.min_completion_trigger = c["main"].as_int("min_completion_trigger") |
| 189 | + # a hack, pending a better way to handle settings and state |
| 190 | + repl_package.MIN_COMPLETION_TRIGGER = self.min_completion_trigger |
| 191 | + self.last_prompt_message = to_formatted_text('') |
| 192 | + self.last_custom_toolbar_message = to_formatted_text('') |
| 193 | + |
| 194 | + # Register custom special commands. |
| 195 | + self.register_special_commands() |
| 196 | + |
| 197 | + # Load .mylogin.cnf if it exists. |
| 198 | + mylogin_cnf_path = main_module.get_mylogin_cnf_path() |
| 199 | + if mylogin_cnf_path: |
| 200 | + mylogin_cnf = main_module.open_mylogin_cnf(mylogin_cnf_path) |
| 201 | + if mylogin_cnf_path and mylogin_cnf: |
| 202 | + # .mylogin.cnf gets read last, even if defaults_file is specified. |
| 203 | + self.cnf_files.append(mylogin_cnf) |
| 204 | + elif mylogin_cnf_path and not mylogin_cnf: |
| 205 | + # There was an error reading the login path file. |
| 206 | + print("Error: Unable to read login path file.") |
| 207 | + |
| 208 | + self.my_cnf = main_module.read_config_files(self.cnf_files, list_values=False) |
| 209 | + ensure_my_cnf_sections(self.my_cnf) |
| 210 | + prompt_cnf = self.read_my_cnf(self.my_cnf, ["prompt"])["prompt"] |
| 211 | + configure_prompt_state(self, c, prompt, prompt_cnf, toolbar_format) |
| 212 | + self.prompt_session = None |
| 213 | + self.destructive_keywords = destructive_keywords_from_config(c) |
| 214 | + special.set_destructive_keywords(self.destructive_keywords) |
| 215 | + |
| 216 | + def close(self) -> None: |
| 217 | + if hasattr(self, 'schema_prefetcher'): |
| 218 | + self.schema_prefetcher.stop() |
| 219 | + if self.sqlexecute is not None: |
| 220 | + self.sqlexecute.close() |
| 221 | + |
| 222 | + def run_cli(self) -> None: |
| 223 | + from mycli import main as main_module |
| 224 | + |
| 225 | + main_module.main_repl(self) |
0 commit comments