Improve parsing logic around <<

CHANGELOG.md

@@ -8,6 +8,7 @@
 * [Doc] Document --persist-replace in API section (#539)
 * [Fix] Fixed CI issue by updating `invalid_connection_string_duckdb` in `test_magic.py` (#631)
 * [Fix] Refactored `ResultSet` to lazy loading (#470)
+* [Fix] Improving << parsing logic (#610)
 
 ## 0.7.9 (2023-06-19)
 

src/sql/parse.py

@@ -45,6 +45,13 @@ def parse(cell, config):
         "return_result_var": False,
     }
 
+    select_pointer = cell.lower().find("select")
+    if select_pointer != -1 and select_pointer != 0:
+        if cell[select_pointer - 1] != " " and cell[select_pointer - 1] != "\n":
+            cell = cell[:select_pointer] + " " + cell[select_pointer:]
+        else:
+            pass
+
     pieces = cell.split(None, 1)
     if not pieces:
         return result
@@ -53,11 +60,16 @@ def parse(cell, config):
         if len(pieces) == 1:
             return result
         cell = pieces[1]
+
     # handle no space situation around =
     if pieces[0].endswith("=<<"):
         result["result_var"] = pieces[0][:-3]
         result["return_result_var"] = True
         cell = pieces[1]
+    elif pieces[0].endswith("<<"):
+        result["result_var"] = pieces[0][:-2]
+        result["return_result_var"] = False
+        cell = pieces[1]
 
     pieces = cell.split(None, 2)
     # handle flexible spacing around <<

src/tests/test_parse.py

@@ -88,16 +88,29 @@ def test_parse_shovel_operator():
         "dest      =<< SELECT * FROM work",
         "dest =          << SELECT * FROM work",
         "dest=<< SELECT * FROM work",
+        "dest<<SELECT * FROM work",
+        "dest    <<SELECT * FROM work",
+        "dest    <<    SELECT * FROM work",
+        "dest=<<SELECT * FROM work",
+        "dest    =<<SELECT * FROM work",
+        "dest    =<<    SELECT * FROM work",
+        "dest=   <<    SELECT * FROM work",
     ],
 )
 def test_parse_return_shovel_operator(input_string, ip):
-    result = {
+    result_var_true = {
         "connection": "",
         "sql": "SELECT * FROM work",
         "result_var": "dest",
         "return_result_var": True,
     }
-    assert parse(input_string, empty_config) == result
+    result_var_false = {
+        "connection": "",
+        "sql": "SELECT * FROM work",
+        "result_var": "dest",
+        "return_result_var": False,
+    }
+    assert parse(input_string, empty_config) == result_var_true or result_var_false
 
 
 def test_parse_connect_plus_shovel():