Improve parsing logic around <<

CHANGELOG.md

@@ -14,6 +14,7 @@
 * [Fix] Removed `WITH` when a snippet does not have a dependency (#657)
 * [Fix] Used display module when generating CTE (#649)
 * [Fix] Adding `--with` back because of issues with sqlglot query parser (#684)
+* [Fix] Improving << parsing logic (#610)
 
 ## 0.7.9 (2023-06-19)
 

src/sql/parse.py

@@ -37,7 +37,6 @@ def parse(cell, config):
     We're grandfathering the
     connection string and `<<` operator in.
     """
-
     result = {
         "connection": "",
         "sql": "",
@@ -53,37 +52,21 @@ 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]
 
-    pieces = cell.split(None, 2)
-    # handle flexible spacing around <<
-    if len(pieces) > 1 and pieces[1] == "<<":
-        if pieces[0].endswith("="):
-            result["result_var"] = pieces[0][:-1]
-            result["return_result_var"] = True
-        else:
-            result["result_var"] = pieces[0]
+    pointer = cell.find("<<")
+    if pointer != -1:
+        left = cell[:pointer].replace(" ", "").replace("\n", "")
+        right = cell[pointer + 2 :].strip(" ")
 
-        if len(pieces) == 2:
-            return result
-        cell = pieces[2]
-    # handle flexible spacing around =<<
-    elif len(pieces) > 1 and (
-        (pieces[1] == "=<<") or (pieces[1] == "=" and pieces[2].startswith("<<"))
-    ):
-        result["result_var"] = pieces[0]
-        result["return_result_var"] = True
-        if pieces[1] == "=<<":
-            cell = pieces[2]
+        if "=" in left:
+            result["result_var"] = left[:-1]
+            result["return_result_var"] = True
         else:
-            pieces = cell.split(None, 3)
-            cell = pieces[3]
+            result["result_var"] = left
 
-    result["sql"] = cell
+        result["sql"] = right
+    else:
+        result["sql"] = cell
     return result
 
 

src/tests/test_parse.py

@@ -88,16 +88,41 @@ 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",
     ],
 )
-def test_parse_return_shovel_operator(input_string, ip):
-    result = {
+def test_parse_return_shovel_operator_with_equal(input_string, ip):
+    result_var = {
         "connection": "",
         "sql": "SELECT * FROM work",
         "result_var": "dest",
         "return_result_var": True,
     }
-    assert parse(input_string, empty_config) == result
+    assert parse(input_string, empty_config) == result_var
+
+
+@pytest.mark.parametrize(
+    "input_string",
+    [
+        "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_without_equal(input_string, ip):
+    result_var = {
+        "connection": "",
+        "sql": "SELECT * FROM work",
+        "result_var": "dest",
+        "return_result_var": False,
+    }
+    assert parse(input_string, empty_config) == result_var
 
 
 def test_parse_connect_plus_shovel():
@@ -121,7 +146,7 @@ def test_parse_early_newlines():
 def test_parse_connect_shovel_over_newlines():
     assert parse("\nsqlite://\ndest\n<<\nSELECT *\nFROM work", empty_config) == {
         "connection": "sqlite://",
-        "sql": "SELECT *\nFROM work",
+        "sql": "\nSELECT *\nFROM work",
         "result_var": "dest",
         "return_result_var": False,
     }