[#20972] YSQL, Backups: Implement new ROLEs related flags

Summary:
The following **NEW** `yb_backup.py` flags were implemented:
* `--backup_roles`
* `--restore_roles`
* `--ignore_existing_roles`
* `--use_roles`
2 last flags are passed into `ysqlsh -v ignore_existing_roles=true -v use_roles=true ....`.

**SEMANTICS**:
Backup-create: "--backup_roles" ~ backup Roles
Backup-restore: "--restore_roles" ~ restore the stored Roles
Backup-restore: "--ignore_existing_roles" ~ do not create the role if it already exists
Backup-restore: "--use_roles" ~ use the Roles for the tables via
`ALTER TABLE ... OWNER TO <role>` and
`REVOKE/GRANT ... ON TABLE ... FROM/TO <role>` syntax.

NOTE: The API is changed only when `ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR`
is set into `True` (now it's `False`).

In the YSQL dump the new logic is implemented:
```
\if :{?ignore_existing_roles}
\else
\set ignore_existing_roles false
\endif

\set role_exists false
\if :ignore_existing_roles
    SELECT EXISTS(SELECT 1 FROM pg_roles WHERE rolname = 'role_admin') AS role_exists \gset
\endif
\if :role_exists
    \echo 'Role role_admin already exists.'
\else
    CREATE ROLE role_admin;
\endif
```

```
\if :{?use_roles}
\else
\set use_roles true
\endif

CREATE TABLE ...

\if :use_roles
ALTER TABLE t OWNER TO role_admin;
\endif

\if :use_roles
GRANT SELECT ON TABLE t TO role_admin;
\endif
```
Jira: DB-9947

Test Plan:
Set ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR=True in yb_backup.py.

New tests:
ybd --java-test org.yb.pgsql.TestYbBackup#testBackupRestoreRoles
ybd --java-test org.yb.pgsql.TestYbBackup#testBackupRolesWithoutUseRoles
ybd --java-test org.yb.pgsql.TestYbBackup#testBackupRolesWithoutRestoreRoles

Other related tests:
ybd --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpWithYbMetadata
ybd --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpWithoutYbMetadata
ybd --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpAllWithoutYbMetadata
ybd --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpAllWithYbMetadata
ybd --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpLegacyColocatedDB
ybd --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpColocatedDB

ybd --java-test org.yb.pgsql.TestYbBackup#testExtensionBackupUsingTestExtension
ybd --java-test org.yb.pgsql.TestYbBackup#testBackupRestoreTablespaces
ybd --java-test org.yb.pgsql.TestYbBackup#testFailureOnExistingTablespaces
ybd --java-test org.yb.pgsql.TestYbBackup#testIgnoreExistingTablespaces

ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioning
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningNoRegions
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningOneRegion
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningRestoringIntoExisting
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningWithTablespaces
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningNoRegionsWithTablespaces
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningOneRegionWithTablespaces
ybd --java-test org.yb.pgsql.TestYbBackup#testGeoPartitioningRestoringIntoExistingWithTablespaces

Reviewers: mihnea, tverona, myang, yguan, vkumar, yng

Reviewed By: myang

Subscribers: yugaware, yql

Differential Revision: https://phorge.dev.yugabyte.com/D32252

java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbBackup.java

@@ -56,6 +56,7 @@
 import org.yb.util.YBTestRunnerNonTsanAsan;
 
 import com.google.common.collect.ImmutableMap;
+import com.yugabyte.util.PSQLException;
 
 import static org.yb.AssertionWrappers.assertArrayEquals;
 import static org.yb.AssertionWrappers.assertEquals;
@@ -69,7 +70,15 @@
 public class TestYbBackup extends BasePgSQLTest {
   private static final Logger LOG = LoggerFactory.getLogger(TestYbBackup.class);
 
-  // This value must be synced with the same variable in `yb_backup.py` script.
+  private static final String PERMISSION_DENIED = "permission denied";
+
+  // This constant must be synced with the same variable in `yb_backup.py` script.
+  // This temporary constant enables the features:
+  // 1. "STOP_ON_ERROR" mode in 'ysqlsh' tool
+  // 2. New API: 'backup_tablespaces'/'restore_tablespaces'+'use_tablespaces'
+  //    instead of old 'use_tablespaces'
+  // 3. If the new API 'backup_roles' is NOT used - the YSQL Dump is generated
+  //    with '--no-privileges' flag.
   private static final boolean ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR = false;
 
   @Before
@@ -1445,6 +1454,97 @@ public void testGeoPartitioningDeleteBackup() throws Exception {
     }
   }
 
+  public void doTestBackupRestoreRoles(boolean restoreRoles, boolean useRoles)
+      throws Exception {
+    String backupDir = null;
+    try (Statement stmt = connection.createStatement()) {
+      stmt.execute("CREATE TABLE test_table(id INT PRIMARY KEY)");
+      stmt.execute("INSERT INTO test_table (id) VALUES (1)");
+
+      stmt.execute("CREATE ROLE admin LOGIN NOINHERIT");
+      stmt.execute("REVOKE ALL ON TABLE test_table FROM admin");
+      stmt.execute("GRANT SELECT ON TABLE test_table TO admin");
+
+      backupDir = YBBackupUtil.getTempBackupDir();
+      String output = YBBackupUtil.runYbBackupCreate("--backup_location", backupDir,
+          "--keyspace", "ysql.yugabyte", "--backup_roles");
+      backupDir = new JSONObject(output).getString("snapshot_url");
+    }
+
+    try (Connection connection2 = getConnectionBuilder().withUser("admin").connect();
+         Statement stmt = connection2.createStatement()) {
+      assertQuery(stmt, "SELECT * FROM test_table WHERE id=1", new Row(1));
+
+      runInvalidQuery(stmt, "INSERT INTO test_table (id) VALUES (9)", PERMISSION_DENIED);
+    }
+
+    try (Statement stmt = connection.createStatement()) {
+      stmt.execute("REVOKE ALL ON TABLE test_table FROM admin");
+      stmt.execute("DROP ROLE admin");
+    }
+
+    List<String> args = new ArrayList<>(Arrays.asList(
+        "--keyspace", "ysql.yb2", "--ignore_existing_roles"));
+    if (restoreRoles) {
+      args.add("--restore_roles");
+    }
+    if (useRoles) {
+      args.add("--use_roles");
+    }
+
+    try {
+      YBBackupUtil.runYbBackupRestore(backupDir, args);
+    } catch (YBBackupException ex) {
+      if (ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR && !restoreRoles) {
+        LOG.info("Expected exception", ex);
+        assertTrue(ex.getMessage().contains("ERROR:  role \"admin\" does not exist"));
+        return;
+      } else {
+        throw ex;
+      }
+    }
+
+    try (Connection connection3 = getConnectionBuilder().withDatabase("yb2").connect();
+         Statement stmt = connection3.createStatement()) {
+      assertQuery(stmt, "SELECT * FROM test_table WHERE id=1", new Row(1));
+    }
+
+    try (Connection connection4 =
+             getConnectionBuilder().withDatabase("yb2").withUser("admin").connect();
+         Statement stmt = connection4.createStatement()) {
+      assertQuery(stmt, "SELECT * FROM test_table WHERE id=1", new Row(1));
+
+      runInvalidQuery(stmt, "INSERT INTO test_table (id) VALUES (9)", PERMISSION_DENIED);
+    } catch (PSQLException ex) {
+      if (restoreRoles) {
+        throw ex;
+      } else {
+        LOG.info("Expected exception", ex);
+        assertTrue(ex.getMessage().contains("FATAL: role \"admin\" does not exist"));
+     }
+    }
+
+    // Cleanup.
+    try (Statement stmt = connection.createStatement()) {
+      stmt.execute("DROP DATABASE yb2");
+    }
+  }
+
+  @Test
+  public void testBackupRestoreRoles() throws Exception {
+    doTestBackupRestoreRoles(/* restoreRoles */ true, /* useRoles */ true);
+  }
+
+  @Test
+  public void testBackupRolesWithoutUseRoles() throws Exception {
+    doTestBackupRestoreRoles(/* restoreRoles */ true, /* useRoles */ false);
+  }
+
+  @Test
+  public void testBackupRolesWithoutRestoreRoles() throws Exception {
+    doTestBackupRestoreRoles(/* restoreRoles */ false, /* useRoles */ false);
+  }
+
   @Test
   public void testUserDefinedTypes() throws Exception {
     // TODO(myang): Add ALTER TYPE test after #1893 is fixed.

managed/devops/bin/yb_backup.py

@@ -72,6 +72,8 @@
 # 1. "STOP_ON_ERROR" mode in 'ysqlsh' tool
 # 2. New API: 'backup_tablespaces'/'restore_tablespaces'+'use_tablespaces'
 #    instead of old 'use_tablespaces'
+# 3. If the new API 'backup_roles' is NOT used - the YSQL Dump is generated
+#    with '--no-privileges' flag.
 ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR = False
 
 ROCKSDB_PATH_PREFIX = '/yb-data/tserver/data/rocksdb'
@@ -88,6 +90,7 @@
 MANIFEST_FILE_NAME = 'Manifest'
 METADATA_FILE_NAME = 'SnapshotInfoPB'
 SQL_TBSP_DUMP_FILE_NAME = 'YSQLDump_tablespaces'
+SQL_ROLES_DUMP_FILE_NAME = 'YSQLDump_roles'
 SQL_DUMP_FILE_NAME = 'YSQLDump'
 SQL_DATA_DUMP_FILE_NAME = 'YSQLDump_data'
 CREATE_METAFILES_MAX_RETRIES = 10
@@ -992,6 +995,7 @@ def init(self, snapshot_bucket, pg_based_backup):
         if not ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR:
             properties['use-tablespaces'] = self.backup.args.use_tablespaces
         properties['backup-tablespaces'] = self.backup.args.backup_tablespaces
+        properties['backup-roles'] = self.backup.args.backup_roles
         properties['pg-based-backup'] = pg_based_backup
         properties['start-time'] = self.backup.timer.start_time_str()
         properties['end-time'] = self.backup.timer.end_time_str()
@@ -1343,12 +1347,27 @@ def parse_arguments(self):
             help='Restore YSQL TABLESPACE objects from the backup.')
         parser.add_argument(
             '--ignore_existing_tablespaces', required=False, action='store_true', default=False,
-            help='Whether to skip tablespace creation if tablespace already exists.')
+            help='Ignore the tablespace creation if it already exists.')
         parser.add_argument(
             '--use_tablespaces', required=False, action='store_true', default=False,
             help="Use the restored YSQL TABLESPACE objects for the tables "
                  "via 'SET default_tablespace=...' syntax.")
 
+        parser.add_argument(
+            '--backup_roles', required=False, action='store_true', default=False,
+            help='Backup YSQL ROLE objects into the backup.')
+        parser.add_argument(
+            '--restore_roles', required=False, action='store_true', default=False,
+            help='Restore YSQL ROLE objects from the backup.')
+        parser.add_argument(
+            '--ignore_existing_roles', required=False, action='store_true', default=False,
+            help='Ignore the role creation if it already exists.')
+        parser.add_argument(
+            '--use_roles', required=False, action='store_true', default=False,
+            help="Use the restored YSQL ROLE objects for the tables "
+                 "via 'ALTER TABLE ... OWNER TO <role>' and "
+                 "'REVOKE/GRANT ... ON TABLE ... FROM/TO <role>' syntax.")
+
         parser.add_argument('--upload', dest='upload', action='store_true', default=True)
         # Please note that we have to use this weird naming (i.e. underscore in the argument name)
         # style to keep it in sync with YB processes G-flags.
@@ -3051,6 +3070,9 @@ def create_metadata_files(self):
             if self.args.use_tablespaces:
                 raise BackupException("--use_tablespaces can be used in RESTORE mode only")
 
+        if self.args.use_roles:
+            raise BackupException("--use_roles can be used in RESTORE mode only")
+
         snapshot_id = None
         dump_files = []
         pg_based_backup = self.args.pg_based_backup
@@ -3075,6 +3097,17 @@ def create_metadata_files(self):
             else:
                 ysql_dump_args.append('--no-tablespaces')
 
+            sql_roles_dump_path = os.path.join(self.get_tmp_dir(), SQL_ROLES_DUMP_FILE_NAME)\
+                if self.args.backup_roles else None
+            if sql_roles_dump_path:
+                logging.info("[app] Creating ysql dump for roles to {}".format(
+                    sql_roles_dump_path))
+                self.run_ysql_dumpall(['--roles-only', '--include-yb-metadata',
+                                       '--file=' + sql_roles_dump_path])
+                dump_files.append(sql_roles_dump_path)
+            elif ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR:
+                ysql_dump_args.append('--no-privileges')
+
             logging.info("[app] Creating ysql dump for DB '{}' to {}".format(
                 db_name, sql_dump_path))
             self.run_ysql_dump(ysql_dump_args)
@@ -3439,6 +3472,15 @@ def download_metadata_file(self):
         else:
             sql_tbsp_dump_path = None
 
+        if self.args.restore_roles:
+            src_sql_roles_dump_path = os.path.join(
+                self.args.backup_location, SQL_ROLES_DUMP_FILE_NAME)
+            sql_roles_dump_path = os.path.join(self.get_tmp_dir(), SQL_ROLES_DUMP_FILE_NAME)
+            self.download_file(src_sql_roles_dump_path, sql_roles_dump_path)
+            dump_files.append(sql_roles_dump_path)
+        else:
+            sql_roles_dump_path = None
+
         src_sql_dump_path = os.path.join(self.args.backup_location, SQL_DUMP_FILE_NAME)
         sql_dump_path = os.path.join(self.get_tmp_dir(), SQL_DUMP_FILE_NAME)
         try:
@@ -3532,12 +3574,15 @@ def import_ysql_dump(self, dump_file_path):
         if self.args.ignore_existing_tablespaces:
             ysql_shell_args.append('--variable=ignore_existing_tablespaces=1')
 
-        if ENABLE_STOP_ON_YSQL_DUMP_RESTORE_ERROR:
-            if self.args.use_tablespaces:
-                ysql_shell_args.append('--variable=use_tablespaces=1')
-        else:  # Always enabled in old semantics
+        if self.args.use_tablespaces:
             ysql_shell_args.append('--variable=use_tablespaces=1')
 
+        if self.args.ignore_existing_roles:
+            ysql_shell_args.append('--variable=ignore_existing_roles=1')
+
+        if self.args.use_roles:
+            ysql_shell_args.append('--variable=use_roles=1')
+
         if on_error_stop:
             logging.info(
                 "[app] Enable ON_ERROR_STOP mode when applying '{}'".format(dump_file_path))
@@ -3812,7 +3857,10 @@ def restore_table(self):
 
         for dump_file_path in dump_file_paths:
             dump_file = os.path.basename(dump_file_path)
-            if dump_file == SQL_TBSP_DUMP_FILE_NAME:
+            if dump_file == SQL_ROLES_DUMP_FILE_NAME:
+                logging.info('[app] Create roles from {}'.format(dump_file_path))
+                self.import_ysql_dump(dump_file_path)
+            elif dump_file == SQL_TBSP_DUMP_FILE_NAME:
                 logging.info('[app] Create tablespaces from {}'.format(dump_file_path))
                 self.import_ysql_dump(dump_file_path)
             elif dump_file == SQL_DUMP_FILE_NAME:

src/postgres/src/bin/pg_dump/pg_backup_archiver.c

@@ -3159,7 +3159,12 @@ _doSetFixedOutputState(ArchiveHandle *AH)
 		ahprintf(AH, "\n-- Set variable use_tablespaces (if not already set)\n"
 					 "\\if :{?use_tablespaces}\n"
 					 "\\else\n"
-					 "\\set use_tablespaces false\n"
+					 "\\set use_tablespaces true\n"
+					 "\\endif\n");
+		ahprintf(AH, "\n-- Set variable use_roles (if not already set)\n"
+					 "\\if :{?use_roles}\n"
+					 "\\else\n"
+					 "\\set use_roles true\n"
 					 "\\endif\n");
 	}
 
