Support a COLLATE clause in plpgsql variable declarations.

This allows the usual rules for assigning a collation to a local variable
to be overridden.  Per discussion, it seems appropriate to support this
rather than forcing all local variables to have the argument-derived
collation.

5 files changed, 136 insertions, 12 deletions

diff --git a/doc/src/sgml/plpgsql.sgml b/doc/src/sgml/plpgsql.sgml
index a04ab139123..1866e43e0e6 100644
--- a/doc/src/sgml/plpgsql.sgml
+++ b/doc/src/sgml/plpgsql.sgml
@@ -328,15 +328,17 @@ arow RECORD;
     <para>
      The general syntax of a variable declaration is:
 <synopsis>
-<replaceable>name</replaceable> <optional> CONSTANT </optional> <replaceable>type</replaceable> <optional> NOT NULL </optional> <optional> { DEFAULT | := } <replaceable>expression</replaceable> </optional>;
+<replaceable>name</replaceable> <optional> CONSTANT </optional> <replaceable>type</replaceable> <optional> COLLATE <replaceable>collation_name</replaceable> </optional> <optional> NOT NULL </optional> <optional> { DEFAULT | := } <replaceable>expression</replaceable> </optional>;
 </synopsis>
       The <literal>DEFAULT</> clause, if given, specifies the initial value assigned
       to the variable when the block is entered.  If the <literal>DEFAULT</> clause
       is not given then the variable is initialized to the
       <acronym>SQL</acronym> null value.
       The <literal>CONSTANT</> option prevents the variable from being
-      assigned to, so that its value will remain constant for the duration of
-      the block.
+      assigned to after initialization, so that its value will remain constant
+      for the duration of the block.
+      The <literal>COLLATE</> option specifies a collation to use for the
+      variable (see <xref linkend="plpgsql-declaration-collation">).
       If <literal>NOT NULL</>
       is specified, an assignment of a null value results in a run-time
       error. All variables declared as <literal>NOT NULL</>
@@ -768,9 +770,23 @@ $$ LANGUAGE plpgsql;
    </para>
 
    <para>
-    Explicit <literal>COLLATE</> clauses can be written inside a function
-    if it is desired to force a particular collation to be used regardless
-    of what the function is called with.  For example,
+    A local variable of a collatable data type can have a different collation
+    associated with it by including the <literal>COLLATE</> option in its
+    declaration, for example
+
+<programlisting>
+DECLARE
+    local_a text COLLATE "en_US";
+</programlisting>
+
+    This option overrides the collation that would otherwise be
+    given to the variable according to the rules above.
+   </para>
+
+   <para>
+    Also, of course explicit <literal>COLLATE</> clauses can be written inside
+    a function if it is desired to force a particular collation to be used in
+    a particular operation.  For example,
 
 <programlisting>
 CREATE FUNCTION less_than_c(a text, b text) RETURNS boolean AS $$
@@ -779,6 +795,10 @@ BEGIN
 END;
 $$ LANGUAGE plpgsql;
 </programlisting>
+
+    This overrides the collations associated with the table columns,
+    parameters, or local variables used in the expression, just as would
+    happen in a plain SQL command.
    </para>
   </sect2>
   </sect1>
diff --git a/src/pl/plpgsql/src/gram.y b/src/pl/plpgsql/src/gram.y
index fbd441a1bc9..4e2b7058f0c 100644
--- a/src/pl/plpgsql/src/gram.y
+++ b/src/pl/plpgsql/src/gram.y
@@ -21,6 +21,7 @@
 #include "parser/parse_type.h"
 #include "parser/scanner.h"
 #include "parser/scansup.h"
+#include "utils/builtins.h"
 
 
 /* Location tracking support --- simpler than bison's default */
@@ -122,6 +123,7 @@ static	List			*read_raise_options(void);
 		PLcword					cword;
 		PLwdatum				wdatum;
 		bool					boolean;
