[#24809] YSQL: Fix Default Tablespace Backup/Restore For Colocated Ta…

…bles

Summary:
In commit 0dbe7d6, support was introduced for backup and restore of colocated tables with tablespaces, allowing for backup/restore without explicit tablespace information (without `--use_tablespaces`). This approach included:

 1. Adding a binary upgrade option in YSQLDump (`pg_catalog.binary_upgrade_set_next_tablegroup_default`), which was set true when (1) a backup was done without `--use_tablespaces`, and (2) the table originally belonged to the default tablespace/tablegroup.

 2. Utilizing the existing `pg_catalog.binary_upgrade_set_next_tablegroup_oid` to retain tablegroup IDs across backup and restore, a feature introduced in commit 254b726.

During restoration, tablegroups are named:
 - With `--use_tablespaces`: `colocation_<new_tablespace_oid>`, achieved by utilizing the tablespace data in YSQLDump.
 - Without `--use_tablespaces`: During restoration without `--use_tablespaces`, tablegroups are named as `colocation_restore_<tablegroup_oid>` based on the following conditions
  - If `pg_catalog.binary_upgrade_set_next_tablegroup_oid` has a valid `tablegroup_oid` and `pg_catalog.binary_upgrade_set_next_tablegroup_default` is set to false, then the tablegroup is named `colocation_restore_<tablegroup_oid>`.
  - If `pg_catalog.binary_upgrade_set_next_tablegroup_default` is set to true or if tablegroup_oid is invalid, the tablegroup is named as "default."
 - In both cases, the original "default" tablegroup retains its name after restoration.

**Problem**
Previously, for backups made with `--use_tablespaces`, the YSQLDump structure led to inconsistencies between tables sharing the same colocation tablet (in default tablespace) but assigned different tablegroups on restore. For example, consider the following tables:
```
CREATE DATABASE db WITH COLOCATION=true;
CREATE TABLE t1(col INT) TABLESPACE tsp1;
CREATE TABLE t2(col INT);
CREATE TABLE t3(col INT) TABLESPACE tsp1;
CREATE TABLE t4(col INT);
```

The YSQLDump output looked like this:
```
SET default_tablespace = "tsp1";
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid(<tablegroup_oid1>);
CREATE TABLE t1(col INT);

SET default_tablespace = "";  -- default tablespace
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid(<tablegroup_oid2>);
CREATE TABLE t2(col INT);

SET default_tablespace = "tsp1";
CREATE TABLE t3(col INT);  -- no oid logged again for already touched tablespace

SET default_tablespace = ""; -- default tablespace
CREATE TABLE t4(col INT);  -- no oid logged again for already touched tablespace
```

This inconsistency resulted in tables like `t2` and `t4`, which are part of the same colocation tablet in default tablespace, receiving different tablegroups after restoration:

 - Table t2 was assigned to the tablegroup `colocation_restore_<tablegroup_oid2>`, as it had a valid tablegroup OID and an invalid tablespace OID (indicating the default tablespace).
 - Table t4 was assigned to the default tablegroup, as it had both an invalid tablespace and tablegroup OID.

This discrepancy in assigned tablegroups was due to inconsistencies in the logged tablegroup information during YSQLDump creation.

**Fix**
To address this, the revision modifies YSQLDump to now log `pg_catalog.binary_upgrade_set_next_tablegroup_oid` and `pg_catalog.binary_upgrade_set_next_tablegroup_default` for every table creation, in `--use_tablespaces` . This approach ensures that `t2` and `t4` are consistently assigned the same default tablegroup after restoration.

```
SET default_tablespace = "tsp1";
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid(<tablegroup_oid1>);
CREATE TABLE t1(col INT);

SET default_tablespace = "";  -- default tablespace
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid(<tablegroup_oid2>);
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);  -- added
CREATE TABLE t2(col INT);

SET default_tablespace = "tsp1";
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid(<tablegroup_oid1>);  -- added
CREATE TABLE t3(col INT);

SET default_tablespace = "";
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid(<tablegroup_oid2>);  -- added
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);  -- added
CREATE TABLE t4(col INT);
```

This revision also resolves unrelated lint errors in the pg code.

JIRA: DB-13915

Test Plan:
./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 --cxx-test yb-backup-cross-feature-test --gtest_filter YBBackupTestColocatedTablesWithTablespaces.TestBackupGeoPartitionedColocatedTablesWithTablespaces**: This is a newly introduced  test, which reproduces the bug.
./yb_build.sh --cxx-test yb-backup-cross-feature-test --gtest_filter YBBackupTestColocatedTablesWithTablespaces.TestBackupGeoPartitionedColocatedTablesWithoutTablespaces
./yb_build.sh --java-test org.yb.pgsql.TestYsqlDump#ysqlDumpColocatedTablesWithTablespaces

./yb_build.sh --java-test org.yb.pgsql.TestPgRegressColocatedTablesWithTablespaces#testPgRegressColocatedTablesWithTablespaces

**./yb_build.sh --java-test 'org.yb.pgsql.TestYsqlDump#ysqlDumpColocatedDB'** : Modified this test to accommodate the changes in the revision
./yb_build.sh --java-test 'org.yb.pgsql.TestYsqlDump#ysqlDumpLegacyColocatedDB'

./yb_build.sh --cxx_test yb-backup-test_ent --gtest-filter YBBackupTest.TestColocationDuplication
./yb_build.sh --cxx_test yb-backup-test_ent --gtest-filter YBBackupTest.TestLegacyColocatedDBColocationDuplication

./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testAlteredTableInColocatedDB'
./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testAlteredTableInLegacyColocatedDB'

./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testMixedColocatedDatabase'
./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testMixedLegacyColocatedDatabase'

./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testColocatedDBWithColocationIdAlreadySet'
./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testLegacyColocatedDBWithColocationIdAlreadySet'

./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testColocatedDatabaseRestoreToOriginalDB'
./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testLegacyColocatedDatabaseRestoreToOriginalDB'

./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testColocatedMateralizedViewBackup'
./yb_build.sh --java-test 'org.yb.pgsql.TestYbBackup#testLegacyColocatedMateralizedViewBackup'

Reviewers: skumar, yguan, aagrawal

Reviewed By: yguan, aagrawal

Subscribers: yql, ybase

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

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

@@ -1051,8 +1051,9 @@ DefineIndex(Oid relationId,
 			{
 				/*
 				 * 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.
+				 * binary_upgrade_next_tablegroup_oid is valid, that implies either:
+				 * 1. it is a default tablespace.
+				 * 2. 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
@@ -1067,6 +1068,10 @@ DefineIndex(Oid relationId,
 			}
 			else if (yb_binary_restore && OidIsValid(tablegroupId))
 			{
+				/*
+				 * This case handles Primary Key's tablegroup id. The variable
+				 * tablegroupId stores the tablegroupId of the parent table.
+				 */
 				tablegroup_name = get_tablegroup_name(tablegroupId);
 			}
 			else
@@ -1097,6 +1102,14 @@ DefineIndex(Oid relationId,
 				tablegroupId = CreateTableGroup(tablegroup_stmt);
 			}
 		}
+		/*
+		 * Reset the binary_upgrade params as these are not needed anymore (only
+		 * required in CreateTableGroup), to ensure these parameter values are
+		 * not reused in subsequent unrelated statements.
+		 */
+		binary_upgrade_next_tablegroup_oid = InvalidOid;
+		binary_upgrade_next_tablegroup_default = false;
+
 
 		if (stmt->split_options)
 		{
@@ -2610,8 +2623,8 @@ ComputeIndexAttrs(IndexInfo *indexInfo,
 										 accessMethodId);
 
 		/*
-	 	 * In Yugabyte mode, disallow some built-in operator classes if the column has non-C
-	 	 * collation.
+		 * In Yugabyte mode, disallow some built-in operator classes if the column has non-C
+		 * collation.
 		 */
 		if (IsYugaByteEnabled() &&
 			YBIsCollationValidNonC(attcollation) &&
@@ -2725,9 +2738,9 @@ ComputeIndexAttrs(IndexInfo *indexInfo,
 			}
 			else if (colOptionP[attn] == INDOPTION_HASH)
 				ereport(NOTICE,
-                		(errcode(ERRCODE_WRONG_OBJECT_TYPE),
-                  		 errmsg("nulls sort ordering option is ignored, "
-                        		"NULLS FIRST/NULLS LAST not allowed for a HASH column")));
+						(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+						 errmsg("nulls sort ordering option is ignored, "
+								"NULLS FIRST/NULLS LAST not allowed for a HASH column")));
 			else if (attribute->nulls_ordering == SORTBY_NULLS_FIRST)
 				colOptionP[attn] |= INDOPTION_NULLS_FIRST;
 		}

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

@@ -716,8 +716,9 @@ YBCCreateTable(CreateStmt *stmt, char *tableName, char relkind, TupleDesc desc,
 		{
 			/*
 			 * 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.
+			 * binary_upgrade_next_tablegroup_oid is valid, that implies either:
+			 * 1. it is a default tablespace.
+			 * 2. 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
@@ -758,6 +759,13 @@ YBCCreateTable(CreateStmt *stmt, char *tableName, char relkind, TupleDesc desc,
 			tablegroupId = CreateTableGroup(tablegroup_stmt);
 			stmt->tablegroupname = pstrdup(tablegroup_name);
 		}
+		/*
+		 * Reset the binary_upgrade params as these are not needed anymore (only
+		 * required in CreateTableGroup), to ensure these parameter values are
+		 * not reused in subsequent unrelated statements.
+		 */
+		binary_upgrade_next_tablegroup_oid = InvalidOid;
+		binary_upgrade_next_tablegroup_default = false;
 
 		/* Record dependency between the table and tablegroup. */
 		ObjectAddress myself, tablegroup;
@@ -1023,10 +1031,10 @@ YbUnsafeTruncate(Relation rel)
 /* Utility function to handle split points */
 static void
 CreateIndexHandleSplitOptions(YBCPgStatement handle,
-                              TupleDesc desc,
-                              OptSplit *split_options,
-                              int16 * coloptions,
-                              int numIndexKeyAttrs)
+							  TupleDesc desc,
+							  OptSplit *split_options,
+							  int16 * coloptions,
+							  int numIndexKeyAttrs)
 {
 	/* Address both types of split options */
 	switch (split_options->split_type)
@@ -1175,7 +1183,7 @@ YBCCreateIndex(const char *indexName,
 	/* Handle SPLIT statement, if present */
 	if (split_options)
 		CreateIndexHandleSplitOptions(handle, indexTupleDesc, split_options, coloptions,
-		                              indexInfo->ii_NumIndexKeyAttrs);
+									  indexInfo->ii_NumIndexKeyAttrs);
 
 	/* Create the index. */
 	HandleYBStatus(YBCPgExecCreateIndex(handle));

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

@@ -144,11 +144,6 @@ static SimpleStringList table_exclude_patterns = {NULL, NULL};
 static SimpleOidList table_exclude_oids = {NULL, NULL};
 static SimpleStringList tabledata_exclude_patterns = {NULL, NULL};
 static SimpleOidList tabledata_exclude_oids = {NULL, NULL};
-/*
- * The string list records tablespaces used if the dumped database is
- * a colocated database.
- */
-static SimpleStringList colocated_database_tablespaces = {NULL, NULL};
 
 static SimpleStringList foreign_servers_include_patterns = {NULL, NULL};
 static SimpleOidList foreign_servers_include_oids = {NULL, NULL};
@@ -15664,12 +15659,9 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 		 */
 		if (is_colocated_database && !is_legacy_colocated_database
 			&& (tbinfo->relkind == RELKIND_RELATION || tbinfo->relkind == RELKIND_MATVIEW
-				|| tbinfo->relkind == RELKIND_PARTITIONED_TABLE)
-			&& yb_properties && yb_properties->is_colocated
-			&& (!simple_string_list_member(&colocated_database_tablespaces, tbinfo->reltablespace)
-			|| dopt->outputNoTablespaces))
+				|| tbinfo->relkind == RELKIND_PARTITIONED_TABLE) && yb_properties
+			&& yb_properties->is_colocated)
 		{
-			simple_string_list_append(&colocated_database_tablespaces, tbinfo->reltablespace);
 			/*
 			 * Set the next implicit tablegroup oid in a colocated database.
 			 * It's mandatory to reuse the old tablegroup oid to match tablegroup parent table
@@ -15680,15 +15672,13 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 			appendPQExpBuffer(q,
 							  "SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('%u'::pg_catalog.oid);\n",
 							  yb_properties->tablegroup_oid);
-			if (dopt->outputNoTablespaces)
+
+			if(strcmp(yb_properties->tablegroup_name, "default") == 0)
 			{
-				if(strcmp(yb_properties->tablegroup_name, "default") == 0)
-				{
-					appendPQExpBufferStr(q,
-									"\n-- For YB colocation backup without tablespace information, must preserve default tablegroup tables\n");
-					appendPQExpBuffer(q,
-						"SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);\n");
-				}
+				appendPQExpBufferStr(q,
+									 "\n-- For YB colocation backup without tablespace information, must preserve default tablegroup tables\n");
+				appendPQExpBuffer(q,
+								  "SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);\n");
 			}
 		}
 
@@ -16644,7 +16634,7 @@ dumpIndex(Archive *fout, const IndxInfo *indxinfo)
 			binary_upgrade_set_pg_class_oids(fout, q,
 											 indxinfo->dobj.catId.oid, true);
 
-		if (dopt->outputNoTablespaces && is_colocated_database && !is_legacy_colocated_database)
+		if (is_colocated_database && !is_legacy_colocated_database)
 		{
 			YbTableProperties yb_properties;
 			yb_properties = (YbTableProperties) pg_malloc(sizeof(YbTablePropertiesData));

src/postgres/src/test/regress/data/yb_ysql_dump_colocated_database.data.sql

@@ -54,6 +54,9 @@ SELECT pg_catalog.binary_upgrade_set_next_heap_pg_class_oid('16410'::pg_catalog.
 
 -- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
 SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE TABLE public.htest (
     k1 integer,
     k2 text,
@@ -88,6 +91,12 @@ SELECT pg_catalog.binary_upgrade_set_next_array_pg_type_oid('16414'::pg_catalog.
 SELECT pg_catalog.binary_upgrade_set_next_heap_pg_class_oid('16413'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_heap_relfilenode('16413'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE TABLE public.htest_1 (
     k1 integer,
     k2 text,
@@ -124,6 +133,12 @@ SELECT pg_catalog.binary_upgrade_set_next_heap_relfilenode('16384'::pg_catalog.o
 SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16388'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16388'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE TABLE public.tbl (
     k integer NOT NULL,
     v integer,
@@ -158,6 +173,12 @@ SELECT pg_catalog.binary_upgrade_set_next_heap_relfilenode('16390'::pg_catalog.o
 SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16393'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16393'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE TABLE public.tbl2 (
     k integer NOT NULL,
     v integer,
@@ -259,6 +280,12 @@ SELECT pg_catalog.binary_upgrade_set_next_array_pg_type_oid('16417'::pg_catalog.
 SELECT pg_catalog.binary_upgrade_set_next_heap_pg_class_oid('16416'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_heap_relfilenode('16416'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE TABLE public.tbl5 (
     k integer,
     v integer
@@ -349,6 +376,12 @@ ALTER TABLE ONLY public.tbl5
 SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16398'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16398'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE INDEX NONCONCURRENTLY partial_idx ON public.tbl2 USING lsm (k ASC, v DESC) WITH (colocation_id=40001) WHERE ((k > 10) AND (k < 20) AND (v > 200));
 
 
@@ -361,6 +394,12 @@ CREATE INDEX NONCONCURRENTLY partial_idx ON public.tbl2 USING lsm (k ASC, v DESC
 SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16397'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16397'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE UNIQUE INDEX NONCONCURRENTLY partial_unique_idx ON public.tbl USING lsm (v DESC) WITH (colocation_id=40000) WHERE ((v >= 100) AND (v <= 200));
 
 
@@ -373,6 +412,12 @@ CREATE UNIQUE INDEX NONCONCURRENTLY partial_unique_idx ON public.tbl USING lsm (
 SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16396'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16396'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE INDEX NONCONCURRENTLY tbl2_v2_idx ON public.tbl2 USING lsm (v2 ASC) WITH (colocation_id=20004);
 
 
@@ -397,6 +442,12 @@ CREATE UNIQUE INDEX NONCONCURRENTLY tbl3_v_idx ON public.tbl3 USING lsm (v HASH)
 SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16395'::pg_catalog.oid);
 SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16395'::pg_catalog.oid);
 
+
+-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
+
+-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
+SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
 CREATE UNIQUE INDEX NONCONCURRENTLY tbl_v_idx ON public.tbl USING lsm (v DESC) WITH (colocation_id=20003);