PL/pgSQL: Add support for variables declared with a collation

This requires the data type parsing to correctly set the collation field, which for now we do when the type is defined as "text". In passing we also adjust the hardcoded collection in get_collation_oid to use DEFAULT_COLLATION_OID (100), which seems better than using -1.
geldata · Oct 9, 2024 · 115885c · 115885c
1 parent 975c7a9
commit 115885c
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 7 deletions.
diff --git a/scripts/extract_source.rb b/scripts/extract_source.rb
@@ -573,19 +573,30 @@ def write_out
 }
 ))
 runner.mock('parse_datatype', %(
+#include "catalog/pg_collation_d.h"
 static PLpgSQL_type * parse_datatype(const char *string, int location) {
 	PLpgSQL_type *typ;
+
+	/* Ignore trailing spaces */
+	size_t len = strlen(string);
+	while (len > 0 && scanner_isspace(string[len - 1])) --len;
+
 	typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type));
 	typ->typname = pstrdup(string);
-	typ->ttype = pg_strcasecmp(string, "RECORD") == 0 ? PLPGSQL_TTYPE_REC : PLPGSQL_TTYPE_SCALAR;
-	if (pg_strcasecmp(string, "REFCURSOR") == 0 || pg_strcasecmp(string, "CURSOR") == 0)
+	typ->ttype = pg_strncasecmp(string, "RECORD", len) == 0 ? PLPGSQL_TTYPE_REC : PLPGSQL_TTYPE_SCALAR;
+	if (pg_strncasecmp(string, "REFCURSOR", len) == 0 || pg_strncasecmp(string, "CURSOR", len) == 0)
 	{
 		typ->typoid = REFCURSOROID;
 	}
+	else if (pg_strncasecmp(string, "TEXT", len) == 0)
+	{
+		typ->typoid = TEXTOID;
+		typ->collation = DEFAULT_COLLATION_OID;
+	}
 	return typ;
 }
 ))
-runner.mock('get_collation_oid', 'Oid get_collation_oid(List *name, bool missing_ok) { return -1; }')
+runner.mock('get_collation_oid', 'Oid get_collation_oid(List *name, bool missing_ok) { return DEFAULT_COLLATION_OID; }')
 runner.mock('plpgsql_parse_wordtype', 'PLpgSQL_type * plpgsql_parse_wordtype(char *ident) { return NULL; }')
 runner.mock('plpgsql_parse_wordrowtype', 'PLpgSQL_type * plpgsql_parse_wordrowtype(char *ident) { return NULL; }')
 runner.mock('plpgsql_parse_cwordtype', 'PLpgSQL_type * plpgsql_parse_cwordtype(List *idents) { return NULL; }')
@@ -679,6 +690,7 @@ def write_out
 
 # Other required functions
 runner.deep_resolve('pg_printf')
+runner.deep_resolve('pg_strncasecmp')
 
 # Retain these functions for optional 32-bit support
 # (see BITS_PER_BITMAPWORD checks in bitmapset.c)

diff --git a/src/postgres/src_backend_catalog_namespace.c b/src/postgres/src_backend_catalog_namespace.c
@@ -905,7 +905,7 @@ NameListToString(List *names)
  * Note that this will only find collations that work with the current
  * database's encoding.
  */
-Oid get_collation_oid(List *name, bool missing_ok) { return -1; }
+Oid get_collation_oid(List *name, bool missing_ok) { return DEFAULT_COLLATION_OID; }
 
 
 /*

diff --git a/src/postgres/src_pl_plpgsql_src_pl_gram.c b/src/postgres/src_pl_plpgsql_src_pl_gram.c
@@ -6091,15 +6091,26 @@ plpgsql_sql_error_callback(void *arg)
  * expect that the given string is a copy from the source text.
  */
 
+#include "catalog/pg_collation_d.h"
 static PLpgSQL_type * parse_datatype(const char *string, int location) {
 	PLpgSQL_type *typ;
+
+	/* Ignore trailing spaces */
+	size_t len = strlen(string);
+	while (len > 0 && scanner_isspace(string[len - 1])) --len;
+
 	typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type));
 	typ->typname = pstrdup(string);
