pg_query_normalize: Fix handling of special strings in DefElem (#325)

Fixes #254.

Author: lfittl

src/pg_query_normalize.c

@@ -336,10 +336,20 @@ static void RecordConstLocation(pgssConstLocations *jstate, int location)
 	}
 }
 
+static bool is_string_delimiter(char c)
+{
+	return c == '\'' || c == '$';
+}
+
+static bool is_special_string_start(char c)
+{
+	return c == 'b' || c == 'B' || c == 'x' || c == 'X' || c == 'n' || c == 'N' || c == 'e' || c == 'E';
+}
+
 static void record_defelem_arg_location(pgssConstLocations *jstate, int location)
 {
 	for (int i = location; i < jstate->query_len; i++) {
-		if (jstate->query[i] == '\'' || jstate->query[i] == '$') {
+		if (is_string_delimiter(jstate->query[i]) || (i + 1 < jstate->query_len && is_special_string_start(jstate->query[i]) && is_string_delimiter(jstate->query[i + 1]))) {
 			RecordConstLocation(jstate, i);
 			break;
 		}

test/normalize_tests.c

@@ -40,6 +40,8 @@ const char* tests[] = {
   "ALTER USER MAPPING FOR bob SERVER foo OPTIONS (SET password $1)",
   "MERGE into measurement m USING new_measurement nm ON (m.city_id = nm.city_id and m.logdate=nm.logdate) WHEN MATCHED AND nm.peaktemp IS NULL THEN DELETE WHEN MATCHED THEN UPDATE SET peaktemp = greatest(m.peaktemp, nm.peaktemp), unitsales = m.unitsales + coalesce(nm.unitsales, 0) WHEN NOT MATCHED THEN INSERT (city_id, logdate, peaktemp, unitsales) VALUES (city_id, logdate, peaktemp, unitsales)",
   "MERGE into measurement m USING new_measurement nm ON (m.city_id = nm.city_id and m.logdate=nm.logdate) WHEN MATCHED AND nm.peaktemp IS NULL THEN DELETE WHEN MATCHED THEN UPDATE SET peaktemp = greatest(m.peaktemp, nm.peaktemp), unitsales = m.unitsales + coalesce(nm.unitsales, $1) WHEN NOT MATCHED THEN INSERT (city_id, logdate, peaktemp, unitsales) VALUES (city_id, logdate, peaktemp, unitsales)",
+  "DO language E'plpgsql' 'BEGIN PERFORM 1; END'",
+  "DO language $1 $2",
   // These below are as expected, though questionable if upstream shouldn't be
   // fixed as this could bloat pg_stat_statements
   "DECLARE cursor_b CURSOR FOR SELECT * FROM x WHERE id = 123",

test/normalize_utility_tests.c

@@ -41,6 +41,8 @@ const char* tests[] = {
   "ALTER USER MAPPING FOR bob SERVER foo OPTIONS (SET password $1)",
   "MERGE into measurement m USING new_measurement nm ON (m.city_id = nm.city_id and m.logdate=nm.logdate) WHEN MATCHED AND nm.peaktemp IS NULL THEN DELETE WHEN MATCHED THEN UPDATE SET peaktemp = greatest(m.peaktemp, nm.peaktemp), unitsales = m.unitsales + coalesce(nm.unitsales, 0) WHEN NOT MATCHED THEN INSERT (city_id, logdate, peaktemp, unitsales) VALUES (city_id, logdate, peaktemp, unitsales)",
   "MERGE into measurement m USING new_measurement nm ON (m.city_id = nm.city_id and m.logdate=nm.logdate) WHEN MATCHED AND nm.peaktemp IS NULL THEN DELETE WHEN MATCHED THEN UPDATE SET peaktemp = greatest(m.peaktemp, nm.peaktemp), unitsales = m.unitsales + coalesce(nm.unitsales, 0) WHEN NOT MATCHED THEN INSERT (city_id, logdate, peaktemp, unitsales) VALUES (city_id, logdate, peaktemp, unitsales)",
+  "DO language E'plpgsql' 'BEGIN PERFORM 1; END'",
+  "DO language E'plpgsql' 'BEGIN PERFORM 1; END'",
   // These below are as expected, though questionable if upstream shouldn't be
   // fixed as this could bloat pg_stat_statements
   "DECLARE cursor_b CURSOR FOR SELECT * FROM x WHERE id = 123",