Fix loading of data for TEXT column (NULL allowed) if primary key is defined using non-default collation from PG endpoint for logical database

contrib/babelfishpg_common/src/collation.c

@@ -26,7 +26,7 @@
 #define DATABASE_DEFAULT "database_default"
 #define CATALOG_DEFAULT "catalog_default"
 
-collation_callbacks collation_callbacks_var = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
+collation_callbacks collation_callbacks_var = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
 
 /* Cached values derived from server_collation_name */
 static int	server_collation_collidx = NOT_FOUND;
@@ -36,6 +36,8 @@ static bool db_collation_is_CI = true;
 static Oid database_collation_oid = InvalidOid;
 static int database_collation_collidx = NOT_FOUND;
 
+static bool is_logical_db_name_set = false;
+
 /*
  * Below two vars are defined to store the value of the babelfishpg_tsql.server_collation_name
  * and babelfishpg_tsql.default_locale.
@@ -1249,6 +1251,14 @@ BABELFISH_CLUSTER_COLLATION_OID()
 													 * database or server_collation_oid */
 		return db_coll;
 	}
+
+	/* 
+	 * If logical db collation is required from PG, we can return the server level collation
+	 * this is decided by is_logical_db_name_set, which is set if psql_logical_babelfish_db_name is set
+	 */
+	if (is_logical_db_name_set)
+		return get_database_or_server_collation_oid_internal(false);
+
 	return DEFAULT_COLLATION_OID;
 }
 
@@ -1647,6 +1657,7 @@ get_collation_callbacks(void)
 		collation_callbacks_var.translate_bbf_collation_to_tsql_collation = &translate_bbf_collation_to_tsql_collation;
 		collation_callbacks_var.translate_tsql_collation_to_bbf_collation = &translate_tsql_collation_to_bbf_collation;
 		collation_callbacks_var.set_db_collation = &set_db_collation;
+		collation_callbacks_var.set_logical_db_name_cache = &set_logical_db_name_cache;
 	}
 	return &collation_callbacks_var;
 }
@@ -1749,4 +1760,17 @@ set_db_collation(Oid db_coll)
 	database_collation_oid = db_coll;
 	database_collation_collidx = find_any_collation((lookup_collation_table(database_collation_oid).collname), false);
 	db_collation_is_CI = collation_is_CI(database_collation_oid);
-}
+}
+
+/* 
+ * Set the cache for logical database name
+ * this is decided by "psql_logical_babelfish_db_name" GUC
+ */
+void
+set_logical_db_name_cache(const char* dbname)
+{
+	if (dbname)
+		is_logical_db_name_set = true;
+	else
+		is_logical_db_name_set = false;
+}

contrib/babelfishpg_common/src/collation.h

@@ -121,6 +121,8 @@ typedef struct collation_callbacks
 
 	void		(*set_db_collation) (Oid db_coll);
 
+	void		(*set_logical_db_name_cache) (const char *dbname);
+
 } collation_callbacks;
 
 extern int	find_cs_as_collation(int collidx);
@@ -152,6 +154,7 @@ extern bool has_ilike_node(Node *expr);
 extern bool has_like_node(Node *expr);
 extern Oid	babelfish_define_type_default_collation(Oid typeNamespace);
 extern void	set_db_collation(Oid db_coll);
+extern void set_logical_db_name_cache(const char *dbname);
 
 extern collation_callbacks *get_collation_callbacks(void);
 

contrib/babelfishpg_tsql/src/collation.c

@@ -25,6 +25,7 @@
 #include <unicode/usearch.h>
 #endif
 
+#include "miscadmin.h"
 #include "pltsql.h"
 #include "src/collation.h"
 #include "catalog.h"
@@ -2174,3 +2175,21 @@ patindex_ai_collations(PG_FUNCTION_ARGS)
 
 	PG_RETURN_INT32(result);
 }
+
+
+void
+assign_psql_logical_babelfish_db_name(const char *newval, void *extra)
+{
+	/* 
+	 * No need to tell babelfishpg_common that logical db name has been set if it is TDS connection
+	 * because the psql_logical_babelfish_db_name is to be used for PG dialect only
+	 */
+	if (!IS_TDS_CONN())
+	{
+		/* Initialise collation callbacks */
+		init_and_check_collation_callbacks();
+
+		/* let babelfishpg_common know that psql_logical_babelfish_db_name has been updated */
+		(*collation_callbacks_ptr->set_logical_db_name_cache) (newval);
+	}
+}

contrib/babelfishpg_tsql/src/collation.h

@@ -90,6 +90,8 @@ typedef struct collation_callbacks
 
 	void		(*set_db_collation) (Oid db_coll);
 
+	void		(*set_logical_db_name_cache) (const char *dbname);
+
 } collation_callbacks;
 
 extern collation_callbacks *collation_callbacks_ptr;

contrib/babelfishpg_tsql/src/guc.c

@@ -825,7 +825,7 @@ define_custom_variables(void)
 							   NULL,
 							   PGC_USERSET,
 							   GUC_NOT_IN_SAMPLE | GUC_NO_SHOW_ALL | GUC_NO_RESET_ALL | GUC_DISALLOW_IN_FILE | GUC_DISALLOW_IN_AUTO_FILE,
-							   NULL, NULL, NULL);
+							   NULL, assign_psql_logical_babelfish_db_name, NULL);
 
 	DefineCustomIntVariable("babelfishpg_tsql.datefirst",
 							gettext_noop("Sets the first day of the week to a number from 1 through 7."),

contrib/babelfishpg_tsql/src/guc.h

@@ -45,4 +45,6 @@ void		pltsql_revert_guc(int nest_level);
 extern int	pltsql_new_scope_identity_nest_level(void);
 extern void pltsql_revert_last_scope_identity(int nest_level);
 extern void pltsql_remove_current_query_env(void);
+
+extern void assign_psql_logical_babelfish_db_name(const char *newval, void *extra);
 #endif

contrib/babelfishpg_tsql/src/hooks.c

@@ -5984,8 +5984,15 @@ default_collation_for_builtin_type(Type typ, bool handle_pg_type)
 	 * Special handling for PG datatypes such as TEXT because Babelfish does not define sys.TEXT.
 	 * This is required as Babelfish currently does not handle collation of String Const node correctly.
 	 * TODO: Fix the handling of the collation for String Const node.
+	 * 
+	 * When pltsql_psql_logical_babelfish_db_name is set, 
+	 * return CLUSTER_COLLATION_OID() -- this API will handle whether to return babelfish server/database level
+	 * collation or DEFAULT_COLLTAION_OID based on dialect for TEXT (e.g: string literals) as well
+	 * Otherwise, during collation resolution in merge_collation_state, CLUSTER_COLLATION_OID() will return
+	 * babelfish server/db level collation as pltsql_psql_logical_babelfish_db_name is set but current collation would
+	 * be DEFAULT_COLLATION_OID which would throw collation conflict
 	 */
-	if (handle_pg_type && oid == DEFAULT_COLLATION_OID)
+	if ((handle_pg_type || pltsql_psql_logical_babelfish_db_name) && oid == DEFAULT_COLLATION_OID)
 	{
 		oid = CLUSTER_COLLATION_OID();
 	}