Avoid replacing newlines with spaces early in the query

The parser split off the first three elements of the cell, to check
if they have a connection string or push the result to a variable.
But if these elements weren't found, the components were glued back
together with spaces.  If one of these components was a line comment,
it ends up commenting out more than it should.

This solution tries splitting off bits to see if they match the
expected form.  But if they don't, it goes back to using the pre-
split string, thus preserving the original whitespace.  This should
allow the connection string and redirect to appear after any number
of newline characters, as long as they are still the first non-
whitespace elements in the cell.

src/sql/parse.py

@@ -43,16 +43,23 @@ def parse(cell, config):
 
     result = {"connection": "", "sql": "", "result_var": None}
 
-    pieces = cell.split(None, 3)
+    pieces = cell.split(None, 1)
     if not pieces:
         return result
     result["connection"] = _connection_string(pieces[0], config)
     if result["connection"]:
-        pieces.pop(0)
+        if len(pieces) == 1:
+            return result
+        cell = pieces[1]
+
+    pieces = cell.split(None, 2)
     if len(pieces) > 1 and pieces[1] == "<<":
-        result["result_var"] = pieces.pop(0)
-        pieces.pop(0)  # discard << operator
-    result["sql"] = (" ".join(pieces)).strip()
+        result["result_var"] = pieces[0]
+        if len(pieces) == 2:
+            return result
+        cell = pieces[2]
+
+    result["sql"] = cell
     return result
 
 

src/tests/test_parse.py

@@ -91,6 +91,22 @@ def test_parse_connect_plus_shovel():
     }
 
 
+def test_parse_early_newlines():
+    assert parse("--comment\nSELECT *\n--comment\nFROM work", empty_config) == {
+        "connection": "",
+        "sql": "--comment\nSELECT *\n--comment\nFROM work",
+        "result_var": None
+    }
+
+
+def test_parse_connect_shovel_over_newlines():
+    assert parse("\nsqlite://\ndest\n<<\nSELECT *\nFROM work", empty_config) == {
+        "connection": "sqlite://",
+        "sql": "SELECT *\nFROM work",
+        "result_var": "dest"
+    }
+
+
 class DummyConfig:
     dsn_filename = Path("src/tests/test_dsn_config.ini")