Merge pull request #2038 from dbcli/RW/spaces-in-shell-csv-redirect-target

Let the file target of `$>` redirection be quoted

Author: rolandwalker

changelog.md

@@ -4,6 +4,7 @@ Upcoming (TBD)
 Features
 ---------
 * Subcommand completions for the `/dsn` command.
+* Allow file target of `$>` redirection to be quoted.
 
 
 Bug Fixes

mycli/TIPS

@@ -296,8 +296,8 @@ customize password sources/precedence with "password_sources" in ~/.myclirc!
 
 redirect query output to a shell command with "$| <command>"!
 
-redirect query output to a file with "$> <filename>"!
+redirect query output to a CSV file with "$> <filename>"!
 
-append query output to a file with "$>> <filename>"!
+append query output to a CSV file with "$>> <filename>"!
 
 run a command after shell redirects with "post_redirect_command" in ~/.myclirc!

mycli/packages/hybrid_redirection.py

@@ -1,5 +1,6 @@
 import functools
 import logging
+import shlex
 
 import sqlglot
 
@@ -113,9 +114,6 @@ def invalid_shell_part(
     file_part: str | None,
     command_part: str | None,
 ) -> bool:
-    if file_part and ' ' in file_part:
-        return True
-
     if file_part and '>' in file_part:
         return True
 
@@ -125,6 +123,19 @@ def invalid_shell_part(
     return False
 
 
+def parse_redirect_filename(file_part: str | None) -> str | None:
+    """Return one unquoted redirect filename, or None for an invalid operand."""
+    if file_part is None or not file_part:
+        return None
+    if file_part[0] not in ('\'', '"'):
+        return file_part if not any(character.isspace() for character in file_part) else None
+    try:
+        parts = shlex.split(file_part)
+    except ValueError:
+        return None
+    return parts[0] if len(parts) == 1 else None
+
+
 # todo there are still corner cases combining custom delimiters, caching, and redirection
 @functools.lru_cache(maxsize=1)
 def get_redirect_components(command: str) -> tuple[str | None, str | None, str | None, str | None]:
@@ -174,7 +185,9 @@ def get_redirect_components(command: str) -> tuple[str | None, str | None, str |
     )
 
     if file_part_tokens:
-        file_part = assemble_tokens(file_part_tokens)
+        file_part = parse_redirect_filename(assemble_tokens(file_part_tokens))
+        if file_part is None:
+            return None, None, None, None
     else:
         file_part = None
 

test/pytests/test_hybrid_redirection.py

@@ -30,6 +30,17 @@ def test_find_token_indices_tracks_true_dollars_and_operators() -> None:
     }
 
 
+def test_find_token_indices_ignores_non_redirect_dollars() -> None:
+    tokens = tokenize('select 1 $ $> output.txt')
+
+    assert hybrid_redirection.find_token_indices(tokens) == {
+        'raw_dollar': [2, 3],
+        'true_dollar': [3],
+        'angle_bracket': [4],
+        'pipe': [],
+    }
+
+
 # todo there are still corner cases combining custom delimiters and redirection
 def test_find_sql_part_handles_valid_parse_custom_delimiter_and_invalid_sql(reset_hybrid_redirection) -> None:
     hybrid_redirection.delimiter_command._delimiter = '$$'
@@ -74,10 +85,27 @@ def test_assemble_tokens_quotes_identifier_and_string() -> None:
     assert hybrid_redirection.assemble_tokens(string_tokens) == "'printf'"
 
 
+@pytest.mark.parametrize(
+    ('file_part', 'expected'),
+    [
+        (None, None),
+        ('out.txt', 'out.txt'),
+        (r'C:\Users\alice\output.csv', r'C:\Users\alice\output.csv'),
+        ("'two words.txt'", 'two words.txt'),
+        ('"two words.txt"', 'two words.txt'),
+        ('two words.txt', None),
+        ("'missing quote", None),
+        ('', None),
+    ],
+)
+def test_parse_redirect_filename(file_part: str | None, expected: str | None) -> None:
+    assert hybrid_redirection.parse_redirect_filename(file_part) == expected
+
+
 @pytest.mark.parametrize(
     ('file_part', 'command_part', 'expected'),
     [
-        ('two words.txt', None, True),
+        ('two words.txt', None, False),
         ('bad>file.txt', None, True),
         (None, None, True),
         ('out.txt', None, False),
@@ -101,6 +129,24 @@ def test_get_redirect_components_valid_paths_and_logging() -> None:
         '>',
         'out.txt',
     )
+    assert hybrid_redirection.get_redirect_components('select 1 $> "two words.txt"') == (
+        'select 1',
+        None,
+        '>',
+        'two words.txt',
+    )
+    assert hybrid_redirection.get_redirect_components("select 1 $>> 'two words.txt'") == (
+        'select 1',
+        None,
+        '>>',
+        'two words.txt',
+    )
+    assert hybrid_redirection.get_redirect_components(r'select 1 $> C:\Users\alice\output.csv') == (
+        'select 1',
+        None,
+        '>',
+        r'C:\Users\alice\output.csv',
+    )
 
 
 def test_get_redirect_components_returns_none_on_token_error(monkeypatch) -> None:
@@ -116,7 +162,9 @@ def test_get_redirect_components_rejects_invalid_forms() -> None:
     assert hybrid_redirection.get_redirect_components('select 1 $> out.txt $> other.txt') == (None, None, None, None)
     assert hybrid_redirection.get_redirect_components('select 1 $> out.txt $| cat') == (None, None, None, None)
     assert hybrid_redirection.get_redirect_components('select from $> out.txt') == (None, None, None, None)
-    assert hybrid_redirection.get_redirect_components('select 1 $> "two words.txt"') == (None, None, None, None)
+    assert hybrid_redirection.get_redirect_components('select 1 $> two words.txt') == (None, None, None, None)
+    assert hybrid_redirection.get_redirect_components("select 1 $> 'missing quote") == (None, None, None, None)
+    assert hybrid_redirection.get_redirect_components('select 1 $> "bad>file.txt"') == (None, None, None, None)
 
 
 def test_get_redirect_components_rejects_multiple_pipes_on_windows(monkeypatch) -> None: