diff --git a/changelog.md b/changelog.md index f5c0741f..31a1757a 100644 --- a/changelog.md +++ b/changelog.md @@ -6,6 +6,11 @@ Features * Let favorite queries use Jinja2 templates, allowing optional arguments. +Bugfixes +--------- +* Don't allow saving favorite queries to rewrite comments in `~/.myclirc`. + + Documentation --------- * Document shell completions in `README.md`. diff --git a/mycli/client.py b/mycli/client.py index c7a4034d..722639a7 100644 --- a/mycli/client.py +++ b/mycli/client.py @@ -115,7 +115,7 @@ def __init__( self.beep_after_seconds = float(c["main"]["beep_after_seconds"] or 0) self.default_keepalive_ticks = c['connection'].as_int('default_keepalive_ticks') - FavoriteQueries.instance = FavoriteQueries.from_config(self.config) + FavoriteQueries.instance = FavoriteQueries.from_config(self.config, myclirc) DsnAliases.instance = DsnAliases.from_config(self.config, self) self.dsn_alias: str | None = None diff --git a/mycli/packages/special/favoritequeries.py b/mycli/packages/special/favoritequeries.py index 6ff584ba..0d7722c5 100644 --- a/mycli/packages/special/favoritequeries.py +++ b/mycli/packages/special/favoritequeries.py @@ -1,10 +1,16 @@ from __future__ import annotations +import os import re +from typing import Any from jinja2 import meta, nodes from jinja2.sandbox import SandboxedEnvironment +from mycli.config import read_config_file + +MISSING = object() + favorite_query_template_environment = SandboxedEnvironment(autoescape=False) favorite_query_variable_pattern = re.compile(r'^[A-Za-z_][A-Za-z0-9_-]*$') @@ -108,12 +114,27 @@ class FavoriteQueries: # Class-level variable, for convenience to use as a singleton. instance: FavoriteQueries - def __init__(self, config) -> None: + def __init__(self, config: Any, config_file: str | None = None) -> None: self.config = config + self.config_file = config_file @classmethod - def from_config(cls, config): - return FavoriteQueries(config) + def from_config(cls, config: Any, config_file: str | None = None) -> FavoriteQueries: + return FavoriteQueries(config, config_file) + + def _config_for_write(self) -> Any: + if self.config_file is None: + return self.config + + config = read_config_file(self.config_file) + if config is None: + raise OSError(f"Unable to read config file '{os.path.expanduser(self.config_file)}'.") + return config + + def _set_query(self, config: Any, name: str, query: str) -> None: + if self.section_name not in config: + config[self.section_name] = {} + config[self.section_name][name] = query def list(self) -> list[str | None]: return list(self.config.get(self.section_name, {})) @@ -122,16 +143,41 @@ def get(self, name) -> str | None: return self.config.get(self.section_name, {}).get(name, None) def save(self, name: str, query: str) -> None: - self.config.encoding = "utf-8" - if self.section_name not in self.config: - self.config[self.section_name] = {} - self.config[self.section_name][name] = query - self.config.write() + config = self._config_for_write() + config.encoding = "utf-8" + section_existed = self.section_name in config + previous_query = config.get(self.section_name, {}).get(name, MISSING) + self._set_query(config, name, query) + try: + config.write() + except Exception: + if previous_query is MISSING: + del config[self.section_name][name] + if not section_existed: + del config[self.section_name] + else: + config[self.section_name][name] = previous_query + raise + + if config is not self.config: + self._set_query(self.config, name, query) def delete(self, name: str) -> str: try: - del self.config[self.section_name][name] + self.config[self.section_name][name] except KeyError: return f'{name}: Not Found.' - self.config.write() + + config = self._config_for_write() + if name in config.get(self.section_name, {}): + query = config[self.section_name][name] + del config[self.section_name][name] + try: + config.write() + except Exception: + config[self.section_name][name] = query + raise + + if config is not self.config: + del self.config[self.section_name][name] return f'{name}: Deleted.' diff --git a/test/pytests/test_client.py b/test/pytests/test_client.py index 471c34de..752521de 100644 --- a/test/pytests/test_client.py +++ b/test/pytests/test_client.py @@ -9,6 +9,7 @@ import mycli.client as client_module from mycli.client import MyCli +from mycli.packages.special.favoritequeries import FavoriteQueries def write_myclirc(tmp_path: Path, content: str) -> str: @@ -174,6 +175,16 @@ def test_init_uses_existing_xdg_config_when_myclirc_is_not_given(monkeypatch: py assert cli.config.filename == str(xdg_config) +def test_init_configures_favorite_queries_with_user_config_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + patch_constructor_side_effects(monkeypatch) + myclirc = write_myclirc(tmp_path, '') + + cli = MyCli(myclirc=myclirc) + + assert FavoriteQueries.instance.config is cli.config + assert FavoriteQueries.instance.config_file == myclirc + + def test_init_uses_default_myclirc_when_xdg_config_is_missing(monkeypatch: pytest.MonkeyPatch) -> None: patch_constructor_side_effects(monkeypatch) config_file_args: list[list[str | Any]] = [] diff --git a/test/pytests/test_favoritequeries.py b/test/pytests/test_favoritequeries.py index c3c3aee7..d7fb5eea 100644 --- a/test/pytests/test_favoritequeries.py +++ b/test/pytests/test_favoritequeries.py @@ -1,5 +1,9 @@ from collections.abc import Mapping +from pathlib import Path +import pytest + +import mycli.packages.special.favoritequeries as favoritequeries_module from mycli.packages.special.favoritequeries import FavoriteQueries @@ -13,13 +17,19 @@ def write(self) -> None: self.write_calls += 1 +class FailingConfig(DummyConfig): + def write(self) -> None: + raise OSError('write failed') + + def test_from_config_returns_instance_with_same_config() -> None: config = DummyConfig() - favorites = FavoriteQueries.from_config(config) + favorites = FavoriteQueries.from_config(config, '/tmp/myclirc') assert isinstance(favorites, FavoriteQueries) assert favorites.config is config + assert favorites.config_file == '/tmp/myclirc' def test_list_and_get_use_favorite_queries_section() -> None: @@ -98,3 +108,154 @@ def test_delete_returns_not_found_when_section_is_missing() -> None: assert result == 'missing: Not Found.' assert config == {} assert config.write_calls == 0 + + +def test_save_preserves_user_config_comments_and_excludes_merged_values(tmp_path: Path) -> None: + config_file = tmp_path / 'myclirc' + config_file.write_text( + """# User introduction. +[main] +prompt = custom # Inline comment. + +[favorite_queries] +# Existing favorite. +existing = select 1 +# User footer. +""", + encoding='utf-8', + ) + merged_config = DummyConfig({ + 'main': {'prompt': 'custom', 'package_default': 'do not write'}, + 'favorite_queries': {'existing': 'select 1'}, + }) + favorites = FavoriteQueries(merged_config, str(config_file)) + + favorites.save('new', 'select 2') + + assert ( + config_file.read_text(encoding='utf-8') + == """# User introduction. +[main] +prompt = custom# Inline comment. + +[favorite_queries] +# Existing favorite. +existing = select 1 +new = select 2 +# User footer. +""" + ) + assert merged_config['favorite_queries']['new'] == 'select 2' + assert 'package_default' not in config_file.read_text(encoding='utf-8') + + +def test_save_overwrites_favorite_without_removing_its_comment(tmp_path: Path) -> None: + config_file = tmp_path / 'myclirc' + config_file.write_text( + """[favorite_queries] +# Keep this explanation. +report = select 1 +""", + encoding='utf-8', + ) + merged_config = DummyConfig({'favorite_queries': {'report': 'select 1'}}) + + FavoriteQueries(merged_config, str(config_file)).save('report', 'select 2') + + assert ( + config_file.read_text(encoding='utf-8') + == """[favorite_queries] +# Keep this explanation. +report = select 2 +""" + ) + assert merged_config['favorite_queries']['report'] == 'select 2' + + +def test_delete_preserves_unrelated_user_config_comments(tmp_path: Path) -> None: + config_file = tmp_path / 'myclirc' + config_file.write_text( + """# User introduction. +[favorite_queries] +# Removed with the favorite. +remove = select 1 +# Keep this explanation. +keep = select 2 +# User footer. +""", + encoding='utf-8', + ) + merged_config = DummyConfig({'favorite_queries': {'remove': 'select 1', 'keep': 'select 2'}}) + favorites = FavoriteQueries(merged_config, str(config_file)) + + result = favorites.delete('remove') + + assert result == 'remove: Deleted.' + assert ( + config_file.read_text(encoding='utf-8') + == """# User introduction. +[favorite_queries] +# Keep this explanation. +keep = select 2 +# User footer. +""" + ) + assert merged_config['favorite_queries'] == {'keep': 'select 2'} + + +def test_delete_effective_system_favorite_does_not_rewrite_user_config(tmp_path: Path) -> None: + config_file = tmp_path / 'myclirc' + original = '# User commentary.\n[main]\nprompt = custom\n' + config_file.write_text(original, encoding='utf-8') + merged_config = DummyConfig({'favorite_queries': {'system': 'select 1'}}) + favorites = FavoriteQueries(merged_config, str(config_file)) + + result = favorites.delete('system') + + assert result == 'system: Deleted.' + assert config_file.read_text(encoding='utf-8') == original + assert merged_config['favorite_queries'] == {} + + +def test_save_does_not_update_runtime_config_when_user_config_cannot_be_read( + monkeypatch: pytest.MonkeyPatch, +) -> None: + merged_config = DummyConfig({'favorite_queries': {'existing': 'select 1'}}) + favorites = FavoriteQueries(merged_config, '~/.myclirc') + monkeypatch.setattr(favoritequeries_module, 'read_config_file', lambda _path: None) + + with pytest.raises(OSError, match=r"Unable to read config file '.*/\.myclirc'\."): + favorites.save('new', 'select 2') + + assert merged_config['favorite_queries'] == {'existing': 'select 1'} + + +@pytest.mark.parametrize('initial', ({}, {'favorite_queries': {'existing': 'select 1'}})) +def test_save_restores_runtime_config_after_write_failure(initial: dict[str, object]) -> None: + config = FailingConfig(initial) + favorites = FavoriteQueries(config) + + with pytest.raises(OSError, match='write failed'): + favorites.save('new', 'select 2') + + assert config == initial + + +def test_save_restores_overwritten_runtime_query_after_write_failure() -> None: + config = FailingConfig({'favorite_queries': {'existing': 'select 1'}}) + favorites = FavoriteQueries(config) + + with pytest.raises(OSError, match='write failed'): + favorites.save('existing', 'select 2') + + assert config['favorite_queries']['existing'] == 'select 1' + + +def test_delete_restores_runtime_query_after_write_failure() -> None: + config = FailingConfig({'favorite_queries': {'existing': 'select 1'}}) + favorites = FavoriteQueries(config) + + with pytest.raises(OSError, match='write failed'): + favorites.delete('existing') + + assert config['favorite_queries'] == {'existing': 'select 1'}