[BACKPORT 2024.2][#23312] YSQL: Fix backup/restore for colocated tabl…

…es without tablespace information.

Summary:
Backport Description:
No merge conflicts

Original Description:
Original commit: 0dbe7d6 / D36859

Previously, backup and restore operations for colocated tables failed when multiple tablespaces were present, and the `--use_tablespaces` option was not used. The issue arose because each tablespace could only contain one tablegroup, leading to conflicts during the restore process. Without tablespace information, we could only restore tables to a single tablegroup (corresponding to the default tablespace). However, since a tablegroup represents a (parent) colocated tablet, restoring multiple colocated tablets required the same number of tablegroups.

The revision resolves this issue by enabling multiple tablegroups in the default tablespace during restoration. This is achieved through the following steps:

 * Originally, the tablegroup’s name is `colocation_<tablespace_oid>`, but since the concept of tablespaces doesn’t apply after restoration, we rename the tablegroup as `colocation_restore_<tablegroup_oid>`. The default tablegroup's name remains unchanged after the restore process.
 * To retain the original name of the default tablegroup, we track which tables are associated with it during the backup. This allows us to correctly restore the tablegroup's properties.

**Implementation**

This update enhances the backup and restore functionality for colocated tables with tablespaces by addressing the issue as follows:

 - **Backup Phase**: The backup process now tags each table as either belonging to the default tablegroup or not, along with the tablegroup ID in YSQLDump. The tablegroup name, originally formatted as `colocation_<tablespace_oid>`, will be renamed during restoration as described above.

 - **Restore Phase**: During restoration, the tablegroup name and its association with the default tablegroup are retrieved from YSQLDump. If the corresponding tablegroup doesn’t already exist, it is created in the default tablespace, allowing multiple tablegroups to coexist without conflicts.
 - **ALTER Table**: To move a tablegroup to another tablespace after restoration, the syntax:
```ALTER TABLE ALL IN TABLESPACE pg_default COLOCATED WITH <table_name> SET TABLESPACE <new_tablespace_name> CASCADE;```
should be used, ensuring proper management of tablegroups and tablespaces.

DB-12237

Test Plan:
== Automated ==

./yb_build.sh  --cxx-test yb-backup-cross-feature-test --gtest_filter YBBackupTestColocatedTablesWithTablespaces.TestBackupColocatedTablesWithTablespaces
./yb_build.sh  --cxx-test yb-backup-cross-feature-test --gtest_filter YBBackupTestColocatedTablesWithTablespaces.TestBackupColocatedTablesWithoutTablespaces
./yb_build.sh --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpColocatedTablesWithTablespaces

 1. **YBBackupTestColocatedTablesWithTablespaces.TestBackupColocatedTablesWithoutTablespaces**: This test case increases the number of tablespaces (and consequently the number of tablegroups) to replicate the scenario where multiple tablespaces exist. It validates that backup and restore work correctly when performed without the `--use_tablespaces` flag. This is the primary case that this revision addresses.

 2. **YBBackupTestColocatedTablesWithTablespaces.TestBackupColocatedTablesWithTablespaces**: To ensure consistency, this test also increases the number of tablespaces but verifies that backup and restore function correctly when using the `--use_tablespaces` flag.

 3. **org.yb.pgsql.TestYsqlDump#ysqlDumpColocatedTablesWithTablespaces**: This test checks that the YsqlDump output matches the expected format. Specifically, in the case of the fix, the dump file should include the line:
```SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);```

Reviewers: skumar, aagrawal, yguan

Reviewed By: skumar

Subscribers: ybase, yql, mlillibridge, svc_phabricator

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

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

@@ -45,6 +45,7 @@ public class TestYsqlDump extends BasePgSQLTest {
   private static final Logger LOG = LoggerFactory.getLogger(TestYsqlDump.class);
 
   private static enum IncludeYbMetadata { ON, OFF }
+  private static enum NoTableSpaces { ON, OFF }
 
   @Override
   public int getTestMethodTimeoutSec() {
@@ -96,7 +97,8 @@ public void ysqlDumpWithYbMetadata() throws Exception {
         "expected/yb_ysql_dump_describe.out" /* expectedDescribeFileRelativePath */,
         "results/yb_ysql_dump.out" /* outputFileRelativePath */,
         "results/yb_ysql_dump_describe.out" /* outputDescribeFileRelativePath */,
-        IncludeYbMetadata.ON);
+        IncludeYbMetadata.ON,
+        NoTableSpaces.OFF);
   }
 
   @Test
@@ -111,7 +113,8 @@ public void ysqlDumpAllWithYbMetadata() throws Exception {
         "expected/yb_ysql_dumpall_describe.out" /* expectedDescribeFileRelativePath */,
         "results/yb_ysql_dumpall.out" /* outputFileRelativePath */,
         "results/yb_ysql_dumpall_describe.out" /* outputDescribeFileRelativePath */,
-        IncludeYbMetadata.ON);
+        IncludeYbMetadata.ON,
+        NoTableSpaces.OFF);
   }
 
   @Test
