From 899ae5f83debec4cd1a5e5b13b0aff0066214b68 Mon Sep 17 00:00:00 2001 From: Roland Walker Date: Sat, 10 Oct 2026 11:43:56 -0400 Subject: [PATCH] add "redirecting" query pseudo-state Add "redirecting" query state for shell-style redirects, /once, or /pipe_once commands. Don't try to kill a remote thread when a query is already in the "redirecting" state. The query is finished at that point already. Avoid running success hooks on interrupt as well. Do a better job of cleanup for redirection processes on interrupt. Some incidental Explorer tests are added. --- changelog.md | 6 + mycli/constants.py | 1 + mycli/main_modes/repl.py | 147 ++++-- mycli/packages/execution/background_runner.py | 12 +- .../packages/special_commands/io_commands.py | 105 +++-- mycli_test/pytests/test_background_runner.py | 61 ++- mycli_test/pytests/test_main_modes_repl.py | 446 +++++++++++++++++- .../pytests/test_special_io_commands.py | 147 +++++- 8 files changed, 830 insertions(+), 95 deletions(-) diff --git a/changelog.md b/changelog.md index 69a80bb5..5fe5676d 100644 --- a/changelog.md +++ b/changelog.md @@ -1,12 +1,18 @@ Upcoming (TBD) ============== +Features +-------- +* Add "redirecting" query state for shell-style redirects, `/once`, or `/pipe_once`. + + Bug Fixes -------- * Fix exceptions when interrupting "streaming" unbuffered queries. * Fix exceptions when interrupting "executing" unbuffered queries. * Don't try to kill a remote thread when a query is already in the "rendering" state. * Don't try to kill a remote thread when a query is already in the "transforming" state. +* Don't try to kill a remote thread when a query is already in the "redirecting" state. Documentation diff --git a/mycli/constants.py b/mycli/constants.py index 58ffada2..b80fc68a 100644 --- a/mycli/constants.py +++ b/mycli/constants.py @@ -58,6 +58,7 @@ class QueryState(enum.StrEnum): + REDIRECTING = 'redirecting' RENDERING = 'rendering' STREAMING = 'streaming' TRANSFORMING = 'transforming' diff --git a/mycli/main_modes/repl.py b/mycli/main_modes/repl.py index ee8eab67..7bbf4164 100644 --- a/mycli/main_modes/repl.py +++ b/mycli/main_modes/repl.py @@ -1,7 +1,7 @@ from __future__ import annotations -from collections.abc import Generator, Iterable -from contextlib import nullcontext +from collections.abc import Generator, Iterable, Iterator +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from datetime import datetime import functools @@ -44,6 +44,11 @@ from prompt_toolkit.output import ColorDepth from prompt_toolkit.shortcuts import CompleteStyle, PromptSession import pymysql +from pymysql.constants.CR import ( + CR_CONN_HOST_ERROR, + CR_SERVER_GONE_ERROR, + CR_SERVER_LOST, +) from pymysql.constants.SERVER_STATUS import SERVER_STATUS_IN_TRANS from pymysql.cursors import Cursor @@ -73,8 +78,8 @@ from mycli.packages.prompt_toolkit.style import style_factory_prompt_toolkit from mycli.packages.prompt_toolkit.toolbar import create_toolbar_tokens_func, get_vi_mode from mycli.packages.pygments.mycli_lexer import MyCliLexer -from mycli.packages.redirection.hybrid_redirection import get_redirect_components, is_redirect_command -from mycli.packages.special_commands.io_commands import temporary_redirect +from mycli.packages.redirection.hybrid_redirection import ShellRedirect, get_redirect_components, is_redirect_command +from mycli.packages.special_commands.io_commands import abort_redirect, temporary_redirect from mycli.packages.special_commands.special_command_utils import format_uptime, get_ssl_version, get_uptime, get_warning_count from mycli.packages.sql_result.sql_result import SQLResult from mycli.packages.utils.interactive_utils import confirm, confirm_destructive_query @@ -451,6 +456,67 @@ def _get_continuation( return [('class:continuation', continuation)] +@contextmanager +def _redirecting(mycli: 'MyCli') -> Iterator[None]: + runner = runner_for(mycli) + try: + with runner.rendering(QueryState.REDIRECTING) if runner else nullcontext(): + yield + except KeyboardInterrupt: + if runner is not None: + raise + raise QueryCancelled(False) from None + + +@contextmanager +def _output_pseudo_phase(mycli: 'MyCli') -> Iterator[None]: + runner = runner_for(mycli) + with _redirecting(mycli) if special_commands.is_redirected() else runner.rendering() if runner else nullcontext(): + yield + + +def _is_reconnectable_error(error: BaseException) -> bool: + return isinstance(error, pymysql.err.InterfaceError) or ( + isinstance(error, pymysql.OperationalError) + and bool(error.args) + and error.args[0] + in ( + CR_CONN_HOST_ERROR, + CR_SERVER_GONE_ERROR, + CR_SERVER_LOST, + ) + ) + + +@contextmanager +def _redirect_output(mycli: 'MyCli', redirect: ShellRedirect | None) -> Iterator[None]: + runner = runner_for(mycli) + before_output = runner.stop_rendering if runner else None + if redirect is not None: + with temporary_redirect( + redirect.command, + redirect.file_operator, + redirect.filename, + mycli.post_redirect_command, + progress=partial(_redirecting, mycli), + before_output=before_output, + ): + with runner.rendering(QueryState.REDIRECTING) if runner else nullcontext(): + yield + else: + try: + with runner.rendering(QueryState.REDIRECTING) if runner and special_commands.is_redirected() else nullcontext(): + yield + with _redirecting(mycli) if special_commands.is_redirected() else nullcontext(): + special_commands.unset_once_if_written(mycli.post_redirect_command) + special_commands.flush_pipe_once_if_written(mycli.post_redirect_command, before_output=before_output) + except BaseException as error: + # The caller may reconnect and retry using these same output resources. + if not _is_reconnectable_error(error): + abort_redirect() + raise + + def _output_results( mycli: 'MyCli', state: ReplState, @@ -470,7 +536,7 @@ def _output_results( if first_result.command is not None and first_result.command['name'] == 'source_page': if special_commands.is_redirected(): - _output_results(mycli, state, result_iterator, start) + _output_results(mycli, state, result_iterator, start, raise_interrupts=raise_interrupts) return paged_output = _single_paged_output_results(mycli, state, result_iterator, start) try: @@ -531,7 +597,17 @@ def _output_results( max_width = None runner = runner_for(mycli) - with runner.rendering() if runner else nullcontext(): + if result_count > 0: + if runner: + runner.stop_rendering() + try: + mycli.echo('') + except KeyboardInterrupt: + if special_commands.is_redirected(): + raise QueryCancelled(False) from None + if raise_interrupts: + raise + with _output_pseudo_phase(mycli): formatted = mycli.format_sqlresult( result, is_expanded=special_commands.is_expanded_output(), @@ -547,14 +623,10 @@ def _output_results( assert mycli.prompt_session is not None mycli.prompt_session.output.bell() try: - if result_count > 0: - if runner: - runner.stop_rendering() - mycli.echo('') try: mycli.output(formatted, result) except KeyboardInterrupt: - if raise_interrupts: + if raise_interrupts or special_commands.is_redirected(): raise finally: if runner: @@ -562,7 +634,7 @@ def _output_results( if special_commands.is_timing_enabled(): mycli.output_timing(f'Time: {duration:0.03f}s') except KeyboardInterrupt: - if raise_interrupts: + if raise_interrupts or special_commands.is_redirected(): raise start = time.time() @@ -575,18 +647,21 @@ def _output_results( saw_warning = False for warning in warnings: saw_warning = True - formatted = mycli.format_sqlresult( - warning, - is_expanded=special_commands.is_expanded_output(), - is_redirected=special_commands.is_redirected(), - null_string=mycli.null_string, - numeric_alignment=mycli.numeric_alignment, - binary_display=mycli.binary_display, - max_width=max_width, - is_warnings_style=True, - ) + if runner: + runner.stop_rendering() mycli.echo('') - mycli.output(formatted, warning, is_warnings_style=True) + with _output_pseudo_phase(mycli): + formatted = mycli.format_sqlresult( + warning, + is_expanded=special_commands.is_expanded_output(), + is_redirected=special_commands.is_redirected(), + null_string=mycli.null_string, + numeric_alignment=mycli.numeric_alignment, + binary_display=mycli.binary_display, + max_width=max_width, + is_warnings_style=True, + ) + mycli.output(formatted, warning, is_warnings_style=True) if saw_warning and special_commands.is_timing_enabled(): mycli.output_timing(f'Time: {warnings_duration:0.03f}s', is_warnings_style=True) @@ -895,6 +970,7 @@ def _one_iteration( return original_text = text + redirect: ShellRedirect | None = None try: polars_pipeline = parse_polars_transform(text) except PolarsTransformError as exc: @@ -905,13 +981,7 @@ def _one_iteration( elif is_redirect_command(text): sql_part, command_part, file_operator_part, file_part = get_redirect_components(text) text = sql_part or '' - try: - special_commands.set_redirect(command_part, file_operator_part, file_part) - except (FileNotFoundError, OSError, RuntimeError) as e: - mycli.logger.error('sql: %r, error: %r', text, e) - mycli.logger.error('traceback: %r', traceback.format_exc()) - mycli.echo(str(e), err=True, fg='red') - return + redirect = ShellRedirect(command_part, file_operator_part, file_part) if mycli.sandbox_mode and not is_sandbox_allowed(text): mycli.echo( @@ -985,11 +1055,7 @@ def _one_iteration( special_commands.set_expanded_output(True) redirect = polars_pipeline.shell_redirect try: - with ( - temporary_redirect(redirect.command, redirect.file_operator, redirect.filename, mycli.post_redirect_command) - if redirect is not None - else nullcontext() - ): + with _redirect_output(mycli, redirect): _output_results(mycli, state, iter([polars_result]), start, raise_interrupts=redirect is not None) except KeyboardInterrupt: raise QueryCancelled(False) from None @@ -999,9 +1065,8 @@ def _one_iteration( polars_pipeline.output_path, ) else: - _output_results(mycli, state, results, start) - special_commands.unset_once_if_written(mycli.post_redirect_command) - special_commands.flush_pipe_once_if_written(mycli.post_redirect_command) + with _redirect_output(mycli, redirect): + _output_results(mycli, state, results, start, raise_interrupts=redirect is not None) successful = True except QueryCancelled as exc: mycli.echo('Query cancelled.', err=True, fg='blue') @@ -1011,7 +1076,7 @@ def _one_iteration( except pymysql.err.InterfaceError: if not mycli.reconnect(): return - _one_iteration(mycli, state, original_text if polars_pipeline is not None else text) + _one_iteration(mycli, state, original_text) return except EOFError as e: raise e @@ -1052,10 +1117,10 @@ def _one_iteration( err=True, fg='red', ) - elif e1.args[0] in (2003, 2006, 2013): + elif _is_reconnectable_error(e1): if not mycli.reconnect(): return - _one_iteration(mycli, state, original_text if polars_pipeline is not None else text) + _one_iteration(mycli, state, original_text) return else: mycli.logger.error('sql: %r, error: %r', text, e1) diff --git a/mycli/packages/execution/background_runner.py b/mycli/packages/execution/background_runner.py index afa1dcfd..c5f59cec 100644 --- a/mycli/packages/execution/background_runner.py +++ b/mycli/packages/execution/background_runner.py @@ -138,13 +138,17 @@ def __init__(self, show_state_interval: float = 0.5) -> None: self._render_thread: Thread | None = None @contextmanager - def rendering(self, state: QueryState = QueryState.RENDERING) -> Iterator[None]: + def rendering(self, state: QueryState | None = None) -> Iterator[None]: with self._display_lock: outer = self._render_depth == 0 self._render_depth += 1 previous_state = self._render_state - self._render_state = state - if outer: + self._render_state = state if state is not None else previous_state + if self._render_state == QueryState.REDIRECTING and self.started is None: + self.started = monotonic() + self._next_update = self.started + self.interval + # Explicit phases may resume progress after terminal output stopped it. + if outer or (state is not None and self._render_thread is None): with self._display_lock: self._output_started = False self._render_stop.clear() @@ -154,7 +158,7 @@ def rendering(self, state: QueryState = QueryState.RENDERING) -> Iterator[None]: try: yield except KeyboardInterrupt: - if self._local_state() in (QueryState.RENDERING, QueryState.TRANSFORMING) and not self._busy: + if self._local_state() in (QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING) and not self._busy: raise QueryCancelled(False) from None raise finally: diff --git a/mycli/packages/special_commands/io_commands.py b/mycli/packages/special_commands/io_commands.py index 2dc905ca..26c90142 100644 --- a/mycli/packages/special_commands/io_commands.py +++ b/mycli/packages/special_commands/io_commands.py @@ -1,6 +1,7 @@ from __future__ import annotations -from contextlib import contextmanager +from collections.abc import Callable +from contextlib import AbstractContextManager, contextmanager, nullcontext import locale import logging import os @@ -55,6 +56,7 @@ 'stdin': [], 'stdout_file': None, 'stdout_mode': None, + 'process_group': False, } delimiter_command = DelimiterCommand() favorite_queries = FavoriteQueries(ConfigObj()) @@ -368,37 +370,30 @@ def temporary_redirect( file_operator_part: str | None, file_part: str | None, post_redirect_command: str | None, + *, + progress: Callable[[], AbstractContextManager[None]] = nullcontext, + before_output: Callable[[], None] | None = None, ) -> Iterator[None]: - """Own a transform's redirection without leaking resources into the next query.""" + """Own a hybrid redirection without leaking resources into the next query.""" global once_file, written_to_once_file previous_once, previous_written = once_file, written_to_once_file previous_pipe = dict(PIPE_ONCE) process_group = not WIN _reset_one_time_redirects() try: - set_redirect(command_part, file_operator_part, file_part, start_new_session=process_group) + with progress(): + set_redirect(command_part, file_operator_part, file_part, start_new_session=process_group) yield - if once_file is not None: - filename = once_file.name - once_file.close() - once_file = None - _run_post_redirect_hook(post_redirect_command, filename) - flush_pipe_once_if_written(post_redirect_command, force=True, process_group=process_group) - finally: - try: + with progress(): if once_file is not None: + filename = once_file.name once_file.close() - if process := PIPE_ONCE['process']: - _kill_pipe_process(process, process_group=process_group) - try: - process.communicate(timeout=2) - except subprocess.TimeoutExpired: - for stream in (process.stdin, process.stdout, process.stderr): - if stream is not None: - stream.close() - process.wait(timeout=2) - except Exception: - logging.getLogger(__name__).debug('Redirect cleanup failed', exc_info=True) + once_file = None + _run_post_redirect_hook(post_redirect_command, filename) + flush_pipe_once_if_written(post_redirect_command, force=True, process_group=process_group, before_output=before_output) + finally: + try: + abort_redirect() finally: once_file, written_to_once_file = previous_once, previous_written PIPE_ONCE.clear() @@ -408,7 +403,22 @@ def temporary_redirect( def _reset_one_time_redirects() -> None: global once_file, written_to_once_file once_file, written_to_once_file = None, False - PIPE_ONCE.update(process=None, stdin=[], stdout_file=None, stdout_mode=None) + PIPE_ONCE.update(process=None, stdin=[], stdout_file=None, stdout_mode=None, process_group=False) + + +def abort_redirect() -> None: + """Discard active redirects without running success hooks.""" + try: + try: + if once_file is not None: + once_file.close() + finally: + if process := PIPE_ONCE['process']: + _cleanup_pipe_process(process, process_group=PIPE_ONCE['process_group']) + except (Exception, KeyboardInterrupt): + logging.getLogger(__name__).debug('Redirect cleanup failed', exc_info=True) + finally: + _reset_one_time_redirects() @special_command( @@ -980,14 +990,21 @@ def _run_post_redirect_hook(post_redirect_command: str | None, filename: str) -> return post_cmd = post_redirect_command.format(shlex.quote(filename)) try: - subprocess.run( + process = subprocess.Popen( post_cmd, shell=True, - check=True, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + start_new_session=not WIN, ) + try: + returncode = process.wait() + if returncode: + raise subprocess.CalledProcessError(returncode, post_cmd) + except BaseException: + _cleanup_pipe_process(process, process_group=not WIN) + raise except Exception as e: raise OSError(f"Redirect post hook failed: {e}") from e @@ -999,9 +1016,11 @@ def _run_post_redirect_hook(post_redirect_command: str | None, filename: str) -> aliases=[SpecialCommandAlias("\\|", case_sensitive=False)], completion_snippet='next result to subprocess', ) -def set_pipe_once(arg: str, *, start_new_session: bool = False, **_) -> list[SQLResult]: +def set_pipe_once(arg: str, *, start_new_session: bool | None = None, **_) -> list[SQLResult]: if not arg: raise OSError("pipe_once requires a command") + if start_new_session is None: + start_new_session = not WIN if WIN: # best effort, no chaining pipe_once_cmd = shlex.split(arg) @@ -1009,6 +1028,7 @@ def set_pipe_once(arg: str, *, start_new_session: bool = False, **_) -> list[SQL # to support chaining pipe_once_cmd = ['sh', '-c', arg] PIPE_ONCE['stdin'] = [] + PIPE_ONCE['process_group'] = start_new_session PIPE_ONCE['process'] = subprocess.Popen( pipe_once_cmd, stdin=subprocess.PIPE, @@ -1037,17 +1057,34 @@ def _kill_pipe_process(process: subprocess.Popen[str], *, process_group: bool) - process.kill() +def _cleanup_pipe_process(process: subprocess.Popen[Any], *, process_group: bool) -> None: + try: + _kill_pipe_process(process, process_group=process_group) + try: + process.communicate(timeout=2) + except subprocess.TimeoutExpired: + for stream in (process.stdin, process.stdout, process.stderr): + if stream is not None: + stream.close() + process.wait(timeout=2) + except (Exception, KeyboardInterrupt): + logging.getLogger(__name__).debug('Redirect cleanup failed', exc_info=True) + + def flush_pipe_once_if_written( post_redirect_command: str | None, *, force: bool = False, - process_group: bool = False, + process_group: bool | None = None, + before_output: Callable[[], None] | None = None, ) -> None: """Flush the pipe_once cmd, if lines have been written.""" if not PIPE_ONCE['process']: return if not PIPE_ONCE['stdin'] and not force: return + if process_group is None: + process_group = PIPE_ONCE['process_group'] timed_out = False try: content = '\n'.join(PIPE_ONCE['stdin']) + '\n' if PIPE_ONCE['stdin'] else '' @@ -1064,21 +1101,19 @@ def flush_pipe_once_if_written( print(stdout_data, file=f) _run_post_redirect_hook(post_redirect_command, PIPE_ONCE['stdout_file']) elif stdout_data: + if before_output is not None: + before_output() click.secho(stdout_data.rstrip('\n')) if stderr_data: + if before_output is not None: + before_output() click.secho(stderr_data.rstrip('\n'), err=True, fg='red') if timed_out or returncode: - PIPE_ONCE['process'] = None - PIPE_ONCE['stdin'] = [] - PIPE_ONCE['stdout_file'] = None - PIPE_ONCE['stdout_mode'] = None + PIPE_ONCE.update(process=None, stdin=[], stdout_file=None, stdout_mode=None, process_group=False) if timed_out: raise OSError('process timed out after 60 seconds') raise OSError(f'process exited with nonzero code {returncode}') - PIPE_ONCE['process'] = None - PIPE_ONCE['stdin'] = [] - PIPE_ONCE['stdout_file'] = None - PIPE_ONCE['stdout_mode'] = None + PIPE_ONCE.update(process=None, stdin=[], stdout_file=None, stdout_mode=None, process_group=False) @special_command( diff --git a/mycli_test/pytests/test_background_runner.py b/mycli_test/pytests/test_background_runner.py index 240d6729..bab518dd 100644 --- a/mycli_test/pytests/test_background_runner.py +++ b/mycli_test/pytests/test_background_runner.py @@ -281,7 +281,7 @@ def test_attach_selects_background_cursor(runner: BackgroundRunner, cursorclass: @pytest.mark.parametrize('cursorclass', [Cursor, SSCursor, BackgroundSSCursor]) -@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING]) +@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING]) @pytest.mark.parametrize('periodic', [False, True]) def test_local_state_label_matches_cursor_mode( runner: BackgroundRunner, @@ -803,7 +803,7 @@ def test_non_tty_suppresses_monitor(monkeypatch: pytest.MonkeyPatch) -> None: instance.close() -@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING]) +@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING]) def test_rendering_keeps_elapsed_time_advancing( runner: BackgroundRunner, monkeypatch: pytest.MonkeyPatch, @@ -830,6 +830,31 @@ def display(elapsed: float, state: str) -> None: monitor.assert_not_called() +@pytest.mark.parametrize('show_state', [False, True]) +def test_redirecting_starts_clock_without_database_work( + runner: BackgroundRunner, + monkeypatch: pytest.MonkeyPatch, + show_state: bool, +) -> None: + runner.show_state = show_state + runner.interval = 60 + monkeypatch.setattr(background_runner, 'monotonic', lambda: 10.0) + + with runner.rendering(QueryState.REDIRECTING): + assert runner.started == 10.0 + assert runner._next_update == 70.0 + + +def test_redirecting_preserves_existing_clock_and_deadline(runner: BackgroundRunner, monkeypatch: pytest.MonkeyPatch) -> None: + runner.started = 10.0 + runner._next_update = 11.0 + monkeypatch.setattr(background_runner, 'monotonic', lambda: 20.0) + + with runner.rendering(QueryState.REDIRECTING): + assert runner.started == 10.0 + assert runner._next_update == 11.0 + + def test_fetch_does_not_reset_statement_start(runner: BackgroundRunner, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(background_runner, 'monotonic', lambda: 10.0) runner.call(lambda: None) @@ -911,7 +936,7 @@ def test_new_statement_resets_update_deadline(runner: BackgroundRunner, monkeypa @pytest.mark.parametrize('error', [ValueError('formatting'), KeyboardInterrupt()]) -@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING]) +@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING]) def test_rendering_error_stops_ticker(runner: BackgroundRunner, error: BaseException, state: QueryState) -> None: runner.show_state = True expected = QueryCancelled if isinstance(error, KeyboardInterrupt) else type(error) @@ -984,13 +1009,13 @@ def test_successful_visible_query_hands_off_to_rendering(runner: BackgroundRunne assert runner.visible -@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING]) +@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING]) def test_disabled_rendering_starts_no_thread(runner: BackgroundRunner, state: QueryState) -> None: with runner.rendering(state): assert runner._render_thread is None -@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING]) +@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING]) def test_nested_rendering_interrupt_is_local(runner: BackgroundRunner, state: QueryState) -> None: outer_state = QueryState.TRANSFORMING if state == QueryState.RENDERING else QueryState.RENDERING with runner.rendering(outer_state): @@ -1003,9 +1028,11 @@ def test_nested_rendering_interrupt_is_local(runner: BackgroundRunner, state: Qu assert runner._render_depth == 0 -@pytest.mark.parametrize('phase', ['streaming', 'transforming_database', 'database']) +@pytest.mark.parametrize('phase', ['streaming', 'transforming_database', 'redirecting_database', 'database']) def test_non_rendering_interrupt_keeps_existing_handling(runner: BackgroundRunner, phase: str) -> None: state = QueryState.TRANSFORMING if phase == 'transforming_database' else QueryState.RENDERING + if phase == 'redirecting_database': + state = QueryState.REDIRECTING if phase == 'streaming': runner.attach(Connection(defer_connect=True, cursorclass=SSCursor), Mock()) runner._busy = phase != 'streaming' @@ -1027,7 +1054,7 @@ def test_nested_rendering_uses_one_ticker(runner: BackgroundRunner) -> None: assert not thread.is_alive() -@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING]) +@pytest.mark.parametrize('state', [QueryState.RENDERING, QueryState.TRANSFORMING, QueryState.REDIRECTING]) def test_pager_suppresses_rendering(runner: BackgroundRunner, state: QueryState) -> None: runner.show_state = True with runner.suspend_display(): @@ -1042,7 +1069,7 @@ def test_nested_processing_restores_outer_state_after_error(runner: BackgroundRu thread = runner._render_thread assert vars(runner)['_render_state'] == QueryState.TRANSFORMING with pytest.raises(ValueError, match='format failed'): - with runner.rendering(): + with runner.rendering(QueryState.RENDERING): assert runner._render_thread is thread assert vars(runner)['_render_state'] == QueryState.RENDERING raise ValueError('format failed') @@ -1052,6 +1079,24 @@ def test_nested_processing_restores_outer_state_after_error(runner: BackgroundRu assert not thread.is_alive() +def test_nested_rendering_inherits_redirecting(runner: BackgroundRunner) -> None: + runner.attach(Connection(defer_connect=True, cursorclass=SSCursor), Mock()) + with runner.rendering(QueryState.REDIRECTING): + with runner.rendering(): + assert runner._local_state() == QueryState.REDIRECTING + assert runner._local_state() == QueryState.STREAMING + + +def test_explicit_redirecting_resumes_stopped_progress(runner: BackgroundRunner) -> None: + runner.show_state = True + with runner.rendering(QueryState.REDIRECTING): + runner.stop_rendering() + with runner.rendering(QueryState.REDIRECTING): + assert runner._render_thread is not None + assert not runner._output_started + assert runner._render_thread is None + + def test_database_handoff_uses_active_transform_state(runner: BackgroundRunner, monkeypatch: pytest.MonkeyPatch) -> None: output = StringIO() monkeypatch.setattr(background_runner.sys, 'stderr', output) diff --git a/mycli_test/pytests/test_main_modes_repl.py b/mycli_test/pytests/test_main_modes_repl.py index 68bc4dec..76dba992 100644 --- a/mycli_test/pytests/test_main_modes_repl.py +++ b/mycli_test/pytests/test_main_modes_repl.py @@ -21,6 +21,7 @@ from mycli.output import OutputMixin from mycli.packages.execution.background_runner import BackgroundRunner, rendering_output from mycli.packages.execution.sql_execute import SQLExecute +from mycli.packages.redirection import hybrid_redirection from mycli.packages.special_commands import io_commands from mycli.packages.sql_result.sql_result import SQLResult from mycli_test.utils import make_streaming_cursor # type: ignore[attr-defined] @@ -276,6 +277,397 @@ def patch_repl_runtime_defaults(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(repl_mode, 'is_mutating', lambda status: False) +@pytest.fixture +def redirect_cli(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[Any, BackgroundRunner, Mock]]: + patch_repl_runtime_defaults(monkeypatch) + monkeypatch.setattr(repl_mode, 'is_redirect_command', hybrid_redirection.is_redirect_command) + for name in ('is_redirected', 'unset_once_if_written', 'flush_pipe_once_if_written'): + monkeypatch.setattr(repl_mode.special_commands, name, getattr(io_commands, name)) + monkeypatch.setattr(io_commands, 'once_file', None) + monkeypatch.setattr(io_commands, 'written_to_once_file', False) + monkeypatch.setattr(io_commands, 'PIPE_ONCE', {}) + io_commands._reset_one_time_redirects() + runner = BackgroundRunner(0) + runner.show_state = True + runner.interval = 0.01 + monkeypatch.setattr(runner, '_display', Mock()) + process = Mock(returncode=0) + process.communicate.return_value = ('', '') + monkeypatch.setattr(io_commands.subprocess, 'Popen', Mock(return_value=process)) + monkeypatch.setattr(io_commands, '_kill_pipe_process', Mock()) + + def run(text: str) -> Iterator[SQLResult]: + if text.startswith('/'): + yield from repl_mode.special_commands.execute(None, text) + else: + runner.call(lambda: None) + yield SQLResult(header=['id'], rows=[(1,)], status='1 row in set') + + sql = SimpleNamespace(dbname='db', connection_id=42, run=Mock(side_effect=run), background_runner=runner, connect=Mock()) + cli = make_repl_cli(sql) + cli.reconnect = Mock() + cli.redirect_formatter = TabularOutputFormatter(format_name='csv') + cli.main_formatter = TabularOutputFormatter(format_name='csv') + cli.helpers_style = cli.helpers_warnings_style = cli.prompt_toolkit_style = None + cli.explicit_pager = False + cli.get_output_margin = lambda status: 0 + cli.format_sqlresult = lambda *args, **kwargs: OutputMixin.format_sqlresult(cli, *args, **kwargs) + cli.output = lambda *args, **kwargs: OutputMixin.output(cli, *args, **kwargs) + try: + yield cli, runner, process + finally: + io_commands.abort_redirect() + runner.close() + + +def arm_test_redirect(cli: Any, mode: str, destination: Path, *, transform: bool = False) -> str: + query = 'SELECT 1' + (' .| df' if transform else '') + if mode in ('/once', '/pipe_once'): + argument = str(destination) if mode == '/once' else 'cat' + repl_mode._one_iteration(cli, repl_mode.ReplState(), f'{mode} {argument}') + assert cli.query_history[-1].successful, cli.echo_calls + assert io_commands.is_redirected() + return query + return f'{query} $| cat' if mode == '$|' else f'{query} {mode} "{destination}"' + + +@pytest.mark.parametrize('mode', ['$>', '$>>', '$|', '/once', '/pipe_once']) +@pytest.mark.parametrize('transform', [False, True]) +def test_redirect_output_uses_redirecting_state( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + transform: bool, +) -> None: + cli, runner, _ = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv', transform=transform) + observed: list[repl_mode.QueryState] = [] + original_format = cli.format_sqlresult + + @rendering_output + def format_result(client: Any, *args: Any, **kwargs: Any) -> Any: + observed.append(runner._local_state()) + return original_format(*args, **kwargs) + + cli.format_sqlresult = lambda *args, **kwargs: format_result(cli, *args, **kwargs) + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + + assert cli.query_history[-1].successful, cli.echo_calls + assert observed == [repl_mode.QueryState.REDIRECTING] + assert not io_commands.is_redirected() + assert runner._render_thread is None + + +@pytest.mark.parametrize('mode', ['$|', '/pipe_once']) +def test_redirect_pipeline_wait_has_live_progress( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli, runner, process = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + updated = Event() + monkeypatch.setattr(runner, '_display', lambda elapsed, state: updated.set() if state == 'redirecting' else None) + + def communicate(**kwargs: Any) -> tuple[str, str]: + updated.clear() + assert updated.wait(2) + assert not runner._output_started + return '', '' + + process.communicate.side_effect = communicate + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + assert cli.query_history[-1].successful, cli.echo_calls + assert runner._render_thread is None + + +@pytest.mark.parametrize('mode', ['$|', '$>', '$>>', '/once', '/pipe_once']) +def test_local_special_command_redirect_has_live_progress( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli, runner, process = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + if mode in ('/once', '/pipe_once'): + command = '/favorite list' + else: + command = '/source --special local.sql' + command.removeprefix('SELECT 1') + original_run = cli.sql_execute.run.side_effect + + def run(text: str) -> Iterator[SQLResult]: + assert text == '/source --special local.sql' + yield from original_run('/favorite list') + + cli.sql_execute.run.side_effect = run + cli.post_redirect_command = 'post {}' + monkeypatch.setattr( + io_commands.FavoriteQueries, + 'instance', + io_commands.FavoriteQueries(ConfigObj({'favorite_queries': {'sample': 'SELECT 1'}})), + raising=False, + ) + database_call = Mock(side_effect=AssertionError('Local command should not execute SQL')) + monkeypatch.setattr(runner, 'call', database_call) + updated = Event() + monkeypatch.setattr(runner, '_display', lambda elapsed, state: updated.set() if state == 'redirecting' else None) + + def wait_for_progress(*args: Any, **kwargs: Any) -> tuple[str, str]: + updated.clear() + assert updated.wait(2), 'No progress while finalizing the local redirect' + return '', '' + + wait = Mock(side_effect=wait_for_progress) + if mode in ('$|', '/pipe_once'): + process.communicate.side_effect = wait + else: + monkeypatch.setattr(io_commands, '_run_post_redirect_hook', wait) + + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + + assert cli.query_history[-1].successful, cli.echo_calls + wait.assert_called_once() + database_call.assert_not_called() + assert runner._render_thread is None + + +@pytest.mark.parametrize('mode', ['$>', '$>>', '/once']) +def test_redirect_hook_has_live_progress( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli, runner, _ = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + updated = Event() + monkeypatch.setattr(runner, '_display', lambda elapsed, state: updated.set() if state == 'redirecting' else None) + + def hook(*args: Any) -> None: + updated.clear() + assert updated.wait(2) + + hook_mock = Mock(side_effect=hook) + monkeypatch.setattr(io_commands, '_run_post_redirect_hook', hook_mock) + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + hook_mock.assert_called_once() + assert cli.query_history[-1].successful, cli.echo_calls + assert runner._render_thread is None + + +@pytest.mark.parametrize('mode', ['$>', '$>>', '$|', '/once', '/pipe_once']) +@pytest.mark.parametrize('show_state', [False, True]) +def test_redirect_output_interrupt_is_local( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + show_state: bool, +) -> None: + cli, runner, _ = redirect_cli + runner.show_state = show_state + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + cli.output = Mock(side_effect=KeyboardInterrupt) + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + assert 'Query cancelled.' in cli.echo_calls + assert not cli.query_history[-1].successful + cli.sql_execute.connect.assert_not_called() + cli.reconnect.assert_not_called() + assert runner.control is None + assert not io_commands.is_redirected() + assert runner._render_thread is None + cli.output = Mock() + repl_mode._one_iteration(cli, repl_mode.ReplState(), 'SELECT 2') + assert cli.query_history[-1].successful + + +@pytest.mark.parametrize('mode', ['$|', '/pipe_once', '$>', '/once']) +def test_redirect_finalization_interrupt_is_local( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli, runner, process = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + if mode in ('$|', '/pipe_once'): + process.communicate.side_effect = [KeyboardInterrupt(), ('', '')] + else: + monkeypatch.setattr(io_commands, '_run_post_redirect_hook', Mock(side_effect=KeyboardInterrupt)) + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + assert 'Query cancelled.' in cli.echo_calls + assert not cli.query_history[-1].successful + cli.sql_execute.connect.assert_not_called() + cli.reconnect.assert_not_called() + assert not io_commands.is_redirected() + assert runner._render_thread is None + + +@pytest.mark.parametrize('mode', ['/once', '/pipe_once']) +def test_empty_command_leaves_one_time_redirect_armed( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, +) -> None: + cli, runner, process = redirect_cli + arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + cli.sql_execute.run.side_effect = lambda text: iter([SQLResult(status='OK')]) + repl_mode._one_iteration(cli, repl_mode.ReplState(), 'SET @a = 1') + assert io_commands.is_redirected() + process.communicate.assert_not_called() + assert runner._render_thread is None + + +@pytest.mark.parametrize('mode', ['/once', '/pipe_once']) +@pytest.mark.parametrize( + 'error', + [ + pymysql.err.InterfaceError(0, ''), + pymysql.err.OperationalError(2003, 'cannot connect'), + pymysql.err.OperationalError(2006, 'server has gone away'), + pymysql.err.OperationalError(2013, 'lost connection'), + ], +) +def test_one_time_redirect_survives_database_retry( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + error: Exception, +) -> None: + cli, _, process = redirect_cli + destination = tmp_path / 'rows.csv' + command = arm_test_redirect(cli, mode, destination) + original_run = cli.sql_execute.run.side_effect + pending_file = io_commands.once_file + attempts = 0 + + def run(text: str) -> Iterator[SQLResult]: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise error + if mode == '/once': + assert io_commands.once_file is pending_file + assert pending_file is not None and not pending_file.closed + else: + assert io_commands.PIPE_ONCE['process'] is process + process.communicate.assert_not_called() + yield from original_run(text) + + cli.sql_execute.run.side_effect = run + cli.reconnect.return_value = True + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + + assert cli.query_history[-1].successful, cli.echo_calls + assert attempts == 2 + cli.reconnect.assert_called_once_with() + if mode == '/once': + assert destination.read_text().splitlines() == ['"id"', '"1"'] + else: + process.communicate.assert_called_once_with(input='"id"\n"1"\n', timeout=60) + assert not io_commands.is_redirected() + + +@pytest.mark.parametrize('mode', ['/once', '/pipe_once']) +def test_failed_reconnect_leaves_one_time_redirect_pending( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, +) -> None: + cli, _, process = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + pending_file = io_commands.once_file + + def run(text: str) -> Iterator[SQLResult]: + yield from () + raise pymysql.err.OperationalError(2006, 'server has gone away') + + cli.sql_execute.run.side_effect = run + cli.reconnect.return_value = False + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + + cli.reconnect.assert_called_once_with() + assert io_commands.is_redirected() + if mode == '/once': + assert io_commands.once_file is pending_file + assert pending_file is not None and not pending_file.closed + else: + assert io_commands.PIPE_ONCE['process'] is process + process.communicate.assert_not_called() + + +@pytest.mark.parametrize('mode', ['/once', '/pipe_once']) +@pytest.mark.parametrize('error', [pymysql.err.OperationalError(1064, 'syntax error'), ValueError('failed')]) +def test_nonretryable_error_cleans_up_one_time_redirect( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + mode: str, + error: Exception, +) -> None: + cli, _, _ = redirect_cli + command = arm_test_redirect(cli, mode, tmp_path / 'rows.csv') + + def run(text: str) -> Iterator[SQLResult]: + yield from () + raise error + + cli.sql_execute.run.side_effect = run + repl_mode._one_iteration(cli, repl_mode.ReplState(), command) + + assert not cli.query_history[-1].successful + assert not io_commands.is_redirected() + cli.reconnect.assert_not_called() + + +def test_redirecting_does_not_override_active_database_interrupt( + redirect_cli: tuple[Any, BackgroundRunner, Mock], +) -> None: + cli, runner, _ = redirect_cli + runner._busy = True + try: + with pytest.raises(KeyboardInterrupt): + with repl_mode._redirecting(cli): + raise KeyboardInterrupt + finally: + runner._busy = False + + +@pytest.mark.parametrize('reason', ['sandbox', 'declined']) +def test_rejected_query_does_not_create_redirect( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + reason: str, +) -> None: + cli, _, _ = redirect_cli + destination = tmp_path / 'rows.csv' + if reason == 'sandbox': + cli.sandbox_mode = True + else: + cli.destructive_warning = True + monkeypatch.setattr(repl_mode, 'confirm_destructive_query', lambda *args: False) + repl_mode._one_iteration(cli, repl_mode.ReplState(), f'SELECT 1 $> "{destination}"') + assert not destination.exists() + assert not io_commands.is_redirected() + + +@pytest.mark.parametrize('error', [pymysql.err.InterfaceError(0, ''), pymysql.err.OperationalError(2006, 'lost')]) +def test_hybrid_redirect_retries_complete_command( + redirect_cli: tuple[Any, BackgroundRunner, Mock], + tmp_path: Path, + error: Exception, +) -> None: + cli, _, _ = redirect_cli + cli.reconnect.return_value = True + cli.sql_execute.run.side_effect = [error, iter([SQLResult(header=['id'], rows=[(1,)])])] + destination = tmp_path / 'rows.csv' + repl_mode._one_iteration(cli, repl_mode.ReplState(), f'SELECT 1 $> "{destination}"') + assert cli.query_history[-1].successful, cli.echo_calls + assert destination.read_text().splitlines() == ['"id"', '"1"'] + cli.reconnect.assert_called_once_with() + + def test_complete_while_typing_filter_covers_threshold_and_word_rules(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(repl_mode, 'MIN_COMPLETION_TRIGGER', 3) monkeypatch.setattr(repl_mode, 'get_app', lambda: SimpleNamespace(current_buffer=SimpleNamespace(text='ab'))) @@ -974,6 +1366,56 @@ def timing(message: str) -> None: runner.close() +def test_redirecting_without_runner_converts_interrupt_to_local_cancellation() -> None: + cli = make_repl_cli(SimpleNamespace()) + + with pytest.raises(repl_mode.QueryCancelled) as raised: + with repl_mode._redirecting(cli): + raise KeyboardInterrupt + + assert not raised.value.disconnected + + +@pytest.mark.parametrize('redirected', [False, True]) +def test_output_results_propagates_interrupt_between_results(monkeypatch: pytest.MonkeyPatch, redirected: bool) -> None: + patch_repl_runtime_defaults(monkeypatch) + monkeypatch.setattr(repl_mode.special_commands, 'is_redirected', lambda: redirected) + monkeypatch.setattr(repl_mode.special_commands, 'is_show_warnings_enabled', lambda: False) + cli = make_repl_cli(SimpleNamespace()) + cli.echo = Mock(side_effect=KeyboardInterrupt) + first = SQLResult(status='first result') + second = SQLResult(status='second result') + expected_error = repl_mode.QueryCancelled if redirected else KeyboardInterrupt + + with pytest.raises(expected_error): + repl_mode._output_results(cli, repl_mode.ReplState(), [first, second], 0.0, raise_interrupts=not redirected) + + cli.echo.assert_called_once_with('') + assert [result for _, result, _ in cli.output_calls] == [first] + + +def test_output_results_clears_progress_before_warning_separator(monkeypatch: pytest.MonkeyPatch) -> None: + patch_repl_runtime_defaults(monkeypatch) + monkeypatch.setattr(repl_mode.special_commands, 'is_show_warnings_enabled', lambda: True) + monkeypatch.setattr(repl_mode, 'Cursor', FakeCursorBase) + runner = BackgroundRunner(0) + + def run(text: str) -> Iterator[SQLResult]: + assert text == 'SHOW WARNINGS' + runner.visible = True + yield SQLResult(status='warning', rows=[('warning text',)]) + + cli = make_repl_cli(SimpleNamespace(background_runner=runner, run=run)) + separator_states: list[tuple[str, bool]] = [] + cli.echo = lambda message: separator_states.append((message, runner.visible)) + rows = cast(Any, FakeCursorBase(rowcount=1, warning_count=1)) + try: + repl_mode._output_results(cli, repl_mode.ReplState(), [SQLResult(rows=rows, status='Query OK')], 0.0) + assert separator_states == [('', False)] + finally: + runner.close() + + def test_output_results_covers_watch_warning_timing_beep_and_interrupts(monkeypatch: pytest.MonkeyPatch) -> None: class FakeSQLExecute: def run(self, text: str) -> list[SQLResult]: @@ -1832,7 +2274,7 @@ def run(self, text: str) -> Iterator[SQLResult]: monkeypatch.setattr(repl_mode, 'is_redirect_command', lambda text: text == 'redirect') monkeypatch.setattr(repl_mode, 'get_redirect_components', lambda text: ('dropdb', 'tee', '>', 'out.txt')) redirects: list[tuple[Any, ...]] = [] - monkeypatch.setattr(repl_mode.special_commands, 'set_redirect', lambda *args: redirects.append(args)) + monkeypatch.setattr(io_commands, 'set_redirect', lambda *args, **kwargs: redirects.append(args)) monkeypatch.setattr( repl_mode, 'confirm_destructive_query', @@ -2709,7 +3151,7 @@ def run(self, text: str) -> Iterator[SQLResult]: cli = make_repl_cli(FakeSQLExecute()) monkeypatch.setattr(repl_mode, 'is_redirect_command', lambda text: text == 'redirect-bad') monkeypatch.setattr(repl_mode, 'get_redirect_components', lambda text: ('sql', 'tee', '>', 'out.txt')) - monkeypatch.setattr(repl_mode.special_commands, 'set_redirect', lambda *args: (_ for _ in ()).throw(RuntimeError('redirect boom'))) + monkeypatch.setattr(io_commands, 'set_redirect', lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError('redirect boom'))) repl_mode._one_iteration(cli, repl_mode.ReplState(), 'redirect-bad') assert 'redirect boom' in cli.echo_calls[-1] diff --git a/mycli_test/pytests/test_special_io_commands.py b/mycli_test/pytests/test_special_io_commands.py index ccee3905..71fd70dd 100644 --- a/mycli_test/pytests/test_special_io_commands.py +++ b/mycli_test/pytests/test_special_io_commands.py @@ -510,6 +510,33 @@ def teardown_function(): mycli.packages.special_commands.set_delimiter(";") +@pytest.mark.parametrize('enabled', [False, True]) +def test_explorer_output_setting(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None: + monkeypatch.setattr(io_commands, 'use_explorer_output', not enabled) + + io_commands.set_explorer_output(enabled) + + assert io_commands.is_explorer_output() is enabled + + +@pytest.mark.parametrize( + ('statement', 'expected'), + [ + ('', False), + (' ', False), + (r'\fs sample SELECT 1', True), + (' /fs sample SELECT 1; SELECT 2', True), + (r'\favorite save sample SELECT 1', True), + ('/FAVORITE SAVE sample SELECT 1', True), + ('/favorite run sample', False), + ('/favorite', False), + ('SELECT 1', False), + ], +) +def test_is_favorite_save_command(statement: str, expected: bool) -> None: + assert io_commands.is_favorite_save_command(statement) is expected + + def test_simple_setters_and_toggle_timing() -> None: config = {'favorite_queries': {'demo': 'select 1'}} @@ -805,7 +832,7 @@ def test_timed_out_pipe_does_not_publish_output( assert not destination.exists() hook.assert_not_called() secho.assert_called_once_with('diagnostic', err=True, fg='red') - assert io_commands.PIPE_ONCE == {'process': None, 'stdin': [], 'stdout_file': None, 'stdout_mode': None} + assert io_commands.PIPE_ONCE == {'process': None, 'stdin': [], 'stdout_file': None, 'stdout_mode': None, 'process_group': False} def test_temporary_redirect_closes_streams_after_cleanup_timeout(monkeypatch: pytest.MonkeyPatch) -> None: @@ -1890,9 +1917,9 @@ def fake_run_post_redirect_hook(command: str, filename: str) -> None: def fake_run(*args, **kwargs) -> SimpleNamespace: run_calls.append((args, kwargs)) - return SimpleNamespace(returncode=0) + return SimpleNamespace(wait=lambda: 0) - monkeypatch.setattr(io_commands.subprocess, 'run', fake_run) + monkeypatch.setattr(io_commands.subprocess, 'Popen', fake_run) io_commands._run_post_redirect_hook('', str(target)) assert run_calls == [] @@ -1900,16 +1927,16 @@ def fake_run(*args, **kwargs) -> SimpleNamespace: assert run_calls[0][0] == ('cat ' + io_commands.shlex.quote(str(target)),) assert run_calls[0][1] == { 'shell': True, - 'check': True, 'stdin': io_commands.subprocess.DEVNULL, 'stdout': io_commands.subprocess.DEVNULL, 'stderr': io_commands.subprocess.DEVNULL, + 'start_new_session': not io_commands.WIN, } def raise_run(*_args, **_kwargs): raise RuntimeError('hook failed') - monkeypatch.setattr(io_commands.subprocess, 'run', raise_run) + monkeypatch.setattr(io_commands.subprocess, 'Popen', raise_run) with pytest.raises(OSError, match='Redirect post hook failed: hook failed'): io_commands._run_post_redirect_hook('cat {}', str(target)) @@ -1927,6 +1954,115 @@ def test_run_post_redirect_hook_delegates_to_private_helper(monkeypatch) -> None assert hook_calls == [('post {}', 'output.parquet')] +@pytest.mark.parametrize('windows', [False, True]) +def test_interrupted_post_redirect_hook_is_reaped(monkeypatch: pytest.MonkeyPatch, windows: bool) -> None: + monkeypatch.setattr(io_commands, 'WIN', windows) + process = Mock() + process.wait.side_effect = KeyboardInterrupt + popen = Mock(return_value=process) + kill = Mock() + monkeypatch.setattr(io_commands.subprocess, 'Popen', popen) + monkeypatch.setattr(io_commands, '_kill_pipe_process', kill) + + with pytest.raises(KeyboardInterrupt): + io_commands._run_post_redirect_hook('post {}', 'some file.csv') + + assert popen.call_args.kwargs['start_new_session'] is not windows + kill.assert_called_once_with(process, process_group=not windows) + process.communicate.assert_called_once_with(timeout=2) + + +def test_nonzero_post_redirect_hook_reports_error(monkeypatch: pytest.MonkeyPatch) -> None: + process = Mock() + process.wait.return_value = 7 + monkeypatch.setattr(io_commands.subprocess, 'Popen', Mock(return_value=process)) + monkeypatch.setattr(io_commands, '_kill_pipe_process', Mock()) + with pytest.raises(OSError, match='Redirect post hook failed: .*exit status 7'): + io_commands._run_post_redirect_hook('post {}', 'file.csv') + + +def test_pipe_output_clears_progress_before_printing(monkeypatch: pytest.MonkeyPatch) -> None: + process = FakeProcess(stdout='rows\n', stderr='diagnostic\n', returncode=0) + io_commands.PIPE_ONCE.update(process=process, stdin=['row']) + events: list[str] = [] + monkeypatch.setattr(io_commands.click, 'secho', lambda message, **kwargs: events.append(message)) + + io_commands.flush_pipe_once_if_written(None, before_output=lambda: events.append('clear')) + + assert events == ['clear', 'rows', 'clear', 'diagnostic'] + + +@pytest.mark.parametrize('error', [OSError('close failed'), KeyboardInterrupt()]) +def test_abort_redirect_recovers_from_file_close_failure( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + error: BaseException, +) -> None: + output_file = Mock() + output_file.close.side_effect = error + process = Mock() + monkeypatch.setattr(io_commands, 'once_file', output_file) + monkeypatch.setattr(io_commands, 'written_to_once_file', True) + monkeypatch.setattr( + io_commands, + 'PIPE_ONCE', + {'process': process, 'stdin': ['row'], 'stdout_file': 'rows.csv', 'stdout_mode': 'w', 'process_group': False}, + ) + hook = Mock() + monkeypatch.setattr(io_commands, '_run_post_redirect_hook', hook) + + with caplog.at_level('DEBUG', logger=io_commands.__name__): + io_commands.abort_redirect() + + output_file.close.assert_called_once_with() + process.kill.assert_called_once_with() + process.communicate.assert_called_once_with(timeout=2) + hook.assert_not_called() + assert io_commands.once_file is None + assert io_commands.written_to_once_file is False + assert io_commands.PIPE_ONCE == { + 'process': None, + 'stdin': [], + 'stdout_file': None, + 'stdout_mode': None, + 'process_group': False, + } + assert any( + record.message == 'Redirect cleanup failed' and record.exc_info is not None and record.exc_info[1] is error + for record in caplog.records + ) + + +def test_repeated_interrupt_during_redirect_cleanup_resets_state(monkeypatch: pytest.MonkeyPatch) -> None: + process = Mock() + process.communicate.side_effect = KeyboardInterrupt + io_commands.PIPE_ONCE.update(process=process, stdin=['row']) + monkeypatch.setattr(io_commands, '_kill_pipe_process', Mock()) + io_commands.abort_redirect() + assert not io_commands.is_redirected() + assert io_commands.PIPE_ONCE['stdin'] == [] + + +@pytest.mark.skipif(os.name == 'nt', reason='POSIX shell process groups') +def test_pipe_once_abort_stops_pipeline_children(tmp_path: Path) -> None: + started, finished = tmp_path / 'started', tmp_path / 'finished' + command = f'(printf ready > {shlex.quote(str(started))}; sleep 2; touch {shlex.quote(str(finished))}) | cat' + io_commands.set_pipe_once(command) + process = io_commands.PIPE_ONCE['process'] + try: + deadline = monotonic() + 5 + while not started.exists() and monotonic() < deadline: + sleep(0.01) + assert started.exists(), 'Child did not start' + io_commands.abort_redirect() + assert process.poll() is not None + sleep(2.1) + assert not finished.exists(), 'Child continued executing after abort' + finally: + io_commands.abort_redirect() + process.communicate(timeout=5) + + def test_set_pipe_once_and_flush_short_circuits(monkeypatch) -> None: popen_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] monkeypatch.setattr(io_commands, 'WIN', True) @@ -1988,6 +2124,7 @@ def test_flush_pipe_once_timeout_and_nonzero_exit(monkeypatch, tmp_path: Path) - 'stdin': [], 'stdout_file': None, 'stdout_mode': None, + 'process_group': False, }