Merge pull request #1834 from dbcli/RW/migrate-show-warnings-from-main-to-special

Move show_warnings from `main.py` to `special/iocommands.py`

Author: rolandwalker

changelog.md

@@ -49,6 +49,7 @@ Internal
 * Move CLI utilities to a new `cli_utils.py`.
 * Move keybinding utilities to a new `key_binding_utils.py`.
 * Move interactive utilities to `interactive_utils.py`.
+* Move special commands out of `main.py`.
 * Modernize orthography of prompt_toolkit filters.
 * Pin all GitHub Actions to hashes.
 

mycli/main.py

@@ -144,6 +144,7 @@ def __init__(
         auto_vertical_output: bool = False,
         warn: bool | None = None,
         myclirc: str = "~/.myclirc",
+        show_warnings: bool | None = None,
     ) -> None:
         self.sqlexecute = sqlexecute
         self.logfile = logfile
@@ -178,6 +179,10 @@ def __init__(
         self.vi_ttimeoutlen = c['keys'].as_float('vi_ttimeoutlen')
         special.set_timing_enabled(c["main"].as_bool("timing"))
         special.set_show_favorite_query(c["main"].as_bool("show_favorite_query"))
+        if show_warnings is not None:
+            special.set_show_warnings_enabled(show_warnings)
+        else:
+            special.set_show_warnings_enabled(c['main'].as_bool('show_warnings'))
         self.beep_after_seconds = float(c["main"]["beep_after_seconds"] or 0)
         self.default_keepalive_ticks = c['connection'].as_int('default_keepalive_ticks')
 
@@ -223,7 +228,6 @@ def __init__(
 
         # read from cli argument or user config file
         self.auto_vertical_output = auto_vertical_output or c["main"].as_bool("auto_vertical_output")
-        self.show_warnings = c["main"].as_bool("show_warnings")
 
         # Write user config if system config wasn't the last config loaded.
         if c.filename not in self.system_config_files and not os.path.exists(myclirc):
@@ -328,22 +332,6 @@ def register_special_commands(self) -> None:
             aliases=["\\Tr"],
             case_sensitive=True,
         )
-        special.register_special_command(
-            self.disable_show_warnings,
-            "nowarnings",
-            "nowarnings",
-            "Disable automatic warnings display.",
-            aliases=["\\w"],
-            case_sensitive=True,
-        )
-        special.register_special_command(
-            self.enable_show_warnings,
-            "warnings",
-            "warnings",
-            "Enable automatic warnings display.",
-            aliases=["\\W"],
-            case_sensitive=True,
-        )
         special.register_special_command(
             self.execute_from_file, "source", "source <filename>", "Execute queries from a file.", aliases=["\\."]
         )
@@ -363,16 +351,6 @@ def manual_reconnect(self, arg: str = "", **_) -> Generator[SQLResult, None, Non
         else:
             yield self.change_db(arg).send(None)
 
-    def enable_show_warnings(self, **_) -> Generator[SQLResult, None, None]:
-        self.show_warnings = True
-        msg = "Show warnings enabled."
-        yield SQLResult(status=msg)
-
-    def disable_show_warnings(self, **_) -> Generator[SQLResult, None, None]:
-        self.show_warnings = False
-        msg = "Show warnings disabled."
-        yield SQLResult(status=msg)
-
     def change_table_format(self, arg: str, **_) -> Generator[SQLResult, None, None]:
         try:
             self.main_formatter.format_name = arg
@@ -557,7 +535,6 @@ def connect(
         use_keyring: bool | None = None,
         reset_keyring: bool | None = None,
         keepalive_ticks: int | None = None,
-        show_warnings: bool | None = None,
     ) -> None:
         cnf = {
             "database": None,
@@ -587,8 +564,6 @@ def connect(
         ssl_config: dict[str, Any] = ssl or {}
         user_connection_config = self.config_without_package_defaults.get('connection', {})
         self.keepalive_ticks = keepalive_ticks
-        if show_warnings is not None:
-            self.show_warnings = show_warnings
 
         int_port = port and int(port)
         if not int_port:
@@ -1088,7 +1063,7 @@ def run_query(
                 click.echo(line, nl=new_line)
 
             # get and display warnings if enabled
-            if self.show_warnings and isinstance(result.rows, Cursor) and result.rows.warning_count > 0:
+            if special.is_show_warnings_enabled() and isinstance(result.rows, Cursor) and result.rows.warning_count > 0:
                 warnings = self.sqlexecute.run("SHOW WARNINGS")
                 for warning in warnings:
                     output = self.format_sqlresult(
@@ -1555,6 +1530,7 @@ def get_password_from_file(password_file: str | None) -> str | None:
         auto_vertical_output=cli_args.auto_vertical_output,
         warn=cli_args.warn,
         myclirc=cli_args.myclirc,
+        show_warnings=cli_args.show_warnings,
     )
 
     if cli_args.checkup:
@@ -1916,7 +1892,6 @@ def get_password_from_file(password_file: str | None) -> str | None:
         use_keyring=use_keyring,
         reset_keyring=reset_keyring,
         keepalive_ticks=keepalive_ticks,
-        show_warnings=cli_args.show_warnings,
     )
 
     if combined_init_cmd:

mycli/main_modes/repl.py

@@ -415,7 +415,7 @@ def _output_results(
         result_count += 1
         state.mutating = state.mutating or is_mutating(result.status_plain)
 
-        if mycli.show_warnings and isinstance(result.rows, Cursor) and result.rows.warning_count > 0:
+        if special.is_show_warnings_enabled() and isinstance(result.rows, Cursor) and result.rows.warning_count > 0:
             warnings = sqlexecute.run('SHOW WARNINGS')
             warnings_duration = time.time() - start
             saw_warning = False

mycli/packages/special/__init__.py

@@ -8,7 +8,9 @@
     close_tee,
     copy_query_to_clipboard,
     disable_pager,
+    disable_show_warnings,
     editor_command,
+    enable_show_warnings,
     flush_pipe_once_if_written,
     forced_horizontal,
     get_clip_query,
@@ -19,6 +21,7 @@
     is_pager_enabled,
     is_redirected,
     is_show_favorite_query,
+    is_show_warnings_enabled,
     is_timing_enabled,
     open_external_editor,
     set_delimiter,
@@ -30,6 +33,7 @@
     set_pager_enabled,
     set_redirect,
     set_show_favorite_query,
+    set_show_warnings_enabled,
     set_timing_enabled,
     split_queries,
     unset_once_if_written,
@@ -58,7 +62,9 @@
     'close_tee',
     'copy_query_to_clipboard',
     'disable_pager',
+    'disable_show_warnings',
     'editor_command',
+    'enable_show_warnings',
     'execute',
     'flush_pipe_once_if_written',
     'forced_horizontal',
@@ -71,6 +77,7 @@
     'is_llm_command',
     'is_pager_enabled',
     'is_redirected',
+    'is_show_warnings_enabled',
     'is_timing_enabled',
     'list_databases',
     'list_tables',
@@ -85,6 +92,7 @@
     'set_pager',
     'set_pager_enabled',
     'set_redirect',
+    'set_show_warnings_enabled',
     'set_timing_enabled',
     'set_show_favorite_query',
     'is_show_favorite_query',

mycli/packages/special/iocommands.py

@@ -46,6 +46,7 @@
 delimiter_command = DelimiterCommand()
 favoritequeries = FavoriteQueries(ConfigObj())
 DESTRUCTIVE_KEYWORDS: list[str] = []
+SHOW_WARNINGS_ENABLED: bool = False
 
 
 def set_favorite_queries(config):
@@ -81,6 +82,45 @@ def set_destructive_keywords(val: list[str]) -> None:
     DESTRUCTIVE_KEYWORDS = val
 
 
+def set_show_warnings_enabled(val: bool) -> None:
+    global SHOW_WARNINGS_ENABLED
+    SHOW_WARNINGS_ENABLED = val
+
+
+def is_show_warnings_enabled() -> bool:
+    return SHOW_WARNINGS_ENABLED
+
+
+@special_command(
+    'warnings',
+    'warnings',
+    'Enable automatic warnings display.',
+    arg_type=ArgType.NO_QUERY,
+    aliases=['\\W'],
+    case_sensitive=True,
+)
+def enable_show_warnings() -> Generator[SQLResult, None, None]:
+    global SHOW_WARNINGS_ENABLED
+    SHOW_WARNINGS_ENABLED = True
+    msg = "Show warnings enabled."
+    yield SQLResult(status=msg)
+
+
+@special_command(
+    'nowarnings',
+    'nowarnings',
+    'Disable automatic warnings display.',
+    arg_type=ArgType.NO_QUERY,
+    aliases=['\\w'],
+    case_sensitive=True,
+)
+def disable_show_warnings() -> Generator[SQLResult, None, None]:
+    global SHOW_WARNINGS_ENABLED
+    SHOW_WARNINGS_ENABLED = False
+    msg = 'Show warnings disabled.'
+    yield SQLResult(status=msg)
+
+
 @special_command(
     "pager",
     "pager [command]",

test/pytests/test_main.py

@@ -37,6 +37,8 @@
     PORT,
     TEMPFILE_PREFIX,
     USER,
+    DummyFormatter,
+    FakeCursorBase,
     ReusableLock,
     call_click_entrypoint_direct,
     dbtest,
@@ -2324,3 +2326,42 @@ def test_click_entrypoint_callback_covers_mycnf_underscore_fallback(monkeypatch:
 
     call_click_entrypoint_direct(main.CliArgs())
     assert any('ssl-ca = /tmp/ca.pem' in line for line in click_lines)
+
+
+def test_format_sqlresult_uses_redirect_formatter_when_redirected() -> None:
+    cli = make_bare_mycli()
+    cli.main_formatter = DummyFormatter()
+    cli.redirect_formatter = DummyFormatter()
+
+    result = SQLResult(header=['id'], rows=[(1,)], status='ok')
+    assert list(main.MyCli.format_sqlresult(cli, result, is_redirected=True)) == ['plain output']
+
+    assert cli.main_formatter.calls == []
+    assert len(cli.redirect_formatter.calls) == 1
+
+
+def test_format_sqlresult_materializes_cursor_rows_when_width_is_limited(monkeypatch: pytest.MonkeyPatch) -> None:
+    cli = make_bare_mycli()
+    cli.main_formatter = DummyFormatter()
+    rows = FakeCursorBase(rows=[(1,)], rowcount=1, description=[('id', 3)])
+    monkeypatch.setattr(main, 'Cursor', FakeCursorBase)
+
+    result = SQLResult(header=['id'], rows=cast(Any, rows), status='ok')
+    list(main.MyCli.format_sqlresult(cli, result, max_width=100))
+
+    formatted_rows = cli.main_formatter.calls[-1][0][0]
+    assert formatted_rows == [(1,)]
+
+
+def test_format_sqlresult_appends_postamble() -> None:
+    cli = make_bare_mycli()
+    result = SQLResult(header=['id'], rows=[(1,)], status='ok', postamble='done')
+
+    assert list(main.MyCli.format_sqlresult(cli, result))[-1] == 'done'
+
+
+def test_get_last_query_returns_latest_query() -> None:
+    cli = make_bare_mycli()
+    cli.query_history = [main.Query('select 1', True, False)]
+
+    assert main.MyCli.get_last_query(cli) == 'select 1'

test/pytests/test_main_modes_repl.py

@@ -527,7 +527,6 @@ def run(self, text: str) -> list[SQLResult]:
     cli.auto_vertical_output = True
     cli.prompt_session = FakePromptSession(columns=91)
     cli.beep_after_seconds = 0.1
-    cli.show_warnings = True
     state = repl_mode.ReplState()
     format_widths: list[int | None] = []
 
@@ -540,6 +539,7 @@ def format_sqlresult(result: SQLResult, **kwargs: Any) -> Iterator[str]:
     monkeypatch.setattr(repl_mode.time, 'time', lambda: next(time_values))
     monkeypatch.setattr(repl_mode.special, 'is_expanded_output', lambda: False)
     monkeypatch.setattr(repl_mode.special, 'is_redirected', lambda: False)
+    monkeypatch.setattr(repl_mode.special, 'is_show_warnings_enabled', lambda: True)
     monkeypatch.setattr(repl_mode.special, 'is_timing_enabled', lambda: True)
     monkeypatch.setattr(repl_mode, 'Cursor', FakeCursorBase)
     monkeypatch.setattr(repl_mode, 'is_select', lambda status: False)