Fix parsing error with dollar-quoted procedure bodies (fixes issue83).

2 files changed, 44 insertions, 0 deletions

diff --git a/tests/test_parse.py b/tests/test_parse.py
index 59b8e72..6c9cfa9 100644
--- a/tests/test_parse.py
+++ b/tests/test_parse.py
@@ -113,3 +113,24 @@ def test_quoted_identifier():
     assert isinstance(t[2], sqlparse.sql.Identifier)
     assert t[2].get_name() == 'z'
     assert t[2].get_real_name() == 'y'
+
+
+def test_psql_quotation_marks():  # issue83
+    # regression: make sure plain $$ work
+    t = sqlparse.split("""
+    CREATE OR REPLACE FUNCTION testfunc1(integer) RETURNS integer AS $$
+          ....
+    $$ LANGUAGE plpgsql;
+    CREATE OR REPLACE FUNCTION testfunc2(integer) RETURNS integer AS $$
+          ....
+    $$ LANGUAGE plpgsql;""")
+    assert len(t) == 2
+    # make sure $SOMETHING$ works too
+    t = sqlparse.split("""
+    CREATE OR REPLACE FUNCTION testfunc1(integer) RETURNS integer AS $PROC_1$
+          ....
+    $PROC_1$ LANGUAGE plpgsql;
+    CREATE OR REPLACE FUNCTION testfunc2(integer) RETURNS integer AS $PROC_2$
+          ....
+    $PROC_2$ LANGUAGE plpgsql;""")
+    assert len(t) == 2
diff --git a/tests/test_regressions.py b/tests/test_regressions.py
index e1c2c89..ba156f0 100644
--- a/tests/test_regressions.py
+++ b/tests/test_regressions.py
@@ -134,3 +134,26 @@ def test_issue78():
         for func_name, result in results:
             func = getattr(i, func_name)
             assert func() == result
+
+
+def test_issue83():
+    sql = """
+CREATE OR REPLACE FUNCTION func_a(text)
+  RETURNS boolean  LANGUAGE plpgsql STRICT IMMUTABLE AS
+$_$
+BEGIN
+ ...
+END;
+$_$;
+
+CREATE OR REPLACE FUNCTION func_b(text)
+  RETURNS boolean  LANGUAGE plpgsql STRICT IMMUTABLE AS
+$_$
+BEGIN
+ ...
+END;
+$_$;
+
+ALTER TABLE..... ;"""
+    t = sqlparse.split(sql)
+    assert len(t) == 3