Skip to content

Commit

Permalink
Do not dump babelfish initialize user
Browse files Browse the repository at this point in the history
Signed-off-by: Rishabh Tanwar <[email protected]>
  • Loading branch information
ritanwar committed Sep 13, 2023
1 parent d060983 commit 9bbf95f
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 26 deletions.
65 changes: 48 additions & 17 deletions src/bin/pg_dump/dump_babel_utils.c
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ static const CatalogId nilCatalogId = {0, 0};
static char *escaped_bbf_db_name = NULL;
static int bbf_db_id = 0;
static SimpleOidList catalog_table_include_oids = {NULL, NULL};
static char *babel_init_user = NULL;

static char *getMinOid(Archive *fout);
static bool isBabelfishConfigTable(TableInfo *tbinfo);
Expand Down Expand Up @@ -223,6 +224,7 @@ dumpBabelRestoreChecks(Archive *fout)
char *source_server_version;
char *source_migration_mode;
PQExpBuffer qry;
ArchiveFormat format = ((ArchiveHandle *) fout)->format;

/* Skip if not Babelfish database or binary upgrade */
if (!isBabelfishDatabase(fout) || fout->dopt->binary_upgrade)
Expand All @@ -241,8 +243,10 @@ dumpBabelRestoreChecks(Archive *fout)
/*
* Temporarily enable ON_ERROR_STOP so that whole restore script
* execution fails if the following do block raises an error.
* Note that it can only be used in plain text dump (archNull).
*/
appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP on\n\n");
if (format == archNull)
appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP on\n\n");
appendPQExpBuffer(qry,
"DO $$"
"\nDECLARE"
Expand Down Expand Up @@ -276,7 +280,8 @@ dumpBabelRestoreChecks(Archive *fout)
"\n END IF;"
"\nEND$$;\n\n"
, source_migration_mode);
appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP off\n");
if (format == archNull)
appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP off\n");
PQclear(res);

ArchiveEntry(fout, nilCatalogId, createDumpId(),
Expand Down Expand Up @@ -844,7 +849,7 @@ updateExtConfigArray(Archive *fout, char ***extconfigarray, int nconfigitems)
}

PQclear(res);
resetPQExpBuffer(query);
destroyPQExpBuffer(query);
}

/*
Expand Down Expand Up @@ -890,6 +895,18 @@ prepareForBabelfishDatabaseDump(Archive *fout, SimpleStringList *schema_include_
for (i = 0; i < ntups; i++)
simple_oid_list_append(&catalog_table_include_oids, atooid(PQgetvalue(res, i, 0)));

PQclear(res);
resetPQExpBuffer(query);

/*
* Find out initialize user of current Babelfish database.
*/
appendPQExpBufferStr(query, "SELECT r.rolname FROM pg_roles r "
"INNER JOIN pg_database d ON r.oid = d.datdba "
"WHERE d.datname = current_database()");
res = ExecuteSqlQueryForSingleRow(fout, query->data);
babel_init_user = pstrdup(PQgetvalue(res, 0, 0));

PQclear(res);
destroyPQExpBuffer(query);

