Add support for collation in partition functions and partition keys

contrib/babelfishpg_tsql/antlr/TSqlParser.g4

@@ -1163,7 +1163,7 @@ alter_partition_function
     ;
 
 create_partition_function
-    : CREATE PARTITION FUNCTION partition_function_name=id LR_BRACKET data_type RR_BRACKET AS RANGE (LEFT | RIGHT)? FOR VALUES LR_BRACKET expression_list? RR_BRACKET
+    : CREATE PARTITION FUNCTION partition_function_name=id LR_BRACKET data_type collation? RR_BRACKET AS RANGE (LEFT | RIGHT)? FOR VALUES LR_BRACKET expression_list? RR_BRACKET
     ;
 
 // https://docs.microsoft.com/en-us/sql/t-sql/statements/alter-partition-scheme-transact-sql

contrib/babelfishpg_tsql/runtime/functions.c

@@ -1987,6 +1987,8 @@ search_partition(PG_FUNCTION_ARGS)
 	Datum			arg;
 	Oid			argtypeid;
 	char			*func_param_typname = NULL;
+	char			*func_param_collation = NULL;
+	Oid			collation_oid = InvalidOid;
 	Oid			func_param_typoid;
 	Oid			sqlvariant_typoid;
 	Datum			*range_values;
@@ -1999,6 +2001,13 @@ search_partition(PG_FUNCTION_ARGS)
 	if (!PG_ARGISNULL(2)) /* Database is specified. */
 	{
 		char *db_name = text_to_cstring(PG_GETARG_TEXT_P(2));
+		/* Lowercase the db_name, if needed. */
+		if (pltsql_case_insensitive_identifiers)
+		{
+			char *tmp = db_name;
+			db_name = downcase_identifier(tmp, strlen(tmp), false, false);
+			pfree(tmp);
+		}
 		dbid = get_db_id(db_name);
 		if (!DbidIsValid(dbid))
 			ereport(ERROR,
@@ -2037,6 +2046,7 @@ search_partition(PG_FUNCTION_ARGS)
 	if (HeapTupleIsValid(tuple))
 	{
 		func_param_typname = TextDatumGetCString(heap_getattr(tuple, Anum_bbf_partition_function_input_parameter_type, RelationGetDescr(rel), &isnull));
+		func_param_collation = NameStr(*DatumGetName(heap_getattr(tuple, Anum_bbf_partition_function_input_parameter_collation, RelationGetDescr(rel), &isnull)));
 		values = DatumGetArrayTypeP(heap_getattr(tuple, Anum_bbf_partition_function_range_values, RelationGetDescr(rel), &isnull));
 		deconstruct_array(values, sqlvariant_typoid,
 					-1, false, 'i', &range_values, &nulls, &nelems);
@@ -2069,6 +2079,17 @@ search_partition(PG_FUNCTION_ARGS)
 
 	/* Get OID of partition function parameter type. */
 	func_param_typoid = (*common_utility_plugin_ptr->get_tsql_datatype_oid) (func_param_typname);
+
+	/* Get collation oid from partition function collation. */
+	if (func_param_collation)
+	{
+		collation_oid = get_collation_oid(list_make1(makeString(func_param_collation)), false);
+		/* Sanity check. */
+		if (!OidIsValid(collation_oid))
+			ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+					errmsg("Invalid collation '%s'.", func_param_collation)));
+	}
 
 	/* 
 	 * Implicitly convert input value to parameter type of
@@ -2096,7 +2117,7 @@ search_partition(PG_FUNCTION_ARGS)
 								ObjectIdGetDatum(get_namespace_oid("sys", false)));
 
 	cxt.function_oid = cmpfunction_oid;
-	cxt.colloid = tsql_get_database_or_server_collation_oid_internal(false);
+	cxt.colloid = collation_oid;
 
 	/* Perform binary search on sorted range values. */
 	result = tsql_bsearch_arg(&arg, range_values, nelems, sizeof(Datum), tsql_compare_values, &cxt);

contrib/babelfishpg_tsql/sql/babelfish_partitions.sql

@@ -10,6 +10,7 @@ CREATE TABLE sys.babelfish_partition_function
   range_values sys.sql_variant[] CHECK (array_length(range_values, 1) < 15000), -- boundary values
   create_date SYS.DATETIME NOT NULL,
   modify_date SYS.DATETIME NOT NULL,
+  input_parameter_collation NAME, -- collation of the input parameter
   PRIMARY KEY(dbid, partition_function_name)
 );
 

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--5.1.0--5.2.0.sql