-	typ->ttype = pg_strcasecmp(string, "RECORD") == 0 ? PLPGSQL_TTYPE_REC : PLPGSQL_TTYPE_SCALAR;
-	if (pg_strcasecmp(string, "REFCURSOR") == 0 || pg_strcasecmp(string, "CURSOR") == 0)
+	typ->ttype = pg_strncasecmp(string, "RECORD", len) == 0 ? PLPGSQL_TTYPE_REC : PLPGSQL_TTYPE_SCALAR;
+	if (pg_strncasecmp(string, "REFCURSOR", len) == 0 || pg_strncasecmp(string, "CURSOR", len) == 0)
 	{
 		typ->typoid = REFCURSOROID;
 	}
+	else if (pg_strncasecmp(string, "TEXT", len) == 0)
+	{
+		typ->typoid = TEXTOID;
+		typ->collation = DEFAULT_COLLATION_OID;
+	}
 	return typ;
 }
 

diff --git a/src/postgres/src_port_pgstrcasecmp.c b/src/postgres/src_port_pgstrcasecmp.c
@@ -2,6 +2,7 @@
  * Symbols referenced in this file:
  * - pg_strcasecmp
  * - pg_toupper
+ * - pg_strncasecmp
  *--------------------------------------------------------------------
  */
 
@@ -72,7 +73,34 @@ pg_strcasecmp(const char *s1, const char *s2)
  * Case-independent comparison of two not-necessarily-null-terminated strings.
  * At most n bytes will be examined from each string.
  */
