Merge pull request #32 from mercedes-benz/fix-handling-of-magic-sql-in-sql-cells

Fix handling of magic sql in sql cells

Author: bsamseth

blackbricks/blackbricks.py

@@ -95,14 +95,13 @@ def _format_sql_cell(
     if NOFMT in cell.strip().splitlines()[0]:
         return title_line + cell
 
-    magics = []
     sql_lines = []
-    for line in cell.strip().splitlines():
-        if line.strip().startswith("# MAGIC %sql"):
+    for line in (line.strip() for line in cell.strip().splitlines()):
+        if line == "# MAGIC %sql":
             continue
-        words = line.split()
-        magic, sql = words[:2], words[2:]
-        magics.append(magic)
+        elif line.startswith("# MAGIC %sql"):
+            line = line.replace(" %sql", "", 1)
+        sql = line.split()[2:]  # Remove "# MAGIC".
         sql_lines.append(" ".join(sql).strip())
 
     return (
@@ -111,7 +110,9 @@ def _format_sql_cell(
         + "\n".join(
             f"# MAGIC {sql}"
             for sql in sqlparse.format(
-                "\n".join(sql_lines), reindent=True, keyword_case=sql_keyword_case
+                "\n".join(sql_lines).strip(),
+                reindent=True,
+                keyword_case=sql_keyword_case,
             ).splitlines()
         )
     )

test_notebooks/test.py

@@ -71,4 +71,11 @@ def test_func(input_param):
 # MAGIC           COLLECT_SET(bar.fizz)[0],
 # MAGIC           FIRST(bar.baz)
 # MAGIC    FROM dsa.asd bar
-# MAGIC    GROUP BY bar.fizzbuzz);
+# MAGIC    GROUP BY bar.fizzbuzz);
+
+# COMMAND ----------
+
+# MAGIC %sql
+# MAGIC SELECT id
+# MAGIC FROM test_table
+# MAGIC LIMIT 1