@@ -450,6 +450,29 @@ BEGIN
 END;
 $$;
 
+DO $$
+BEGIN
+IF NOT EXISTS(
+    SELECT 1 FROM pg_class c JOIN pg_attribute a ON a.attrelid = c.oid 
+      WHERE c.relname = 'babelfish_partition_function' COLLATE sys.database_default
+      AND c.relnamespace::regnamespace::text = 'sys' COLLATE sys.database_default
+	  AND a.attname = 'input_parameter_collation' COLLATE sys.database_default)
+THEN
+    -- Add input_parameter_collation column in sys.babelfish_partition_function.
+    SET allow_system_table_mods = on;
+    ALTER TABLE sys.babelfish_partition_function ADD COLUMN input_parameter_collation NAME;
+    RESET allow_system_table_mods;
+
+    -- Update the input_parameter_collation column in sys.babelfish_partition_function
+    -- catalog for collatable datatypes with default database collation.
+    UPDATE sys.babelfish_partition_function pf
+    SET input_parameter_collation = db.default_collation
+    FROM sys.babelfish_sysdatabases db 
+    WHERE pf.dbid = db.dbid 
+    AND pf.input_parameter_type IN ('CHAR', 'VARCHAR', 'NCHAR', 'NVARCHAR');
+END IF;
+END $$;
+
 CREATE OR REPLACE FUNCTION sys.babelfish_try_conv_float_to_string(IN p_datatype TEXT,
 														  IN p_floatval FLOAT,
 														  IN p_style NUMERIC DEFAULT 0)

contrib/babelfishpg_tsql/src/catalog.c

@@ -5314,7 +5314,7 @@ is_partition_function_used(int16 dbid, const char *partition_function_name)
  */
 void
 add_entry_to_bbf_partition_function(int16 dbid, const char *partition_function_name, char *typname,
-					bool partition_option, ArrayType *values)
+					bool partition_option, ArrayType *values, char *collation)
 {
 	Relation	rel;
 	TupleDesc	dsc;
@@ -5339,6 +5339,15 @@ add_entry_to_bbf_partition_function(int16 dbid, const char *partition_function_n
 	new_record[Anum_bbf_partition_function_range_values - 1] = PointerGetDatum(values);
 	new_record[6] = new_record[7] = TimestampGetDatum(GetSQLLocalTimestamp(3));
 
+	if (collation)
+	{
+		NameData input_parameter_collation;
+		namestrcpy(&input_parameter_collation, collation);
+		new_record[Anum_bbf_partition_function_input_parameter_collation - 1] = NameGetDatum(&input_parameter_collation);
+	}
+	else
+		new_record_nulls[Anum_bbf_partition_function_input_parameter_collation - 1] = true;
+
 	tuple = heap_form_tuple(dsc, new_record, new_record_nulls);
 
 	/* insert new record in the bbf_partition_function table */

contrib/babelfishpg_tsql/src/catalog.h

@@ -432,8 +432,9 @@ extern Oid	get_bbf_extended_properties_idx_oid(void);
 #define Anum_bbf_partition_function_input_parameter_type 4
 #define Anum_bbf_partition_function_partition_option 5
 #define Anum_bbf_partition_function_range_values 6
+#define Anum_bbf_partition_function_input_parameter_collation 9
 
-#define BBF_PARTITION_FUNCTION_NUM_COLS 8
+#define BBF_PARTITION_FUNCTION_NUM_COLS 9
 
 extern Oid	bbf_partition_function_oid;
 extern Oid	bbf_partition_function_pk_idx_oid;
@@ -446,7 +447,7 @@ extern Oid	get_bbf_partition_function_id_idx_oid(void);
 extern Oid	get_bbf_partition_function_seq_oid(void);
 extern int32	get_available_partition_function_id(void);
 extern void	add_entry_to_bbf_partition_function(int16 dbid, const char *partition_function_name,
-						char *typname, bool partition_option, ArrayType *values);
+						char *typname, bool partition_option, ArrayType *values, char *collation);
 extern void	remove_entry_from_bbf_partition_function(int16 dbid, const char *partition_function_name);
 extern bool	partition_function_exists(int16 dbid, const char *partition_function_name);
 extern int	get_partition_count(int16 dbid, const char *partition_function_name);

contrib/babelfishpg_tsql/src/pl_exec-2.c

@@ -4298,6 +4298,9 @@ exec_stmt_partition_function(PLtsql_execstate *estate, PLtsql_stmt_partition_fun
 	int32			valtypmod;
 	Datum			tsql_type_datum;
 	char			*tsql_typename = NULL;
+	char			*collation = NULL;
+	Oid			collation_oid = InvalidOid;
+	bool			type_is_collatable;
 	Datum			*input_values;
 	Datum			*sql_variant_values;
 	ArrayType		*arr_value = NULL;
@@ -4337,6 +4340,20 @@ exec_stmt_partition_function(PLtsql_execstate *estate, PLtsql_stmt_partition_fun
 				errmsg("The identifier that starts with '%.128s' is too long. Maximum length is 128.", partition_function_name)));
 	}
 