+int
+pg_strncasecmp(const char *s1, const char *s2, size_t n)
+{
+	while (n-- > 0)
+	{
+		unsigned char ch1 = (unsigned char) *s1++;
+		unsigned char ch2 = (unsigned char) *s2++;
 
+		if (ch1 != ch2)
+		{
+			if (ch1 >= 'A' && ch1 <= 'Z')
+				ch1 += 'a' - 'A';
+			else if (IS_HIGHBIT_SET(ch1) && isupper(ch1))
+				ch1 = tolower(ch1);
+
+			if (ch2 >= 'A' && ch2 <= 'Z')
+				ch2 += 'a' - 'A';
+			else if (IS_HIGHBIT_SET(ch2) && isupper(ch2))
+				ch2 = tolower(ch2);
+
+			if (ch1 != ch2)
+				return (int) ch1 - (int) ch2;
+		}
+		if (ch1 == 0)
+			break;
+	}
+	return 0;
+}
 
 /*
  * Fold a character to upper case.

diff --git a/test/plpgsql_samples.expected.json b/test/plpgsql_samples.expected.json
@@ -25,5 +25,6 @@
 {"PLpgSQL_function":{"datums":[{"PLpgSQL_var":{"refname":"$1","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.float4"}}}},{"PLpgSQL_var":{"refname":"found","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.\"boolean\""}}}}],"action":{"PLpgSQL_stmt_block":{"lineno":4,"body":[{"PLpgSQL_stmt_return":{"lineno":5,"expr":{"PLpgSQL_expr":{"query":"subtotal * 0.06","parseMode":2}}}}]}}}},
 {"PLpgSQL_function":{"datums":[{"PLpgSQL_var":{"refname":"s","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.int4"}}}},{"PLpgSQL_var":{"refname":"e","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.int4"}}}},{"PLpgSQL_var":{"refname":"id","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.int4"}}}},{"PLpgSQL_var":{"refname":"found","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.\"boolean\""}}}}],"action":{"PLpgSQL_stmt_block":{"lineno":2,"body":[{"PLpgSQL_stmt_assign":{"lineno":3,"varno":2,"expr":{"PLpgSQL_expr":{"query":"id := s","parseMode":3}}}},{"PLpgSQL_stmt_loop":{"lineno":4,"body":[{"PLpgSQL_stmt_exit":{"lineno":5,"is_exit":true,"cond":{"PLpgSQL_expr":{"query":"id\u003ee","parseMode":2}}}},{"PLpgSQL_stmt_return_next":{"lineno":6}},{"PLpgSQL_stmt_assign":{"lineno":7,"varno":2,"expr":{"PLpgSQL_expr":{"query":"id := id + 1","parseMode":3}}}}]}},{"PLpgSQL_stmt_return":{}}]}}}},
 {"PLpgSQL_function":{"new_varno":1,"old_varno":2,"datums":[{"PLpgSQL_var":{"refname":"found","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.\"boolean\""}}}},{"PLpgSQL_rec":{"refname":"new","dno":1}},{"PLpgSQL_rec":{"refname":"old","dno":2}}],"action":{"PLpgSQL_stmt_block":{"lineno":5,"body":[{"PLpgSQL_stmt_return":{"lineno":6}}]}}}},
-{"PLpgSQL_function":{"datums":[{"PLpgSQL_var":{"refname":"found","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.\"boolean\""}}}},{"PLpgSQL_var":{"refname":"ref","lineno":3,"datatype":{"PLpgSQL_type":{"typname":"refcursor"}}}}],"action":{"PLpgSQL_stmt_block":{"lineno":4,"body":[{"PLpgSQL_stmt_open":{"lineno":5,"curvar":1,"cursor_options":256,"query":{"PLpgSQL_expr":{"query":"SELECT col FROM test","parseMode":0}}}},{"PLpgSQL_stmt_return":{"lineno":6,"expr":{"PLpgSQL_expr":{"query":"ref","parseMode":2}}}}]}}}}
+{"PLpgSQL_function":{"datums":[{"PLpgSQL_var":{"refname":"found","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.\"boolean\""}}}},{"PLpgSQL_var":{"refname":"ref","lineno":3,"datatype":{"PLpgSQL_type":{"typname":"refcursor"}}}}],"action":{"PLpgSQL_stmt_block":{"lineno":4,"body":[{"PLpgSQL_stmt_open":{"lineno":5,"curvar":1,"cursor_options":256,"query":{"PLpgSQL_expr":{"query":"SELECT col FROM test","parseMode":0}}}},{"PLpgSQL_stmt_return":{"lineno":6,"expr":{"PLpgSQL_expr":{"query":"ref","parseMode":2}}}}]}}}},
+{"PLpgSQL_function":{"datums":[{"PLpgSQL_var":{"refname":"a","datatype":{"PLpgSQL_type":{"typname":"text"}}}},{"PLpgSQL_var":{"refname":"b","datatype":{"PLpgSQL_type":{"typname":"text"}}}},{"PLpgSQL_var":{"refname":"found","datatype":{"PLpgSQL_type":{"typname":"pg_catalog.\"boolean\""}}}},{"PLpgSQL_var":{"refname":"local_a","lineno":3,"datatype":{"PLpgSQL_type":{"typname":"text "}},"default_val":{"PLpgSQL_expr":{"query":"a","parseMode":2}}}},{"PLpgSQL_var":{"refname":"local_b","lineno":4,"datatype":{"PLpgSQL_type":{"typname":"text "}},"default_val":{"PLpgSQL_expr":{"query":"b","parseMode":2}}}}],"action":{"PLpgSQL_stmt_block":{"lineno":5,"body":[{"PLpgSQL_stmt_return":{"lineno":6,"expr":{"PLpgSQL_expr":{"query":"local_a \u003c local_b","parseMode":2}}}}]}}}}
 ]
diff --git a/test/plpgsql_samples.sql b/test/plpgsql_samples.sql
@@ -600,3 +600,13 @@ BEGIN
     RETURN ref;
 END;
 ' LANGUAGE plpgsql;
+
+-- Example from https://www.postgresql.org/docs/16/plpgsql-declarations.html#PLPGSQL-DECLARATION-COLLATION
+CREATE FUNCTION less_than(a text, b text) RETURNS boolean AS $$
+DECLARE
+    local_a text COLLATE "en_US" := a;
+    local_b text := b;
+BEGIN
+    RETURN local_a < local_b;
+END;
+$$ LANGUAGE plpgsql;