+		Oid						oid;
 		struct
 		{
 			char *name;
@@ -167,6 +169,7 @@ static	List			*read_raise_options(void);
 %type <boolean>	decl_const decl_notnull exit_type
 %type <expr>	decl_defval decl_cursor_query
 %type <dtype>	decl_datatype
+%type <oid>		decl_collate
 %type <datum>	decl_cursor_args
 %type <list>	decl_cursor_arglist
 %type <nsitem>	decl_aliasitem
@@ -245,6 +248,7 @@ static	List			*read_raise_options(void);
 %token <keyword>	K_BY
 %token <keyword>	K_CASE
 %token <keyword>	K_CLOSE
+%token <keyword>	K_COLLATE
 %token <keyword>	K_CONSTANT
 %token <keyword>	K_CONTINUE
 %token <keyword>	K_CURSOR
@@ -428,10 +432,27 @@ decl_stmt		: decl_statement
 					}
 				;
 
-decl_statement	: decl_varname decl_const decl_datatype decl_notnull decl_defval
+decl_statement	: decl_varname decl_const decl_datatype decl_collate decl_notnull decl_defval
 					{
 						PLpgSQL_variable	*var;
 
+						/*
+						 * If a collation is supplied, insert it into the
+						 * datatype.  We assume decl_datatype always returns
+						 * a freshly built struct not shared with other
+						 * variables.
+						 */
+						if (OidIsValid($4))
+						{
+							if (!OidIsValid($3->collation))
+								ereport(ERROR,
+										(errcode(ERRCODE_DATATYPE_MISMATCH),
+										 errmsg("collations are not supported by type %s",
+												format_type_be($3->typoid)),
+										 parser_errposition(@4)));
+							$3->collation = $4;
+						}
+
 						var = plpgsql_build_variable($1.name, $1.lineno,
 													 $3, true);
 						if ($2)
@@ -444,10 +465,10 @@ decl_statement	: decl_varname decl_const decl_datatype decl_notnull decl_defval
 										 errmsg("row or record variable cannot be CONSTANT"),
 										 parser_errposition(@2)));
 						}
-						if ($4)
+						if ($5)
 						{
 							if (var->dtype == PLPGSQL_DTYPE_VAR)
-								((PLpgSQL_var *) var)->notnull = $4;
+								((PLpgSQL_var *) var)->notnull = $5;
 							else
 								ereport(ERROR,
 										(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
@@ -455,10 +476,10 @@ decl_statement	: decl_varname decl_const decl_datatype decl_notnull decl_defval
 										 parser_errposition(@4)));
 
 						}