+	/*
+	 * Get collation oid if collation is specified.
+	 */
+	if (stmt->collation)
+	{
+		collation_oid = tsql_get_oid_from_collidx(tsql_find_collation_internal(tsql_translate_tsql_collation_to_bbf_collation(stmt->collation)));
+
+		/* raise an error if specified collation is invalid */
+		if (!OidIsValid(collation_oid))
+			ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+					errmsg("Invalid collation '%s'.", stmt->collation)));
+	}
+
 	/* check if there is existing partition function with the given name in the current database */
 	if (partition_function_exists(dbid, partition_function_name))
 	{
@@ -4405,6 +4422,28 @@ exec_stmt_partition_function(PLtsql_execstate *estate, PLtsql_stmt_partition_fun
 			(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				errmsg("The type '%s' is not yet supported for partition function in Babelfish.", tsql_typename)));
 
+	type_is_collatable = OidIsValid(typ->collation);
+
+	/*
+	 * Raise an error if collate clause is specified and datatype is not collatable.
+	 */
+	if (stmt->collation && !type_is_collatable)
+	{
+
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+			 	errmsg("Expression type '%s' is invalid for COLLATE clause.", tsql_typename)));
+	}
+	/*
+	 * Use default database collation if collate clause is not specified and datatype is collatable.
+	 */
+	else if (stmt->collation == NULL && type_is_collatable)
+		collation_oid = tsql_get_database_or_server_collation_oid_internal(false);
+
+	/* get collation name from collation oid when type is collatable */
+	if (type_is_collatable)
+		collation = get_collation_name(collation_oid);
+
 	/* check if the given number of boundaries are exceeding allowed limit */
 	nargs = list_length(arg);
 	if (nargs >= MAX_PARTITIONS_LIMIT)