@@ -127,7 +130,8 @@ public void ysqlDumpWithoutYbMetadata() throws Exception {
         "results/yb_ysql_dump_without_ybmetadata.out" /* outputFileRelativePath */,
         "results/yb_ysql_dump_without_ybmetadata_describe.out"
         /* outputDescribeFileRelativePath */,
-        IncludeYbMetadata.OFF);
+        IncludeYbMetadata.OFF,
+        NoTableSpaces.OFF);
   }
 
   @Test
@@ -143,7 +147,8 @@ public void ysqlDumpAllWithoutYbMetadata() throws Exception {
         "results/yb_ysql_dumpall_without_ybmetadata.out" /* outputFileRelativePath */,
         "results/yb_ysql_dumpall_without_ybmetadata_describe.out"
         /* outputDescribeFileRelativePath */,
-        IncludeYbMetadata.OFF);
+        IncludeYbMetadata.OFF,
+        NoTableSpaces.OFF);
   }
 
   @Test
@@ -158,7 +163,42 @@ public void ysqlDumpColocatedDB() throws Exception {
         /* expectedDescribeFileRelativePath */,
         "results/yb_ysql_dump_colocated_database.out" /* outputFileRelativePath */,
         "results/yb_ysql_dump_describe_colocated_database.out" /* outputDescribeFileRelativePath */,
-        IncludeYbMetadata.ON);
+        IncludeYbMetadata.ON,
+        NoTableSpaces.OFF);
+  }
+
+  @Test
+  public void ysqlDumpColocatedTablesWithTablespaces() throws Exception {
+    markClusterNeedsRecreation();
+    restartClusterWithClusterBuilder(cb -> {
+      cb.addCommonFlag(
+          "allowed_preview_flags_csv", "ysql_enable_colocated_tables_with_tablespaces");
+      cb.addCommonFlag("ysql_enable_colocated_tables_with_tablespaces", "true");
+      cb.addCommonTServerFlag("placement_cloud", "testCloud");
+      cb.addCommonTServerFlag("placement_region", "testRegion");
+      cb.perTServerFlags(Arrays.asList(
+          Collections.singletonMap("placement_zone", "testZone1"),
+          Collections.singletonMap("placement_zone", "testZone2"),
+          Collections.singletonMap("placement_zone", "testZone3")));
+    });
+    LOG.info("created mini cluster");
+    ysqlDumpTester(
+        "ysql_dump" /* binaryName */,
+        "colo_tables" /* dumpedDatabaseName */,
+        "sql/yb_ysql_dump_colocated_tables_with_tablespaces.sql"
+        /* inputFileRelativePath */,
+        "sql/yb_ysql_dump_describe_colocated_tables_with_tablespaces.sql"
+        /* inputDescribeFileRelativePath */,
+        "data/yb_ysql_dump_colocated_tables_with_tablespaces.data.sql"
+        /* expectedDumpRelativePath */,
+        "expected/yb_ysql_dump_describe_colocated_tables_with_tablespaces.out"
+        /* expectedDescribeFileRelativePath */,
+        "results/yb_ysql_dump_colocated_tables_with_tablespaces.out"
+        /* outputFileRelativePath */,
+        "results/yb_ysql_dump_describe_colocated_tables_with_tablespaces.out"
+        /* outputDescribeFileRelativePath */,
+        IncludeYbMetadata.ON,
+        NoTableSpaces.ON);
   }
 
   @Test