Expand Down Expand Up @@ -1090,8 +1107,8 @@ addFromClauseForPhysicalDatabaseDump(PQExpBuffer buf, TableInfo *tbinfo)
}
else if(strcmp(tbinfo->dobj.name, "babelfish_authid_login_ext") == 0)
appendPQExpBuffer(buf, " FROM ONLY %s a "
"WHERE a.rolname!='sysadmin'",
fmtQualifiedDumpable(tbinfo));
"WHERE a.rolname NOT IN ('sysadmin', '%s')", /* Do not dump sysadmin and Babelfish initialize user */
fmtQualifiedDumpable(tbinfo), babel_init_user);
else if(strcmp(tbinfo->dobj.name, "babelfish_domain_mapping") == 0 ||
strcmp(tbinfo->dobj.name, "babelfish_function_ext") == 0 ||
strcmp(tbinfo->dobj.name, "babelfish_view_def") == 0 ||
Expand All @@ -1117,6 +1134,7 @@ fixCursorForBbfCatalogTableData(Archive *fout, TableInfo *tbinfo, PQExpBuffer bu
int i;
bool is_builtin_db = false;
bool is_bbf_usr_ext_tab = false;
bool is_bbf_sysdatabases_tab = false;

/*
* Return if not a Babelfish database, or if the table is not a Babelfish
Expand All @@ -1132,9 +1150,11 @@ fixCursorForBbfCatalogTableData(Archive *fout, TableInfo *tbinfo, PQExpBuffer bu
pg_strcasecmp(bbf_db_name, "msdb") == 0)
? true : false;

/* Remember if it is babelfish_authid_user_ext catalog table. */
/* Remember if it is babelfish_authid_user_ext and babelfish_sysdatabases catalog table. */
if (strcmp(tbinfo->dobj.name, "babelfish_authid_user_ext") == 0)
is_bbf_usr_ext_tab = true;
if (strcmp(tbinfo->dobj.name, "babelfish_sysdatabases") == 0)
is_bbf_sysdatabases_tab = true;

resetPQExpBuffer(buf);
appendPQExpBufferStr(buf, "DECLARE _pg_dump_cursor CURSOR FOR SELECT ");
Expand All @@ -1153,6 +1173,13 @@ fixCursorForBbfCatalogTableData(Archive *fout, TableInfo *tbinfo, PQExpBuffer bu
*/
if (bbf_db_name != NULL && !is_builtin_db && strcmp(tbinfo->attnames[i], "dbid") == 0)
continue;
/*
* We need to skip owner column of babelfish_sysdatabases table as it might be
* referencing Babelfish initialize user which we do not include in dump. We will
* populate this column during restore with the initialize user of target database.
*/
else if (is_bbf_sysdatabases_tab && strcmp(tbinfo->attnames[i], "owner") == 0)
continue;
if (*nfields > 0)
appendPQExpBufferStr(buf, ", ");
/*
Expand Down Expand Up @@ -1196,6 +1223,7 @@ fixCopyCommand(Archive *fout, PQExpBuffer copyBuf, TableInfo *tbinfo, bool isFro
bool is_builtin_db = false;
bool needComma = false;
bool is_bbf_usr_ext_tab = false;
bool is_bbf_sysdatabases_tab = false;

/*
* Return if not a Babelfish database, or if the table is not a Babelfish
Expand All @@ -1210,9 +1238,11 @@ fixCopyCommand(Archive *fout, PQExpBuffer copyBuf, TableInfo *tbinfo, bool isFro
pg_strcasecmp(bbf_db_name, "msdb") == 0)
? true : false;

/* Remember if it is babelfish_authid_user_ext catalog table. */
/* Remember if it is babelfish_authid_user_ext and babelfish_sysdatabases catalog table. */
if (strcmp(tbinfo->dobj.name, "babelfish_authid_user_ext") == 0)
is_bbf_usr_ext_tab = true;
if (strcmp(tbinfo->dobj.name, "babelfish_sysdatabases") == 0)
is_bbf_sysdatabases_tab = true;

q = createPQExpBuffer();
for (i = 0; i < tbinfo->numatts; i++)
Expand All @@ -1228,6 +1258,13 @@ fixCopyCommand(Archive *fout, PQExpBuffer copyBuf, TableInfo *tbinfo, bool isFro
*/
if (bbf_db_name != NULL && !is_builtin_db && strcmp(tbinfo->attnames[i], "dbid") == 0)
continue;
/*
* We need to skip owner column of babelfish_sysdatabases table as it might be
* referencing Babelfish initialize user which we do not include in dump. We will
* populate this column during restore with the initialize user of target database.
*/
else if (is_bbf_sysdatabases_tab && strcmp(tbinfo->attnames[i], "owner") == 0)
continue;
if (needComma)
appendPQExpBufferStr(q, ", ");

Expand Down Expand Up @@ -1278,19 +1315,13 @@ fixCopyCommand(Archive *fout, PQExpBuffer copyBuf, TableInfo *tbinfo, bool isFro
* Returns true if table in Babelfish Database is to be dumped with INSERT mode.
* Currently we dump tables with sql_variant columns with INSERT operations to
* correctly restore the metadata of the base datatype, which is not directly
* posible with COPY statements. We also dump sys.babelfish_authid_login_ext
* with INSERT statements so that if target database already has a login with
* same name as in the source database, only that INSERT query with fail and won't
* affect the other entries of the catalog table.
* posible with COPY statements.
*/
bool bbfIsDumpWithInsert(Archive *fout, TableInfo *tbinfo)
bool
bbfIsDumpWithInsert(Archive *fout, TableInfo *tbinfo)
{
return (isBabelfishDatabase(fout) &&
(hasSqlvariantColumn(tbinfo) ||
pg_strcasecmp(fmtQualifiedDumpable(tbinfo),
quote_all_identifiers ?
"\"sys\".\"babelfish_authid_login_ext\"" :
"sys.babelfish_authid_login_ext") == 0));
hasSqlvariantColumn(tbinfo));
}

/*
Expand Down
52 changes: 47 additions & 5 deletions src/bin/pg_dump/dumpall_babel_utils.c
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,30 @@ executeQuery(PGconn *conn, const char *query)
return res;
}

/*
* getBabelfishInitUser
* Returns initialize user of currentl Babelfish database
* which is essentially same as owner of Babelfish database.
*/
static char *
getBabelfishInitUser(PGconn *conn)
{
PQExpBuffer qry;
PGresult *res;
char *babel_init_user;

qry = createPQExpBuffer();
appendPQExpBufferStr(qry, "SELECT r.rolname FROM pg_roles r "
"INNER JOIN pg_database d ON r.oid = d.datdba "
"WHERE d.datname = current_database()");
res = executeQuery(conn, qry->data);
babel_init_user = pstrdup(PQgetvalue(res, 0, 0));
PQclear(res);
destroyPQExpBuffer(qry);

return babel_init_user;
}

/*
* isBabelfishDatabase:
* returns true if current database has "babelfishpg_tsql"
Expand Down Expand Up @@ -172,10 +196,21 @@ getBabelfishRolesQuery(PGconn *conn, PQExpBuffer buf, char *role_catalog,

resetPQExpBuffer(buf);
appendPQExpBufferStr(buf, "WITH bbf_catalog AS (");
/* Include logins only in case of Babelfish physical database dump. */
/*
* Include logins only in case of Babelfish physical database dump.
* Note that we will not dump Babelfish initialize user as it might
* already be present on the target server.
*/
if (!bbf_db_name)
appendPQExpBufferStr(buf,
"SELECT rolname FROM sys.babelfish_authid_login_ext UNION ");
{
char *babel_init_user = getBabelfishInitUser(conn);
appendPQExpBuffer(buf,
"SELECT rolname FROM sys.babelfish_authid_login_ext "
"WHERE rolname != '%s' " /* Do not dump Babelfish initialize user */
"UNION ",
babel_init_user);
pfree(babel_init_user);
}
appendPQExpBufferStr(buf,
"SELECT rolname FROM sys.babelfish_authid_user_ext ");
/* Only dump users of the specific logical database we are currently dumping. */
Expand Down Expand Up @@ -241,8 +276,15 @@ getBabelfishRoleMembershipQuery(PGconn *conn, PQExpBuffer buf,
appendPQExpBufferStr(buf, "WITH bbf_catalog AS (");
/* Include logins only in case of Babelfish physical database dump. */
if (!bbf_db_name)
appendPQExpBufferStr(buf,
"SELECT rolname FROM sys.babelfish_authid_login_ext UNION ");
{
char *babel_init_user = getBabelfishInitUser(conn);
appendPQExpBuffer(buf,
"SELECT rolname FROM sys.babelfish_authid_login_ext "
"WHERE rolname != '%s' " /* Do not dump Babelfish initialize user */
"UNION ",
babel_init_user);
pfree(babel_init_user);
}
appendPQExpBuffer(buf,
"SELECT rolname FROM sys.babelfish_authid_user_ext ");
/* Only dump users of the specific logical database we are currently dumping. */
Expand Down
7 changes: 3 additions & 4 deletions src/bin/pg_dump/pg_dump.c
Original file line number Diff line number Diff line change
Expand Up @@ -2149,8 +2149,7 @@ dumpTableData_insert(Archive *fout, const void *dcontext)
int rows_this_statement = 0;

/*
* For tables in Babelfish Database with sql_variant datatype columns and
* sys.babelfish_authid_login_ext Babelfish catalog table, we want to
* For tables in Babelfish Database with sql_variant datatype columns, we want to
* surpass dopt->column_inserts check since these tables need to be dumped
* as when column_inserts is true.
*/
Expand Down Expand Up @@ -2544,8 +2543,8 @@ dumpTableData(Archive *fout, const TableDataInfo *tdinfo)
if (bbfIsDumpWithInsert(fout, tbinfo))
{
/*
* dump tables in Babelfish Database with sql_variant datatype columns and
* sys.babelfish_authid_login_ext Babelfish catalog table using INSERT only
* dump tables in Babelfish Database with sql_variant
* datatype columns using INSERT only.
*/
dumpFn = dumpTableData_insert;
copyStmt = NULL;
Expand Down

0 comments on commit 9bbf95f

Please sign in to comment.