@@ -4464,7 +4503,7 @@ exec_stmt_partition_function(PLtsql_execstate *estate, PLtsql_stmt_partition_fun
 
 	/* set the function oid of operator in tsql comparator context */
 	cxt.function_oid = cmpfunction_oid;
-	cxt.colloid = tsql_get_database_or_server_collation_oid_internal(false);
+	cxt.colloid = collation_oid;
 	cxt.contains_duplicate = false;
 
 	/* 
@@ -4498,12 +4537,14 @@ exec_stmt_partition_function(PLtsql_execstate *estate, PLtsql_stmt_partition_fun
 					-1, false, 'i');
 
 	/* add entry in the sys.babelfish_partition_function catalog */
-	add_entry_to_bbf_partition_function(dbid, partition_function_name, tsql_typename, stmt->is_right, arr_value);
+	add_entry_to_bbf_partition_function(dbid, partition_function_name, tsql_typename, stmt->is_right, arr_value, collation);
 
 	pfree(tsql_typename);
 	pfree(input_values);
 	pfree(sql_variant_values);
 	pfree(arr_value);
+	if (collation)
+		pfree(collation);
 
 	/* cleanup estate */
 	exec_eval_cleanup(estate);

contrib/babelfishpg_tsql/src/pltsql.h

@@ -1110,6 +1110,7 @@ typedef struct PLtsql_stmt_partition_function
 	bool			is_right;
 	PLtsql_type		*datatype;
 	List			*args;		/* the arguments (list of exprs) */
+	char			*collation;
 } PLtsql_stmt_partition_function;
 
 /*

contrib/babelfishpg_tsql/src/pltsql_partition.c

@@ -65,15 +65,18 @@ bbf_create_partition_tables(CreateStmt *stmt)
 	SysScanDesc	scan;
 	ScanKeyData	scanKey[2];
 	char		*input_parameter_type;
+	char		*input_parameter_collation;
 	char		*partition_function_name;
 	Datum		*range_values;
 	Datum		*datum_values;
 	bool		*nulls;
+	bool		type_is_collatable = InvalidOid;
 	int		nelems;
 	Oid		sql_variant_type_oid;
 	Oid		input_type_oid;
 	ListCell	*elements;
 	Oid		partition_column_typoid = InvalidOid;
+	Oid		partition_column_colloid = InvalidOid;
 	Oid		partition_column_basetypoid = InvalidOid;
 	char		*partition_column_typname = NULL;
 	bool		is_binary_datatype = false;
@@ -130,6 +133,11 @@ bbf_create_partition_tables(CreateStmt *stmt)
 			partition_column_typoid = pg_type->oid;
 			partition_column_basetypoid = pg_type->typbasetype;
 			partition_column_typname = pstrdup(NameStr(pg_type->typname));
+			type_is_collatable = OidIsValid(pg_type->typcollation);
+			if (coldef->collClause)
+				partition_column_colloid = get_collation_oid(coldef->collClause->collname, false);
+			else
+				partition_column_colloid = tsql_get_database_or_server_collation_oid_internal(false);;
 			ReleaseSysCache(ctype);
 			break;
 		}
@@ -170,6 +178,7 @@ bbf_create_partition_tables(CreateStmt *stmt)
 	}
 
 	input_parameter_type = TextDatumGetCString(heap_getattr(tuple, Anum_bbf_partition_function_input_parameter_type, RelationGetDescr(rel), &isnull));
+	input_parameter_collation = NameStr(*DatumGetName(heap_getattr(tuple, Anum_bbf_partition_function_input_parameter_collation, RelationGetDescr(rel), &isnull)));
 	values = DatumGetArrayTypeP(heap_getattr(tuple, Anum_bbf_partition_function_range_values, RelationGetDescr(rel), &isnull));
 	deconstruct_array(values, sql_variant_type_oid, -1, false, 'i', &datum_values, &nulls, &nelems);
 
@@ -210,6 +219,16 @@ bbf_create_partition_tables(CreateStmt *stmt)
 	if ((*common_utility_plugin_ptr->is_tsql_binary_datatype) (input_type_oid) ||
 		(*common_utility_plugin_ptr->is_tsql_varbinary_datatype) (input_type_oid))
 			is_binary_datatype = true;
+
+	if (type_is_collatable)
+	{
+		Assert(input_parameter_collation != NULL);
+		if (partition_column_colloid != tsql_get_oid_from_collidx(tsql_find_collation_internal(input_parameter_collation)))
+			ereport(ERROR, 
+				(errcode(ERRCODE_UNDEFINED_OBJECT), 
+					errmsg("Collation of partition column '%s' does not match collation of corresponding parameter in partition function '%s'.",
+					partition_colname, partition_function_name)));
+	}
 
 	/* Convert each sql_variant values to CString. */
 	range_values = palloc(nelems * sizeof(Datum));

contrib/babelfishpg_tsql/src/tsqlIface.cpp

@@ -7845,6 +7845,10 @@ makeCreatePartitionFunction(TSqlParser::Create_partition_functionContext *ctx)
 	std::string typeStr = ::getFullText(ctx->data_type());
 	PLtsql_type *type = parse_datatype(typeStr.c_str(), 0);
 
+	if (ctx->collation())
+		stmt->collation = pstrdup(getFullText(ctx->collation()->id()).c_str());
+	else
+		stmt->collation = NULL;
 	stmt->function_name = pstrdup(stripQuoteFromId(ctx->id()).c_str());
 	stmt->datatype = type;
 	stmt->lineno = getLineNo(ctx);