diff --git a/changelog.md b/changelog.md index 5a7777d0..803b6fd1 100644 --- a/changelog.md +++ b/changelog.md @@ -9,6 +9,7 @@ Features Bugfixes --------- * Don't allow saving favorite queries to rewrite comments in `~/.myclirc`. +* Don't allow saving favorite queries to rewrite quoting in `~/.myclirc`. * Don't allow saving named DSNs to rewrite comments in `~/.myclirc`. diff --git a/mycli/config.py b/mycli/config.py index b42f2b92..c90b78bf 100644 --- a/mycli/config.py +++ b/mycli/config.py @@ -14,6 +14,19 @@ logger = logging.getLogger(__name__) +class LimiitedQuotePreservingConfigObj(ConfigObj): + """Useful for saving individual items without modifying the whole file.""" + + def __init__(self, *args, **kwargs): + ConfigObj.__init__(self, *args, **kwargs) + + def _unquote(self, value): + return value + + def _quote(self, value, multiline=True): + return value + + def log(logger: logging.Logger, level: int, message: str) -> None: """Logs message to stderr if logging isn't initialized.""" @@ -23,7 +36,11 @@ def log(logger: logging.Logger, level: int, message: str) -> None: logger.log(level, message) -def read_config_file(f: str | IO[str], list_values: bool = True) -> ConfigObj | None: +def read_config_file( + f: str | IO[str], + list_values: bool = True, + preserve_quotes: bool = False, +) -> ConfigObj | LimiitedQuotePreservingConfigObj | None: """Read a config file. *list_values* set to `True` is the default behavior of ConfigObj. @@ -38,7 +55,10 @@ def read_config_file(f: str | IO[str], list_values: bool = True) -> ConfigObj | f = os.path.expanduser(f) try: - config = ConfigObj(f, interpolation=False, encoding="utf8", list_values=list_values) + if preserve_quotes: + config = LimiitedQuotePreservingConfigObj(f, interpolation=False, encoding="utf8", list_values=False) + else: + config = ConfigObj(f, interpolation=False, encoding="utf8", list_values=list_values) except ConfigObjError as e: log(logger, logging.WARNING, "Unable to parse line {0} of config file '{1}'.".format(e.line_number, f)) log(logger, logging.WARNING, "Using successfully parsed config values.") diff --git a/mycli/packages/special/favoritequeries.py b/mycli/packages/special/favoritequeries.py index 0d7722c5..8b578bfa 100644 --- a/mycli/packages/special/favoritequeries.py +++ b/mycli/packages/special/favoritequeries.py @@ -126,7 +126,7 @@ def _config_for_write(self) -> Any: if self.config_file is None: return self.config - config = read_config_file(self.config_file) + config = read_config_file(self.config_file, preserve_quotes=True) if config is None: raise OSError(f"Unable to read config file '{os.path.expanduser(self.config_file)}'.") return config @@ -144,10 +144,15 @@ def get(self, name) -> str | None: def save(self, name: str, query: str) -> None: config = self._config_for_write() + query = query.rstrip(' \t\n\r;') + if '\n' in query: + manually_quoted_query = f"'''{query}'''" + else: + manually_quoted_query = query 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) + self._set_query(config, name, manually_quoted_query) try: config.write() except Exception: @@ -160,6 +165,7 @@ def save(self, name: str, query: str) -> None: raise if config is not self.config: + # use the unquoted query for the current session self._set_query(self.config, name, query) def delete(self, name: str) -> str: diff --git a/test/pytests/test_favoritequeries.py b/test/pytests/test_favoritequeries.py index d7fb5eea..29a193ce 100644 --- a/test/pytests/test_favoritequeries.py +++ b/test/pytests/test_favoritequeries.py @@ -77,6 +77,28 @@ def test_save_updates_existing_section_and_writes_config() -> None: assert config.write_calls == 1 +def test_save_quotes_multiline_query_for_disk_and_keeps_runtime_query_unquoted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + write_config = DummyConfig() + runtime_config = DummyConfig() + favorites = FavoriteQueries(runtime_config, '/tmp/myclirc') + read_calls: list[tuple[str, bool]] = [] + + def read_config_file(path: str, preserve_quotes: bool = False) -> DummyConfig: + read_calls.append((path, preserve_quotes)) + return write_config + + monkeypatch.setattr(favoritequeries_module, 'read_config_file', read_config_file) + + favorites.save('report', 'select 1;\nselect 2;\n') + + assert read_calls == [('/tmp/myclirc', True)] + assert write_config['favorite_queries']['report'] == "'''select 1;\nselect 2'''" + assert write_config.write_calls == 1 + assert runtime_config['favorite_queries']['report'] == 'select 1;\nselect 2' + + def test_delete_removes_existing_favorite_and_writes_config() -> None: config = DummyConfig({'favorite_queries': {'demo': 'select 1'}}) favorites = FavoriteQueries(config) @@ -222,7 +244,7 @@ def test_save_does_not_update_runtime_config_when_user_config_cannot_be_read( ) -> None: merged_config = DummyConfig({'favorite_queries': {'existing': 'select 1'}}) favorites = FavoriteQueries(merged_config, '~/.myclirc') - monkeypatch.setattr(favoritequeries_module, 'read_config_file', lambda _path: None) + monkeypatch.setattr(favoritequeries_module, 'read_config_file', lambda _path, **_kwargs: None) with pytest.raises(OSError, match=r"Unable to read config file '.*/\.myclirc'\."): favorites.save('new', 'select 2')