Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

84 changes: 84 additions & 0 deletions contrib/babelfishpg_common/src/babelfishpg_common.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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 *
Expand Down Expand Up @@ -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);
}
8 changes: 8 additions & 0 deletions contrib/babelfishpg_common/src/babelfishpg_common.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
31 changes: 29 additions & 2 deletions contrib/babelfishpg_common/src/collation.c
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include "encoding/encoding.h"
#include "typecode.h"
#include "sqlvariant.h"
#include "babelfishpg_common.h"

#define NOT_FOUND -1

Expand Down Expand Up @@ -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()
{
Expand All @@ -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);

Expand Down
Loading
Loading