@@ -3689,7 +3694,14 @@ _printTocEntry(ArchiveHandle *AH, TocEntry *te, bool isData)
 			appendPQExpBufferStr(temp, "ALTER ");
 			_getObjectDescription(temp, te, AH);
 			appendPQExpBuffer(temp, " OWNER TO %s;", fmtId(te->owner));
-			ahprintf(AH, "%s\n\n", temp->data);
+
+			if (AH->public.dopt->include_yb_metadata)
+				ahprintf(AH, "\\if :use_roles\n"
+							 "    %s\n"
+							 "\\endif\n\n", temp->data);
+			else
+				ahprintf(AH, "%s\n\n", temp->data);
+
 			destroyPQExpBuffer(temp);
 		}
 		else if (strcmp(te->desc, "CAST") == 0 ||

src/postgres/src/bin/pg_dump/pg_dump.c

@@ -15255,20 +15255,30 @@ dumpACL(Archive *fout, CatalogId objCatId, DumpId objDumpId,
 	if (sql->len > 0)
 	{
 		PQExpBuffer tag = createPQExpBuffer();
+		PQExpBuffer use_roles_sql;
 
 		if (subname)
 			appendPQExpBuffer(tag, "COLUMN %s.%s", name, subname);
 		else
 			appendPQExpBuffer(tag, "%s %s", type, name);
 
+		if (dopt->include_yb_metadata)
+		{
+			use_roles_sql = createPQExpBuffer();
+			appendPQExpBuffer(use_roles_sql, "\\if :use_roles\n%s\\endif\n", sql->data);
+		}
+
 		ArchiveEntry(fout, nilCatalogId, createDumpId(),
 					 tag->data, nspname,
 					 NULL,
 					 owner ? owner : "",
 					 false, "ACL", SECTION_NONE,
-					 sql->data, "", NULL,
-					 &(objDumpId), 1,
+					 dopt->include_yb_metadata ? use_roles_sql->data : sql->data,
+					 "", NULL, &(objDumpId), 1,
 					 NULL, NULL);
+		if (dopt->include_yb_metadata)
+			destroyPQExpBuffer(use_roles_sql);
+
 		destroyPQExpBuffer(tag);
 	}
 

src/postgres/src/bin/pg_dump/pg_dumpall.c

@@ -853,11 +853,21 @@ dumpRoles(PGconn *conn)
 	i_is_current_user = PQfnumber(res, "is_current_user");
 
 	if (PQntuples(res) > 0)
+	{
 		fprintf(OPF, "--\n-- Roles\n--\n\n");
 
+		if (include_yb_metadata)
+			fprintf(OPF, "-- Set variable ignore_existing_roles (if not already set)\n"
+						 "\\if :{?ignore_existing_roles}\n"
+						 "\\else\n"
+						 "\\set ignore_existing_roles false\n"
+						 "\\endif\n\n");
+	}
+
 	for (i = 0; i < PQntuples(res); i++)
 	{
 		const char *rolename;
+		char	   *frolename;
 		Oid			auth_oid;
 
 		auth_oid = atooid(PQgetvalue(res, i, i_oid));
@@ -880,6 +890,7 @@ dumpRoles(PGconn *conn)
 							  auth_oid);
 		}
 
+		frolename = pg_strdup(fmtId(rolename));
 		/*
 		 * We dump CREATE ROLE followed by ALTER ROLE to ensure that the role
 		 * will acquire the right properties even if it already exists (ie, it
@@ -890,8 +901,25 @@ dumpRoles(PGconn *conn)
 		 */
 		if (!binary_upgrade ||
 			strcmp(PQgetvalue(res, i, i_is_current_user), "f") == 0)
-			appendPQExpBuffer(buf, "CREATE ROLE %s;\n", fmtId(rolename));
-		appendPQExpBuffer(buf, "ALTER ROLE %s WITH", fmtId(rolename));
+		{
+			if (include_yb_metadata)
+				appendPQExpBuffer(buf,
+					"\\set role_exists false\n"
+					"\\if :ignore_existing_roles\n"
+					"    SELECT EXISTS(SELECT 1 FROM pg_roles WHERE rolname = '%s')"
+					" AS role_exists \\gset\n"
+					"\\endif\n"
+					"\\if :role_exists\n"
+					"    \\echo 'Role %s already exists.'\n"
+					"\\else\n    ", frolename, frolename);
+
+			appendPQExpBuffer(buf, "CREATE ROLE %s;\n", frolename);
+
+			if (include_yb_metadata)
+				appendPQExpBufferStr(buf, "\\endif\n");
+		}
+
+		appendPQExpBuffer(buf, "ALTER ROLE %s WITH", frolename);
 
 		if (strcmp(PQgetvalue(res, i, i_rolsuper), "t") == 0)
 			appendPQExpBufferStr(buf, " SUPERUSER");
@@ -947,7 +975,7 @@ dumpRoles(PGconn *conn)
 
 		if (!no_comments && !PQgetisnull(res, i, i_rolcomment))
 		{
-			appendPQExpBuffer(buf, "COMMENT ON ROLE %s IS ", fmtId(rolename));
+			appendPQExpBuffer(buf, "COMMENT ON ROLE %s IS ", frolename);
 			appendStringLiteralConn(buf, PQgetvalue(res, i, i_rolcomment), conn);
 			appendPQExpBufferStr(buf, ";\n");
 		}
@@ -957,7 +985,11 @@ dumpRoles(PGconn *conn)
 							 "ROLE", rolename,
 							 buf);
 
+		if (include_yb_metadata)
+			appendPQExpBufferStr(buf, "\n");
+
 		fprintf(OPF, "%s", buf->data);
+		free(frolename);
 	}
 
 	/*