diff --git a/changelog.md b/changelog.md index 41523698..09df413b 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,11 @@ Upcoming (TBD) ============== +Bugfixes +--------- +* Avoid an exception when exiting while completions are being refreshed. + + Documentation --------- * Shorten/clarify `pyproject.toml` project description. diff --git a/mycli/client.py b/mycli/client.py index 1d34f2c5..15d44762 100644 --- a/mycli/client.py +++ b/mycli/client.py @@ -241,6 +241,10 @@ def _invalidate_prompt_session(self) -> None: self.prompt_session.app.invalidate() def close(self) -> None: + try: + self.completion_refresher.stop() + except Exception: + pass try: self.schema_prefetcher.stop() except Exception: diff --git a/mycli/completion_refresher.py b/mycli/completion_refresher.py index bdb00472..32875b9d 100644 --- a/mycli/completion_refresher.py +++ b/mycli/completion_refresher.py @@ -18,6 +18,9 @@ class CompletionRefresher: def __init__(self, invalidate_app: Callable[[], None] | None = None) -> None: self._completer_thread: threading.Thread | None = None self._restart_refresh = threading.Event() + self._stop_refresh = threading.Event() + self._executor_lock = threading.Lock() + self._active_executor: SQLExecute | None = None self._refresh_visible_until = 0.0 self._visibility_timer: threading.Timer | None = None self._invalidate_app = invalidate_app @@ -49,6 +52,7 @@ def refresh( if self._visibility_timer is not None: self._visibility_timer.cancel() self._visibility_timer = None + self._stop_refresh.clear() self._refresh_visible_until = monotonic() + MIN_COMPLETION_REFRESH_MESSAGE_SECONDS self._completer_thread = threading.Thread( target=self._bg_refresh, args=(executor, callbacks, completer_options), name="completion_refresh" @@ -63,6 +67,29 @@ def is_refreshing(self) -> bool: def _thread_is_alive(self) -> bool: return bool(self._completer_thread and self._completer_thread.is_alive()) + def stop(self) -> None: + """Stop and wait for an in-flight completion refresh.""" + self._stop_refresh.set() + self._restart_refresh.clear() + self._refresh_visible_until = 0.0 + if self._visibility_timer is not None: + self._visibility_timer.cancel() + self._visibility_timer = None + + with self._executor_lock: + executor = self._active_executor + if executor is not None: + try: + executor.close() + except Exception: + pass + + thread = self._completer_thread + if thread is not None and thread.is_alive() and thread is not threading.current_thread(): + thread.join() + if thread is not None and not thread.is_alive(): + self._completer_thread = None + def _bg_refresh( self, sqlexecute: SQLExecute, @@ -89,7 +116,12 @@ def _bg_refresh( self._finish_refreshing() return + with self._executor_lock: + self._active_executor = executor try: + if self._stop_refresh.is_set(): + return + # If callbacks is a single function then push it into a list. if callable(callbacks): callbacks = [callbacks] @@ -97,6 +129,8 @@ def _bg_refresh( while 1: for refresher in self.refreshers.values(): refresher(completer, executor) + if self._stop_refresh.is_set(): + return if self._restart_refresh.is_set(): self._restart_refresh.clear() break @@ -109,11 +143,23 @@ def _bg_refresh( # break statement. continue - for callback in callbacks: - callback(completer) + if not self._stop_refresh.is_set(): + for callback in callbacks: + callback(completer) + except Exception: + if not self._stop_refresh.is_set(): + raise finally: - executor.close() - self._finish_refreshing() + with self._executor_lock: + if self._active_executor is executor: + self._active_executor = None + try: + executor.close() + except Exception: + if not self._stop_refresh.is_set(): + raise + finally: + self._finish_refreshing() def _finish_refreshing(self) -> None: self._invalidate() diff --git a/test/pytests/test_client.py b/test/pytests/test_client.py index 6687ab55..410e63e8 100644 --- a/test/pytests/test_client.py +++ b/test/pytests/test_client.py @@ -391,23 +391,18 @@ def fake_read_config_file(handle: object, list_values: bool = True) -> dict[str, assert cli.mylogin_cnf == {'client': 'config'} -def test_close_stops_schema_prefetcher_and_closes_sqlexecute() -> None: +def test_close_stops_refreshers_before_closing_connection_and_tunnels() -> None: cli = MyCli.__new__(MyCli) - stopped: list[bool] = [] - closed: list[bool] = [] - ssh_tunnel_closed: list[bool] = [] - boundary_tunnel_closed: list[bool] = [] - cli.schema_prefetcher = SimpleNamespace(stop=lambda: stopped.append(True)) - cli.sqlexecute = SimpleNamespace(close=lambda: closed.append(True)) # type: ignore[assignment] - cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: ssh_tunnel_closed.append(True)) - cli.boundary_tunnel = SimpleNamespace(close=lambda: boundary_tunnel_closed.append(True)) # type: ignore[assignment] + calls: list[str] = [] + cli.completion_refresher = SimpleNamespace(stop=lambda: calls.append('completion')) + cli.schema_prefetcher = SimpleNamespace(stop=lambda: calls.append('prefetch')) + cli.sqlexecute = SimpleNamespace(close=lambda: calls.append('connection')) # type: ignore[assignment] + cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: calls.append('ssh')) + cli.boundary_tunnel = SimpleNamespace(close=lambda: calls.append('boundary')) # type: ignore[assignment] MyCli.close(cli) - assert stopped == [True] - assert closed == [True] - assert ssh_tunnel_closed == [True] - assert boundary_tunnel_closed == [True] + assert calls == ['completion', 'prefetch', 'connection', 'ssh', 'boundary'] def test_close_swallows_cleanup_errors() -> None: @@ -416,6 +411,7 @@ def test_close_swallows_cleanup_errors() -> None: def fail() -> None: raise RuntimeError('cleanup failed') + cli.completion_refresher = SimpleNamespace(stop=fail) cli.schema_prefetcher = SimpleNamespace(stop=fail) cli.sqlexecute = SimpleNamespace(close=fail) # type: ignore[assignment] cast(Any, cli).ssh_tunnel = SimpleNamespace(close=fail) @@ -425,6 +421,8 @@ def fail() -> None: def test_close_swallows_boundary_tunnel_close_error() -> None: cli = MyCli.__new__(MyCli) + cli.completion_refresher = SimpleNamespace(stop=lambda: None) + cli.schema_prefetcher = SimpleNamespace(stop=lambda: None) cli.sqlexecute = None tunnel_closed: list[bool] = [] cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: tunnel_closed.append(True)) diff --git a/test/pytests/test_completion_refresher.py b/test/pytests/test_completion_refresher.py index 063005fe..41340e5e 100644 --- a/test/pytests/test_completion_refresher.py +++ b/test/pytests/test_completion_refresher.py @@ -256,6 +256,108 @@ def test_refresh_cancels_pending_visibility_timer(monkeypatch, refresher) -> Non assert refresher._visibility_timer is None +def test_stop_interrupts_and_joins_active_refresh(monkeypatch, refresher) -> None: + refresh_started = completion_refresher.threading.Event() + executor_closed = completion_refresher.threading.Event() + callback = Mock() + timer = Mock() + + class FakeCompleter: + def __init__(self, **options) -> None: + pass + + class FakeExecutor: + def __init__(self, *args) -> None: + pass + + def close(self) -> None: + executor_closed.set() + + def blocking_refresh(completer, executor) -> None: + refresh_started.set() + assert executor_closed.wait(timeout=1) + raise completion_refresher.pymysql.err.OperationalError(2013, 'connection closed') + + monkeypatch.setattr(completion_refresher, 'SQLCompleter', FakeCompleter) + monkeypatch.setattr(completion_refresher, 'SQLExecute', FakeExecutor) + refresher.refreshers = {'blocking': blocking_refresh} + + refresher.refresh(make_sqlexecute(), callback) + assert refresh_started.wait(timeout=1) + refresher._visibility_timer = timer + + refresher.stop() + + assert executor_closed.is_set() + assert refresher._completer_thread is None + assert refresher._active_executor is None + assert refresher._restart_refresh.is_set() is False + assert refresher.is_refreshing() is False + timer.cancel.assert_called_once_with() + callback.assert_not_called() + + +def test_stop_before_executor_is_ready_prevents_refresh_and_callback(monkeypatch, refresher) -> None: + constructor_started = completion_refresher.threading.Event() + release_constructor = completion_refresher.threading.Event() + stop_finished = completion_refresher.threading.Event() + refresh = Mock() + callback = Mock() + + class FakeCompleter: + def __init__(self, **options) -> None: + pass + + class FakeExecutor: + def __init__(self, *args) -> None: + constructor_started.set() + assert release_constructor.wait(timeout=1) + + def close(self) -> None: + pass + + monkeypatch.setattr(completion_refresher, 'SQLCompleter', FakeCompleter) + monkeypatch.setattr(completion_refresher, 'SQLExecute', FakeExecutor) + refresher.refreshers = {'refresh': refresh} + + refresher.refresh(make_sqlexecute(), callback) + assert constructor_started.wait(timeout=1) + stop_thread = completion_refresher.threading.Thread(target=lambda: (refresher.stop(), stop_finished.set())) + stop_thread.start() + assert refresher._stop_refresh.wait(timeout=1) + assert stop_finished.is_set() is False + release_constructor.set() + stop_thread.join(timeout=1) + + assert stop_finished.is_set() + refresh.assert_not_called() + callback.assert_not_called() + assert refresher._completer_thread is None + + +def test_stop_tolerates_executor_close_error_without_worker(refresher) -> None: + executor = Mock() + executor.close.side_effect = RuntimeError('close failed') + refresher._active_executor = executor + + refresher.stop() + + executor.close.assert_called_once_with() + assert refresher._completer_thread is None + + +def test_stop_does_not_join_current_worker(monkeypatch, refresher) -> None: + thread = Mock() + thread.is_alive.return_value = True + refresher._completer_thread = thread + monkeypatch.setattr(completion_refresher.threading, 'current_thread', lambda: thread) + + refresher.stop() + + thread.join.assert_not_called() + assert refresher._completer_thread is thread + + def test_finish_refreshing_schedules_delayed_invalidation_before_deadline(monkeypatch, refresher) -> None: now = 10.0 monkeypatch.setattr(completion_refresher, 'monotonic', lambda: now) @@ -474,6 +576,72 @@ def __init__(self, *args) -> None: callback.assert_not_called() +def test_bg_refresh_stops_after_current_refresher(monkeypatch, refresher) -> None: + callback = Mock() + executor = Mock() + + def stop_refresh(completer, active_executor) -> None: + refresher._stop_refresh.set() + + monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock()) + monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor)) + refresher.refreshers = {'stop': stop_refresh} + + refresher._bg_refresh(make_sqlexecute(), callback, {}) + + callback.assert_not_called() + executor.close.assert_called_once_with() + + +def test_bg_refresh_skips_callbacks_when_stopped_after_refresh(monkeypatch, refresher) -> None: + callback = Mock() + executor = Mock() + is_stopped = Mock(side_effect=[False, True]) + monkeypatch.setattr(refresher._stop_refresh, 'is_set', is_stopped) + monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock()) + monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor)) + refresher.refreshers = {} + + refresher._bg_refresh(make_sqlexecute(), callback, {}) + + callback.assert_not_called() + assert is_stopped.call_count == 2 + + +def test_bg_refresh_propagates_unexpected_refresher_error(monkeypatch, refresher) -> None: + executor = Mock() + + def fail_refresh(completer, active_executor) -> None: + refresher._active_executor = None + raise RuntimeError('refresh failed') + + monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock()) + monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor)) + refresher.refreshers = {'fail': fail_refresh} + + with pytest.raises(RuntimeError, match='refresh failed'): + refresher._bg_refresh(make_sqlexecute(), Mock(), {}) + + executor.close.assert_called_once_with() + + +@pytest.mark.parametrize('stopping', [False, True]) +def test_bg_refresh_only_suppresses_executor_close_error_when_stopping(monkeypatch, refresher, stopping) -> None: + executor = Mock() + executor.close.side_effect = RuntimeError('close failed') + monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock()) + monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor)) + refresher.refreshers = {} + if stopping: + refresher._stop_refresh.set() + + if stopping: + refresher._bg_refresh(make_sqlexecute(), Mock(), {}) + else: + with pytest.raises(RuntimeError, match='close failed'): + refresher._bg_refresh(make_sqlexecute(), Mock(), {}) + + def test_refresher_decorator_registers_function() -> None: refreshers: dict[str, object] = {}