[Fix 954]:

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 0.10.8dev
 
+* [Fix] Fixed bug where the `%sql` line magic would not parse properly due to expectations from shlex(#954)
+
 ## 0.10.7 (2023-12-23)
 
 * [Feature] Add Spark Connection as a dialect for Jupysql ([#965](https://github.com/ploomber/jupysql/issues/965)) (by [@gilandose](https://github.com/gilandose))

src/sql/parse.py

@@ -228,11 +228,14 @@ def without_sql_comment(parser, line):
     """
 
     args = _option_strings_from_parser(parser)
+    pattern = re.compile(r'([\'"])(.*?)\1')
+    line = pattern.sub(lambda match: match.group().replace(" ", "@@SPACE@@"), line)
     result = itertools.takewhile(
-        lambda word: (not word.startswith("--")) or (word in args),
-        shlex.split(line, posix=False),
+        lambda word: (not word.startswith("--")) or (word in args), line.split()
     )
-    return " ".join(result)
+    result = " ".join(result)
+    result = result.replace("@@SPACE@@", " ")
+    return result
 
 
 def split_args_and_sql(line):
@@ -274,20 +277,26 @@ def split_args_and_sql(line):
     # If any SQL commands are found in the line, we split the line into args and sql.
     #   Note: lines without SQL commands will not be split
     #       ex. %sql duckdb:// or %sqlplot boxplot --table data.csv
-    if not any(cmd in line_no_filenames for cmd in SQL_COMMANDS):
-        return arg_line, sql_line
-
-    # Identify beginning of sql query using keywords
-    split_idx = -1
-    for token in line.split():
-        if token.lower() in SQL_COMMANDS:
-            # Found index at which to split line
-            split_idx = line.find(token)
-            break
-
-    # Split line into args and sql, beginning at sql keyword
-    if split_idx != -1:
-        arg_line, sql_line = line[:split_idx], line[split_idx:]
+    if "<<" in line:
+        [before_assign, after_assign] = line.split("<<")
+        result_var = before_assign.split()[-1]
+        arg_line = " ".join(before_assign.split()[:-1])
+        sql_line = result_var + " << " + after_assign
+
+    elif any(cmd in line_no_filenames for cmd in SQL_COMMANDS) or any(
+        cmd.upper() in line_no_filenames for cmd in SQL_COMMANDS
+    ):
+        # Identify beginning of sql query using keywords
+        split_idx = -1
+        for token in line.split():
+            if token.lower() in SQL_COMMANDS:
+                # Found index at which to split line
+                split_idx = line.find(token)
+                break
+
+        # Split line into args and sql, beginning at sql keyword
+        if split_idx != -1:
+            arg_line, sql_line = line[:split_idx], line[split_idx:]
 
     return arg_line, sql_line
 
@@ -301,14 +310,13 @@ def magic_args(magic_execute, line, cmd_from, allowed_duplicates=None):
     arg_line, sql_line = split_args_and_sql(line)
 
     args = shlex.split(arg_line, posix=False)
-
     if len(args) > 1:
         check_duplicate_arguments(magic_execute, cmd_from, args, allowed_duplicates)
 
     parsed = magic_execute.parser.parse_args(args)
 
     if sql_line:
-        parsed.line = shlex.split(sql_line, posix=False)
+        parsed.line = sql_line.split()
 
     return parsed
 

src/tests/test_column_guesser.py

@@ -12,7 +12,8 @@ def __init__(self, connectstr):
         self.connectstr = connectstr
 
     def query(self, txt):
-        return ip.run_line_magic("sql", "%s %s" % (self.connectstr, txt))
+        ip.run_line_magic("sql", self.connectstr)
+        return ip.run_line_magic("sql", "%s" % txt)
 
 
 sql_env = SqlEnv("sqlite://")

src/tests/test_parse.py

@@ -263,6 +263,12 @@ def test_without_sql_comment_dashes_in_string():
     assert without_sql_comment(parser=parser_stub, line=line) == expected
 
 
+def test_without_sql_comment_dashes_in_string_with_spaces():
+    line = "SELECT ' --very --confusing ' FROM author -- uff da"
+    expected = "SELECT ' --very --confusing ' FROM author"
+    assert without_sql_comment(parser=parser_stub, line=line) == expected
+
+
 def test_without_sql_comment_with_arg_and_leading_comment():
     line = "--file moo.txt --persist --comment, not arg"
     expected = "--file moo.txt --persist"
@@ -897,6 +903,11 @@ def test_escape_string_slicing_notation(query, expected_escaped, expected_found)
             "-p --save snippet -N ",
             "insert into authors values('[100]'::json->0)",
         ),
+        (
+            "--save snippet SELECT TRIM(' padded ')",
+            "--save snippet ",
+            "SELECT TRIM(' padded ')",
+        ),
     ],
     ids=[
         "no-query",
@@ -910,6 +921,7 @@ def test_escape_string_slicing_notation(query, expected_escaped, expected_found)
         "update",
         "delete",
         "insert",
+        "select-uppercase",
     ],
 )
 def test_split_args_and_sql(line, expected_args, expected_sql):