String prefixes and dollar-quoted string delimiters lexing in Postgresql. Also closes issue #886.

2 files changed, 59 insertions, 10 deletions

diff --git a/pygments/lexers/sql.py b/pygments/lexers/sql.py
index 05503c3a..948d876a 100644
--- a/pygments/lexers/sql.py
+++ b/pygments/lexers/sql.py
@@ -57,11 +57,14 @@ line_re  = re.compile('.*?\n')
 
 language_re = re.compile(r"\s+LANGUAGE\s+'?(\w+)'?", re.IGNORECASE)
 
+do_re = re.compile(r'\bDO\b', re.IGNORECASE) 
+
 
 def language_callback(lexer, match):
     """Parse the content of a $-string using a lexer
 
-    The lexer is chosen looking for a nearby LANGUAGE.
+    The lexer is chosen looking for a nearby LANGUAGE or assumed as
+    plpgsql if inside a DO statement and no LANGUAGE has been found.
     """
     l = None
     m = language_re.match(lexer.text[match.end():match.end()+100])
@@ -72,15 +75,26 @@ def language_callback(lexer, match):
             lexer.text[max(0, match.start()-100):match.start()]))
         if m:
             l = lexer._get_lexer(m[-1].group(1))
-
+        else:
+            m = list(do_re.finditer(
+                lexer.text[max(0, match.start()-25):match.start()]))
+            if m:
+                l = lexer._get_lexer('plpgsql')
+    
+    # 1 = $, 2 = delimiter, 3 = $
+    yield (match.start(1), String, match.group(1))
+    yield (match.start(2), String.Delimiter, match.group(2))
+    yield (match.start(3), String, match.group(3))
+    # 4 = string contents
     if l:
-        yield (match.start(1), String, match.group(1))
-        for x in l.get_tokens_unprocessed(match.group(2)):
+        for x in l.get_tokens_unprocessed(match.group(4)):
             yield x
-        yield (match.start(3), String, match.group(3))
-
     else:
-        yield (match.start(), String, match.group())
+        yield (match.start(4), String, match.group(4))
+    # 5 = $, 6 = delimiter, 7 = $
+    yield (match.start(5), String, match.group(5))
+    yield (match.start(6), String.Delimiter, match.group(6))
+    yield (match.start(7), String, match.group(7))
 
 
 class PostgresBase(object):
@@ -148,9 +162,10 @@ class PostgresLexer(PostgresBase, RegexLexer):
             (r'\$\d+', Name.Variable),
             (r'([0-9]*\.[0-9]*|[0-9]+)(e[+-]?[0-9]+)?', Number.Float),
             (r'[0-9]+', Number.Integer),
-            (r"(E|U&)?'(''|[^'])*'", String.Single),
-            (r'(U&)?"(""|[^"])*"', String.Name),  # quoted identifier
-            (r'(?s)(\$[^$]*\$)(.*?)(\1)', language_callback),
+            (r"((?:E|U&)?)('(?:''|[^'])*')", bygroups(String.Affix, String.Single)),
+            # quoted identifier
+            (r'((?:U&)?)("(?:""|[^"])*")', bygroups(String.Affix, String.Name)),
+            (r'(?s)(\$)([^$]*)(\$)(.*?)(\$)(\2)(\$)', language_callback),
             (r'[a-z_]\w*', Name),
 
             # psql variable in SQL
diff --git a/tests/examplefiles/postgresql_test.txt b/tests/examplefiles/postgresql_test.txt
index 190d184f..28db5ee3 100644
--- a/tests/examplefiles/postgresql_test.txt
+++ b/tests/examplefiles/postgresql_test.txt
@@ -45,3 +45,37 @@ $$;
 SELECT U&'\0441\043B\043E\043D'
 FROM U&"\0441\043B\043E\043D";
 
+-- Escapes
+SELECT E'1\n2\n3';
+
+-- DO example from postgresql documentation
+/*
+ * PostgreSQL is Copyright © 1996-2016 by the PostgreSQL Global Development Group.
+ * 
+ * Postgres95 is Copyright © 1994-5 by the Regents of the University of California.
+ * 
+ * Permission to use, copy, modify, and distribute this software and its
+ * documentation for any purpose, without fee, and without a written agreement
+ * is hereby granted, provided that the above copyright notice and this paragraph
+ * and the following two paragraphs appear in all copies.
+ * 
+ * IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA BE LIABLE TO ANY PARTY FOR
+ * DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING
+ * LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION,
+ * EVEN IF THE UNIVERSITY OF CALIFORNIA HAS BEEN ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ * 
+ * THE UNIVERSITY OF CALIFORNIA SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING,
+ * BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE. THE SOFTWARE PROVIDED HEREUNDER IS ON AN "AS-IS" BASIS,
+ * AND THE UNIVERSITY OF CALIFORNIA HAS NO OBLIGATIONS TO PROVIDE MAINTENANCE,
+ * SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
+ */
+DO $$DECLARE r record;
+BEGIN
+    FOR r IN SELECT table_schema, table_name FROM information_schema.tables
+             WHERE table_type = 'VIEW' AND table_schema = 'public'
+    LOOP
+        EXECUTE 'GRANT ALL ON ' || quote_ident(r.table_schema) || '.' || quote_ident(r.table_name) || ' TO webuser';
+    END LOOP;
+END$$;