@@ -180,7 +220,8 @@ public void ysqlDumpLegacyColocatedDB() throws Exception {
         "results/yb_ysql_dump_legacy_colocated_database.out" /* outputFileRelativePath */,
         "results/yb_ysql_dump_describe_legacy_colocated_database.out"
         /* outputDescribeFileRelativePath */,
-        IncludeYbMetadata.ON);
+        IncludeYbMetadata.ON,
+        NoTableSpaces.OFF);
   }
 
   void ysqlDumpTester(final String binaryName,
@@ -191,7 +232,8 @@ void ysqlDumpTester(final String binaryName,
                       final String expectedDescribeFileRelativePath,
                       final String outputFileRelativePath,
                       final String outputDescribeFileRelativePath,
-                      final IncludeYbMetadata includeYbMetadata) throws Exception {
+                      final IncludeYbMetadata includeYbMetadata,
+                      final NoTableSpaces noTableSpaces) throws Exception {
     // Location of Postgres regression tests
     File pgRegressDir = PgRegressBuilder.PG_REGRESS_DIR;
 
@@ -225,6 +267,9 @@ void ysqlDumpTester(final String binaryName,
     if (includeYbMetadata == IncludeYbMetadata.ON) {
       args.add("--include-yb-metadata");
     }
+    if (noTableSpaces == NoTableSpaces.ON) {
+      args.add("--no-tablespaces");
+    }
     if (!dumpedDatabaseName.isEmpty()) {
       Collections.addAll(args, "-d", dumpedDatabaseName);
     }

src/postgres/src/backend/commands/indexcmds.c

@@ -68,6 +68,7 @@
 #include "utils/tqual.h"
 
 /* YB includes. */
+#include "catalog/binary_upgrade.h"
 #include "catalog/pg_database.h"
 #include "commands/progress.h"
 #include "commands/tablegroup.h"
@@ -764,7 +765,7 @@ DefineIndex(Oid relationId,
 		char *tablegroup_name = NULL;
 
 		if (OidIsValid(tablespaceId)) 
-    {
+		{
 			/*
 			 * We look in pg_shdepend rather than directly use the derived name,
 			 * as later we might need to associate an existing implicit tablegroup to a tablespace
@@ -778,9 +779,31 @@ DefineIndex(Oid relationId,
 			tablegroup_name = OidIsValid(tablegroupId) ? get_tablegroup_name(tablegroupId) : 
 				get_implicit_tablegroup_name(tablespaceId);
 
-		} 
-    else 
-    {
+		}
+		else if (yb_binary_restore && OidIsValid(binary_upgrade_next_tablegroup_oid))
+		{
+			/*
+			 * In yb_binary_restore if tablespaceId is not valid but
+			 * binary_upgrade_next_tablegroup_oid is valid, that implies we are
+			 * restoring without tablespace information.
+			 * In this case all tables are restored to default tablespace,
+			 * while maintaining the colocation properties, and tablegroup's name
+			 * will be colocation_restore_tablegroupId, while default tablegroup's
+			 * name would still be default.
+			 */
+			tablegroup_name = binary_upgrade_next_tablegroup_default ?
+								DEFAULT_TABLEGROUP_NAME :
+								get_restore_tablegroup_name(
+									binary_upgrade_next_tablegroup_oid);
+			binary_upgrade_next_tablegroup_default = false;
+			tablegroupId = get_tablegroup_oid(tablegroup_name, true);
+		}
+		else if (yb_binary_restore && OidIsValid(tablegroupId))
+		{
+			tablegroup_name = get_tablegroup_name(tablegroupId);
+		}
+		else
+		{
 			tablegroup_name = DEFAULT_TABLEGROUP_NAME;
 			tablegroupId = get_tablegroup_oid(tablegroup_name, true);
 		}

src/postgres/src/backend/commands/tablecmds.c

@@ -12702,6 +12702,21 @@ AlterTableMoveAll(AlterTableMoveAllStmt *stmt)
 	yb_new_tablegroup_name = get_implicit_tablegroup_name(new_tablespaceoid);
 	yb_orig_tablegroup_oid = get_tablegroup_oid(yb_orig_tablegroup_name, true);
 	yb_new_tablegroup_oid = get_tablegroup_oid(yb_new_tablegroup_name, true);
+	/*
+	 * If a relation name is passed with the ALTER TABLE ALL ... COLOCATED WITH
+	 * ... SET TABLESPACE ... CASCADE command then we get the relation being
+	 * passed.
+	 */
+	if (stmt->yb_relation != NULL)
+	{
+		yb_table_oid = RangeVarGetRelid(stmt->yb_relation, NoLock, false);
+		yb_table_rel = RelationIdGetRelation(yb_table_oid);
+		yb_colocated_with_tablegroup_oid =
+			YbGetTableProperties(yb_table_rel)->tablegroup_oid;
+		RelationClose(yb_table_rel);
+		yb_orig_tablegroup_oid = yb_colocated_with_tablegroup_oid;
+		yb_orig_tablegroup_name = get_tablegroup_name(yb_orig_tablegroup_oid);
+	}
 
 	/*
 	 * The new tablespace must not have any colocated relations present in
@@ -12727,20 +12742,6 @@ AlterTableMoveAll(AlterTableMoveAllStmt *stmt)
 						" tablespace %s", stmt->orig_tablespacename),
 				 errhint("Use ALTER ... CASCADE to move colcated relations.")));
 
-	/*
-	 * If a relation name is passed with the ALTER TABLE ALL ... COLOCATED WITH
-	 * ... SET TABLESPACE ... CASCADE command then we get the relation being
-	 * passed.
-	 */
-	if (stmt->yb_relation != NULL)
-	{
-		yb_table_oid = RangeVarGetRelid(stmt->yb_relation, NoLock, false);
-		yb_table_rel = RelationIdGetRelation(yb_table_oid);
-		yb_colocated_with_tablegroup_oid =
-			YbGetTableProperties(yb_table_rel)->tablegroup_oid;
-		RelationClose(yb_table_rel);
-	}
-
 	if (OidIsValid(yb_table_oid) && !MyDatabaseColocated)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
@@ -12964,7 +12965,7 @@ AlterTableMoveAll(AlterTableMoveAllStmt *stmt)
 	 * database.
 	 */
 	if (yb_cascade && OidIsValid(new_tablespaceoid) &&
-		OidIsValid(orig_tablespaceoid) && MyDatabaseColocated &&
+		MyDatabaseColocated &&
 		!OidIsValid(yb_new_tablegroup_oid) &&
 		OidIsValid(yb_orig_tablegroup_oid))
 	{

src/postgres/src/backend/commands/tablegroup.c

@@ -83,6 +83,7 @@
 #include "pg_yb_utils.h"
 
 Oid binary_upgrade_next_tablegroup_oid = InvalidOid;
+bool binary_upgrade_next_tablegroup_default = false;
 
 /*
  * Create a table group.
@@ -430,6 +431,20 @@ get_implicit_tablegroup_name(Oid oidSuffix)
 	return tablegroup_name_from_tablespace;
 }
 
+char*
+get_restore_tablegroup_name(Oid oidSuffix)
+{
+	char *restore_tablegroup_name = (char*) palloc(
+		(
+			19 /* strlen("colocation_restore_") */ + 10 /* Max digits in OID */ +
+			1 /* Null Terminator */
+		) * sizeof(char)
+	);
+
+	sprintf(restore_tablegroup_name, "colocation_restore_%u", oidSuffix);
+	return restore_tablegroup_name;
+}
+
 /*
  * Rename tablegroup
  */

src/postgres/src/backend/commands/ybccmds.c

@@ -73,6 +73,7 @@
 #include "parser/parse_utilcmd.h"
 
 /* Yugabyte includes */
+#include "catalog/binary_upgrade.h"
 #include "catalog/pg_yb_tablegroup.h"
 #include "optimizer/clauses.h"
 
@@ -539,7 +540,7 @@ CreateTableHandleSplitOptions(YBCPgStatement handle, TupleDesc desc,
 	}
 }
 
-/* 
+/*
  * The fields pgTableId and oldRelfileNodeId are used during table rewrite.
  * During table rewrite, pgTableId is used to relay to DocDB the pg table OID,
  * so that DocDB has a mapping from the table to the correct pg table OID.
@@ -552,7 +553,7 @@ CreateTableHandleSplitOptions(YBCPgStatement handle, TupleDesc desc,
  * However, during table rewrites, stmt->relation points to the new transient
  * PG relation (named pg_temp_xxxx), but we want to create a DocDB table with
  * the same name as the original PG relation.
- */ 
+ */
 void
 YBCCreateTable(CreateStmt *stmt, char *tableName, char relkind, TupleDesc desc,
 			   Oid relationId, Oid namespaceId, Oid tablegroupId,
@@ -744,6 +745,26 @@ YBCCreateTable(CreateStmt *stmt, char *tableName, char relkind, TupleDesc desc,
 								  get_tablegroup_name(tablegroupId) :
 								  get_implicit_tablegroup_name(tablespaceId);
 		}
+		else if (yb_binary_restore &&
+				 is_colocated_tables_with_tablespace_enabled &&
+				 OidIsValid(binary_upgrade_next_tablegroup_oid))
+		{
+			/*
+			 * In yb_binary_restore if tablespaceId is not valid but
+			 * binary_upgrade_next_tablegroup_oid is valid, that implies we are
+			 * restoring without tablespace information.
+			 * In this case all tables are restored to default tablespace,
+			 * while maintaining the colocation properties, and tablegroup's name
+			 * will be colocation_restore_tablegroupId, while default tablegroup's
+			 * name would still be default.
+			 */
+			tablegroup_name = binary_upgrade_next_tablegroup_default ?
+								DEFAULT_TABLEGROUP_NAME :
+								get_restore_tablegroup_name(
+									binary_upgrade_next_tablegroup_oid);
+			binary_upgrade_next_tablegroup_default = false;
+			tablegroupId = get_tablegroup_oid(tablegroup_name, true);
+		}
 		else
 		{
 			tablegroup_name = DEFAULT_TABLEGROUP_NAME;

src/postgres/src/backend/utils/adt/pg_upgrade_support.c

@@ -219,3 +219,14 @@ binary_upgrade_set_next_tablegroup_oid(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+Datum
+binary_upgrade_set_next_tablegroup_default(PG_FUNCTION_ARGS)
+{
+	bool		next_tablegroup_default = PG_GETARG_BOOL(0);
+
+	CHECK_IS_BINARY_UPGRADE;
+	binary_upgrade_next_tablegroup_default = next_tablegroup_default;
+
+	PG_RETURN_VOID();
+}

src/postgres/src/backend/utils/adt/ruleutils.c

@@ -1561,26 +1561,6 @@ pg_get_indexdef_worker(Oid indexrelid, int colno,
 						appendStringInfo(&buf, " %s", range_split_clause);
 				}
 			}
-
-			Relation indrel = heap_open(indrelid, AccessShareLock);
-
-			/*
-			 * If the indexed table's tablegroup mismatches that of an
-			 * index table, this is a leftover from beta days of tablegroup
-			 * feature. We cannot replicate this via DDL statement anymore.
-			 */
-			if (YbGetTableProperties(indexrel)->tablegroup_oid !=
-				YbGetTableProperties(indrel)->tablegroup_oid)
-			{
-				ereport(ERROR,
-						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						 errmsg("tablegroup of an index %s does not match its "
-								"indexed table, this is no longer supported",
-								NameStr(idxrelrec->relname)),
-						 errhint("Please drop and re-create the index.")));
-			}
-
-			heap_close(indrel, AccessShareLock);
 		}
 
 		/*