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/babelfishpg_common.c

@@ -20,6 +20,7 @@
 #include "typecode.h"
 #include "varchar.h"
 #include "datetimeoffset.h"
+#include "tcop/utility.h"
 
 common_utility_plugin common_utility_plugin_var = {NULL};
 static common_utility_plugin *get_common_utility_plugin(void);
@@ -33,6 +34,21 @@ char	   *pltsql_server_collation_name = NULL;
 /* Dump and Restore */
 char	   *babelfish_restored_server_collation_name = NULL;
 
+/* Global cache variable */
+LogicalDBCache g_logical_db_cache = {NULL, false};
+
+/* Previous ProcessUtility hook if any */
+static ProcessUtility_hook_type prev_ProcessUtility = NULL;
+static void
+babelfishpg_logicalDB_ProcessUtility(PlannedStmt *pstmt,
+						const char *queryString,
+						bool readOnlyTree,
+						ProcessUtilityContext context,
+						ParamListInfo params,
+						QueryEnvironment *queryEnv,
+						DestReceiver *dest,
+						QueryCompletion *qc);
+
 const char *BabelfishTranslateCollation(
 										const char *collname,
 										Oid collnamespace,
@@ -149,6 +165,14 @@ _PG_init(void)
  	prev_get_like_collation_hook = get_like_collation_hook;
 	get_like_collation_hook = bbf_get_like_collation;
 
+	/* Initialize logical db name */
+	g_logical_db_cache.logical_db_name = NULL;
+	g_logical_db_cache.is_valid = false;
+
+	/* Save previous hook and install hook for logical db name */
+	prev_ProcessUtility = ProcessUtility_hook;
+	ProcessUtility_hook = babelfishpg_logicalDB_ProcessUtility;
+
 }
 void
 _PG_fini(void)
@@ -158,6 +182,8 @@ _PG_fini(void)
 	PreCreateCollation_hook = prev_PreCreateCollation_hook;
 	set_like_collation_hook = prev_set_like_collation_hook;
 	get_like_collation_hook = prev_get_like_collation_hook;
+	/* Restore previous hook */
+	ProcessUtility_hook = prev_ProcessUtility;
 }
 
 common_utility_plugin *
@@ -219,3 +245,61 @@ get_common_utility_plugin(void)
 	}
 	return &common_utility_plugin_var;
 }
+
+/* Hook function to catch utility commands */
+static void
+babelfishpg_logicalDB_ProcessUtility(PlannedStmt *pstmt,
+						const char *queryString,
+						bool readOnlyTree,
+						ProcessUtilityContext context,
+						ParamListInfo params,
+						QueryEnvironment *queryEnv,
+						DestReceiver *dest,
+						QueryCompletion *qc)
+{
+	Node *parsetree = pstmt->utilityStmt;
+
+	switch (nodeTag(parsetree))
+	{
+		/* SET [ SESSION | LOCAL ] configuration_parameter */
+		case T_VariableSetStmt:
+		{
+			VariableSetStmt *stmt = (VariableSetStmt *) parsetree;
+			if (strcasecmp(stmt->name, "psql_logical_babelfish_db_name") == 0)
+			{
+				g_logical_db_cache.is_valid = false;
+			}
+			break;
+		}
+		/* ALTER DATABASE name SET configuration_parameter */
+		case T_AlterDatabaseSetStmt:
+		{
+			AlterDatabaseSetStmt *stmt = (AlterDatabaseSetStmt *) parsetree;
+			if (strcasecmp(stmt->setstmt->name, "psql_logical_babelfish_db_name") == 0)
+			{
+				g_logical_db_cache.is_valid = false;
+			}
+			break;
+		}
+		/* ALTER SYSTEM SET configuration_parameter */
+		case T_AlterSystemStmt:
+		{
+			AlterSystemStmt *stmt = (AlterSystemStmt *) parsetree;
+			if (strcasecmp(stmt->setstmt->name, "psql_logical_babelfish_db_name") == 0)
+			{
+				g_logical_db_cache.is_valid = false;
+			}
+			break;
+		}
+		default:
+			break;
+	}
+
+	/* Chain to previous hook if any */
+	if (prev_ProcessUtility)
+		prev_ProcessUtility(pstmt, queryString, readOnlyTree,
+						context, params, queryEnv, dest, qc);
+	else
+		standard_ProcessUtility(pstmt, queryString, readOnlyTree,
+							context, params, queryEnv, dest, qc);
+}

contrib/babelfishpg_common/src/babelfishpg_common.h

@@ -96,3 +96,11 @@ typedef struct common_utility_plugin
 	void		(*TsqlUTF8toUTF16StringInfo) (StringInfo utf16_data, const void *data, size_t len);
 	int32_t		(*tsql_numeric_get_typmod) (Numeric num);
 } common_utility_plugin;
+
+/* struct to save the state of the psql_logical_babelfish_db_name GUC */
+typedef struct LogicalDBCache {
+	const char *logical_db_name;
+	bool is_valid;
+} LogicalDBCache;
+
+extern LogicalDBCache g_logical_db_cache;

contrib/babelfishpg_common/src/collation.c

@@ -20,6 +20,7 @@
 #include "encoding/encoding.h"
 #include "typecode.h"
 #include "sqlvariant.h"
+#include "babelfishpg_common.h"
 
 #define NOT_FOUND -1
 
@@ -1240,6 +1241,32 @@ get_database_or_server_collation_oid_internal(bool missingOk)
 	return server_collation_oid;
 }
 
+static const char *
+get_cached_logical_db_name(void)
+{
+	if (!g_logical_db_cache.is_valid)
+	{
+		MemoryContext oldcontext;
+		const char* logical_db = NULL;
+		oldcontext = MemoryContextSwitchTo(CacheMemoryContext);
+
+		logical_db = GetConfigOption("psql_logical_babelfish_db_name", true, true);
+		if (logical_db)
+		{
+			g_logical_db_cache.logical_db_name = MemoryContextStrdup(CacheMemoryContext, logical_db);
+		}
+		else
+		{
+			g_logical_db_cache.logical_db_name = NULL;
+		}
+		g_logical_db_cache.is_valid = true;
+		MemoryContextSwitchTo(oldcontext);
+	}
+
+	return g_logical_db_cache.logical_db_name;
+}
+
+
 Oid
 BABELFISH_CLUSTER_COLLATION_OID()
 {
@@ -1252,11 +1279,11 @@ BABELFISH_CLUSTER_COLLATION_OID()
 	}
 
 	/* 
-	 * If DMS, we can return the server level collation
+	 * If logical db collation is required from PG, we can return the server level collation
 	 * We assume that psql_logical_babelfish_db_name is a valid TSQL database name
 	 * Caller is reponsible to supply valid psql_logical_babelfish_db_name
 	 */
-	logical_babelfish_db_name = GetConfigOption("psql_logical_babelfish_db_name", true, true);
+	logical_babelfish_db_name = get_cached_logical_db_name();
 	if (logical_babelfish_db_name)
 		return get_database_or_server_collation_oid_internal(false);