unrestrict psql meta-commands for using backslash command for pg_dump and pg_dumpall in Babelfish

Signed-off-by: yashneet vinayak <yashneet@amazon.com>

Author: Yashneet Vinayak

src/bin/pg_dump/dump_babel_utils.c

@@ -270,7 +270,11 @@ dumpBabelRestoreChecks(Archive *fout)
 	 * Note that it can only be used in plain text dump (archNull).
 	 */
 	if (format == archNull)
+	{
+		appendPQExpBuffer(qry, "\\unrestrict %s\n", fout->dopt->restrict_key);
 		appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP on\n\n");
+		appendPQExpBuffer(qry, "\\restrict %s\n", fout->dopt->restrict_key);
+	}
 	appendPQExpBuffer(qry,
 					  "DO $$"
 					  "\nDECLARE"
@@ -307,7 +311,11 @@ dumpBabelRestoreChecks(Archive *fout)
 					  "\nEND$$;\n\n"
 					  , source_migration_mode);
 	if (format == archNull)
+	{
+		appendPQExpBuffer(qry, "\\unrestrict %s\n", fout->dopt->restrict_key);
 		appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP off\n");
+		appendPQExpBuffer(qry, "\\restrict %s\n", fout->dopt->restrict_key);
+	}
 	PQclear(res);
 
 	ArchiveEntry(fout, nilCatalogId, createDumpId(),

src/bin/pg_dump/dumpall_babel_utils.c

@@ -123,7 +123,7 @@ isBabelfishDatabase(PGconn *conn)
  * or not.
  */
 void
-dumpBabelRestoreChecks(FILE *OPF, PGconn *conn, int binary_upgrade)
+dumpBabelRestoreChecks(FILE *OPF, PGconn *conn, int binary_upgrade, char *restrict_key)
 {
 	PGresult	*res;
 	int     	source_server_version_num;
@@ -151,7 +151,9 @@ dumpBabelRestoreChecks(FILE *OPF, PGconn *conn, int binary_upgrade)
 	 * Temporarily enable ON_ERROR_STOP so that whole restore script
 	 * execution fails if the following do block raises an error.
 	 */
+	appendPQExpBuffer(qry, "\\unrestrict %s\n", restrict_key);
 	appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP on\n\n");
+	appendPQExpBuffer(qry, "\\restrict %s\n", restrict_key);
 	appendPQExpBuffer(qry,
 					  "DO $$"
 					  "\nDECLARE"
@@ -187,7 +189,10 @@ dumpBabelRestoreChecks(FILE *OPF, PGconn *conn, int binary_upgrade)
 					  "\n    END IF;"
 					  "\nEND$$;\n\n"
 					  , source_migration_mode);
+
+	appendPQExpBuffer(qry, "\\unrestrict %s\n", restrict_key);
 	appendPQExpBufferStr(qry, "\\set ON_ERROR_STOP off\n");
+	appendPQExpBuffer(qry, "\\restrict %s\n", restrict_key);
 	PQclear(res);
 
 	fprintf(OPF, "%s", qry->data);

src/bin/pg_dump/dumpall_babel_utils.h

@@ -22,6 +22,6 @@ extern void getBabelfishRolesQuery(PGconn *conn, PQExpBuffer buf,
 extern void getBabelfishRoleMembershipQuery(PGconn *conn, PQExpBuffer buf,
         char *role_catalog, int binary_upgrade);
 extern bool isBabelfishDatabase(PGconn *conn);
-extern void dumpBabelRestoreChecks(FILE *OPF, PGconn *conn, int binary_upgrade);
+extern void dumpBabelRestoreChecks(FILE *OPF, PGconn *conn, int binary_upgrade, char *restrict_key);
 
 #endif

src/bin/pg_dump/pg_dumpall.c

@@ -623,7 +623,7 @@ main(int argc, char *argv[])
 		fprintf(OPF, "SET escape_string_warning = off;\n");
 	fprintf(OPF, "\n");
 
-	dumpBabelRestoreChecks(OPF, conn, binary_upgrade);
+	dumpBabelRestoreChecks(OPF, conn, binary_upgrade, restrict_key);
 
 	if (!data_only && !statistics_only && !no_schema)
 	{