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

andialbrecht · andialbrecht · commit 5bc57bda7607 · 2012-12-22T21:32:44.000+01:00
diff --git a/CHANGES b/CHANGES
@@ -1,6 +1,9 @@
 Development Version
 -------------------
 
+Bug Fixes
+ * Fix parsing error with dollar-quoted procedure bodies (issue83).
+
 Other
  * py3k fixes (by vthriller).
  * py3k fixes in setup.py (by Florian Bauer).
diff --git a/sqlparse/lexer.py b/sqlparse/lexer.py
@@ -177,7 +177,7 @@ class Lexer(object):
             (r'CASE\b', tokens.Keyword),  # extended CASE(foo)
             (r"`(``|[^`])*`", tokens.Name),
             (r"&acute;(&acute;&acute;|[^&acute;])*&acute;", tokens.Name),
-            (r'\$([^\W\d_]\w*)?\$', tokens.Name.Builtin),
+            (r'\$([^\W\d]\w*)?\$', tokens.Name.Builtin),
             (r'\?{1}', tokens.Name.Placeholder),
             (r'[$:?%]\w+', tokens.Name.Placeholder),
             # FIXME(andi): VALUES shouldn't be listed here
diff --git 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
@@ -134,3 +134,26 @@ def _get_identifier(sql):
         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
