[BACKPORT pg15-cherrypicks][#23312] YSQL: Fix backup/restore for colo…

…cated tables without tablespace information.

Summary:
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

Original commit: 0dbe7d6 / D36859
- src/postgres/src/backend/commands/indexcmds.c
  - DefineIndex, in if (IsYugabyteEnabled()) block after is_colocated assignment:
    - This YB master commit adjusts indentation of braces.
    - pg15 initial merge 55782d5 replaced the block with assignment of tablegroupId.
    - Take pg15 change. Note some diff tools may resolve this automatically because the master change is whitespace-only.
  - DefineIndex, after declaration of is_colocated_tables_with_tablespace_enabled:
    - This YB master commit adds logic, and fixes indentation of braces.
    - pg15 merge commit 4e9241b indents logic around if (is_colocated_tables_with_tablespace_enabled && is_colocated_via_database && !MyColocatedDatabaseLegacy) to be inside IsYugabyteEnabled().
    - Take the if (is_colocated_tables_with_tablespace_enabled && is_colocated_via_database && !MyColocatedDatabaseLegacy) condition first. Then pull the entire block beginning char *tablegroup_name = NULL;, adjusting the indentation, up to the end of the block for the OidIsValid(tablespaceId) condition. Then take the master change of the remaining conflict, but indent it to be correctly indented within the IsYugabyteEnabled().
- src/postgres/src/backend/utils/adt/ruleutils.c
  - pg_get_indexdef_worker
   - This YB master commit removed piece of code.
   - YB Pg15 commit 55782d5 changed heap_open (this was present in code which was removed in this master commit) to table_open.
   - Resolved by accepting incoming master commit hence completely removing piece of code.
- src/postgres/src/bin/pg_dump/pg_dump.c
 - getYbTablePropertiesAndReloptions:
  - This YB master commit introduces new piece of code.
  - YB Pg15 commit 55782d5, changed exit_horribly to fatal and introduced #ifdef YB_TODO and #endif
  - Resolved conflict by accepting a combination of both the changes, by keeping it fatal with ifdef and then introducing new code snippet

- yb_ysql_dump_colocated_tables_with_tablespaces.data.sql
 - Upstream PG commit 9a974cbcba005256a19991203583a94b4f9a21a9 preserves relfilenodes.
 - Upstream PG commit 9fbc3f318d039c3e1e8614c38e40843cf8fcffde adds dump of pg_stat_statements_info. YB role-related code added in YB master commit 5659b73.
 - Other minor PG15 changes.
- yb_ysql_dump_describe_colocated_tables_with_tablespaces.out
 - Change is likely due to changes in the execution plan, similar to schema ordering differences found in YB pg15 ab32fa2

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: jason, tfoucher

Reviewed By: tfoucher

Subscribers: ybase, yql, mlillibridge, svc_phabricator

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

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

@@ -71,6 +71,7 @@
 #include "utils/syscache.h"
 
 /* YB includes. */
+#include "catalog/binary_upgrade.h"
 #include "catalog/pg_database.h"
 #include "commands/progress.h"
 #include "commands/tablegroup.h"
@@ -1012,9 +1013,9 @@ DefineIndex(Oid relationId,
 			is_colocated_via_database && !MyColocatedDatabaseLegacy)
 		{
 			char *tablegroup_name = NULL;
-		
-			if (OidIsValid(tablespaceId)) 
-		{
+
+			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
@@ -1028,9 +1029,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

@@ -15587,6 +15587,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
@@ -15612,20 +15627,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),
@@ -15848,7 +15849,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

@@ -82,6 +82,7 @@
 #include "pg_yb_utils.h"
 
 Oid binary_upgrade_next_tablegroup_oid = InvalidOid;
+bool binary_upgrade_next_tablegroup_default = false;
 
 /*
  * Create a table group.
@@ -432,6 +433,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

@@ -76,6 +76,7 @@
 #include "parser/parse_utilcmd.h"
 
 /* Yugabyte includes */
+#include "catalog/binary_upgrade.h"
 #include "catalog/pg_yb_tablegroup.h"
 #include "optimizer/clauses.h"
 
@@ -502,7 +503,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.
@@ -515,7 +516,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,
@@ -703,6 +704,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

@@ -274,3 +274,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

@@ -1622,26 +1622,6 @@ pg_get_indexdef_worker(Oid indexrelid, int colno,
 						appendStringInfo(&buf, " %s", range_split_clause);
 				}
 			}
-
-			Relation indrel = table_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.")));
-			}
-
-			table_close(indrel, AccessShareLock);
 		}
 
 		/*