-						if ($5 != NULL)
+						if ($6 != NULL)
 						{
 							if (var->dtype == PLPGSQL_DTYPE_VAR)
-								((PLpgSQL_var *) var)->default_val = $5;
+								((PLpgSQL_var *) var)->default_val = $6;
 							else
 								ereport(ERROR,
 										(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
@@ -685,6 +706,19 @@ decl_datatype	:
 					}
 				;
 
+decl_collate	:
+					{ $$ = InvalidOid; }
+				| K_COLLATE T_WORD
+					{
+						$$ = get_collation_oid(list_make1(makeString($2.ident)),
+											   false);
+					}
+				| K_COLLATE T_CWORD
+					{
+						$$ = get_collation_oid($2.idents, false);
+					}
+				;
+
 decl_notnull	:
 					{ $$ = false; }
 				| K_NOT K_NULL
@@ -2432,7 +2466,8 @@ read_datatype(int tok)
 				yyerror("incomplete data type declaration");
 		}
 		/* Possible followers for datatype in a declaration */
-		if (tok == K_NOT || tok == '=' || tok == COLON_EQUALS || tok == K_DEFAULT)
+		if (tok == K_COLLATE || tok == K_NOT ||
+			tok == '=' || tok == COLON_EQUALS || tok == K_DEFAULT)
 			break;
 		/* Possible followers for datatype in a cursor_arg list */
 		if ((tok == ',' || tok == ')') && parenlevel == 0)
diff --git a/src/pl/plpgsql/src/pl_scanner.c b/src/pl/plpgsql/src/pl_scanner.c
index e8a2628f2f1..e1c0b625954 100644
--- a/src/pl/plpgsql/src/pl_scanner.c
+++ b/src/pl/plpgsql/src/pl_scanner.c
@@ -64,6 +64,7 @@ static const ScanKeyword reserved_keywords[] = {
 	PG_KEYWORD("by", K_BY, RESERVED_KEYWORD)
 	PG_KEYWORD("case", K_CASE, RESERVED_KEYWORD)
 	PG_KEYWORD("close", K_CLOSE, RESERVED_KEYWORD)
+	PG_KEYWORD("collate", K_COLLATE, RESERVED_KEYWORD)
 	PG_KEYWORD("continue", K_CONTINUE, RESERVED_KEYWORD)
 	PG_KEYWORD("declare", K_DECLARE, RESERVED_KEYWORD)
 	PG_KEYWORD("default", K_DEFAULT, RESERVED_KEYWORD)
diff --git a/src/test/regress/expected/collate.linux.utf8.out b/src/test/regress/expected/collate.linux.utf8.out
index f0008ddf14b..f225b487306 100644
--- a/src/test/regress/expected/collate.linux.utf8.out
+++ b/src/test/regress/expected/collate.linux.utf8.out
@@ -825,6 +825,46 @@ ORDER BY a.b, b.b;
  bbc | bbc | f  | f    | f              | f
 (16 rows)
 
+-- collation override in plpgsql
+CREATE FUNCTION mylt2 (x text, y text) RETURNS boolean LANGUAGE plpgsql AS $$
+declare
+  xx text := x;
+  yy text := y;
+begin
+  return xx < yy;
+end
+$$;
+SELECT mylt2('a', 'B' collate "en_US") as t, mylt2('a', 'B' collate "C") as f;
+ t | f 
+---+---
+ t | f
+(1 row)
+
+CREATE OR REPLACE FUNCTION
+  mylt2 (x text, y text) RETURNS boolean LANGUAGE plpgsql AS $$
+declare
+  xx text COLLATE "POSIX" := x;
+  yy text := y;
+begin
+  return xx < yy;
+end
+$$;
+SELECT mylt2('a', 'B') as f;
+ f 
+---
+ f
+(1 row)
+
+SELECT mylt2('a', 'B' collate "C") as fail; -- conflicting collations
+ERROR:  could not determine which collation to use for string comparison
+HINT:  Use the COLLATE clause to set the collation explicitly.
+CONTEXT:  PL/pgSQL function "mylt2" line 6 at RETURN
+SELECT mylt2('a', 'B' collate "POSIX") as f;
+ f 
+---
+ f
+(1 row)
+
 -- polymorphism
 SELECT * FROM unnest((SELECT array_agg(b ORDER BY b) FROM collate_test1)) ORDER BY 1;
  unnest 
diff --git a/src/test/regress/sql/collate.linux.utf8.sql b/src/test/regress/sql/collate.linux.utf8.sql
index 51d65cf0da8..dfb10e4d15b 100644
--- a/src/test/regress/sql/collate.linux.utf8.sql
+++ b/src/test/regress/sql/collate.linux.utf8.sql
@@ -256,6 +256,34 @@ FROM collate_test1 a, collate_test1 b
 ORDER BY a.b, b.b;
 
 
+-- collation override in plpgsql
+
+CREATE FUNCTION mylt2 (x text, y text) RETURNS boolean LANGUAGE plpgsql AS $$
+declare
+  xx text := x;
+  yy text := y;
+begin
+  return xx < yy;
+end
+$$;
+
+SELECT mylt2('a', 'B' collate "en_US") as t, mylt2('a', 'B' collate "C") as f;
+
+CREATE OR REPLACE FUNCTION
+  mylt2 (x text, y text) RETURNS boolean LANGUAGE plpgsql AS $$
+declare
+  xx text COLLATE "POSIX" := x;
+  yy text := y;
+begin
+  return xx < yy;
+end
+$$;
+
+SELECT mylt2('a', 'B') as f;
+SELECT mylt2('a', 'B' collate "C") as fail; -- conflicting collations
+SELECT mylt2('a', 'B' collate "POSIX") as f;
+
+
 -- polymorphism
 
 SELECT * FROM unnest((SELECT array_agg(b ORDER BY b) FROM collate_test1)) ORDER BY 1;