Skip to content

Commit 5e5c168

Browse files
committed
[#28781] YSQL: Fix for missing unique index data of colocated db after restore
Summary: As part of the diff https://phorge.dev.yugabyte.com/D21271, support was added for backup restore of colocated database with/without //--use_tablespaces// option. It required preserving implicit tablegroup oid right before the CREATE TABLE/INDEX statement of each colocated table/indexes. This diff fixes an issue where indexes created by //unique constraints// in colocated databases weren't properly preserving their implicit tablegroup in above case. Test Plan: Added new tests Reviewers: stiwary, aagrawal, yguan Reviewed By: yguan Subscribers: yql, ybase Differential Revision: https://phorge.dev.yugabyte.com/D48081
1 parent e8a9097 commit 5e5c168

File tree

3 files changed

+293
-31
lines changed

3 files changed

+293
-31
lines changed

java/yb-pgsql/src/test/resources/TestYsqlDump/yb_ysql_dump_colocated_database.data.sql

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,12 @@ CREATE UNIQUE INDEX NONCONCURRENTLY tbl3_v_idx ON public.tbl3 USING lsm (v HASH)
540540
SELECT pg_catalog.binary_upgrade_set_next_index_pg_class_oid('16419'::pg_catalog.oid);
541541
SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('16419'::pg_catalog.oid);
542542

543+
544+
-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid
545+
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('16387'::pg_catalog.oid);
546+
547+
-- For YB colocation backup without tablespace information, must preserve default tablegroup tables
548+
SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);
543549
CREATE UNIQUE INDEX NONCONCURRENTLY tbl5_v_key ON public.tbl5 USING lsm (v ASC) WITH (colocation_id=20006);
544550

545551
ALTER TABLE ONLY public.tbl5

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

Lines changed: 60 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,10 @@ static void binary_upgrade_set_type_oids_by_rel(Archive *fout,
340340
static void binary_upgrade_set_pg_class_oids(Archive *fout,
341341
PQExpBuffer upgrade_buffer,
342342
Oid pg_class_oid, bool is_index);
343+
static void yb_binary_upgrade_preserve_index_tablegroup_oid(Archive *fout,
344+
PQExpBuffer buffer,
345+
const IndxInfo *indxinfo,
346+
const TableInfo *tbinfo);
343347
static void binary_upgrade_extension_member(PQExpBuffer upgrade_buffer,
344348
const DumpableObject *dobj,
345349
const char *objtype,
@@ -374,6 +378,7 @@ static void getYbTablePropertiesAndReloptions(Archive *fout,
374378
YbcTableProperties properties,
375379
PQExpBuffer reloptions_buf, Oid reloid, const char *relname,
376380
char relkind);
381+
static void freeYbcTablePropertiesIfRequired(YbcTableProperties yb_properties);
377382
static void isDatabaseColocated(Archive *fout);
378383
static char *getYbSplitClause(Archive *fout, const TableInfo *tbinfo);
379384
static void ybDumpUpdatePgExtensionCatalog(Archive *fout);
@@ -5345,6 +5350,53 @@ binary_upgrade_set_pg_class_oids(Archive *fout,
53455350
destroyPQExpBuffer(upgrade_query);
53465351
}
53475352

5353+
/*
5354+
* For YB colocated databases with tablespaces, preserve the tablegroup OID
5355+
* for an index's implicit tablegroup during binary upgrade. This ensures that
5356+
* the index and its data are correctly restored to the same tablegroup.
5357+
*/
5358+
static void
5359+
yb_binary_upgrade_preserve_index_tablegroup_oid(Archive *fout,
5360+
PQExpBuffer buffer,
5361+
const IndxInfo *indxinfo,
5362+
const TableInfo *tbinfo)
5363+
{
5364+
YbcTableProperties yb_properties;
5365+
PQExpBuffer yb_reloptions;
5366+
5367+
/* Only applies to colocated databases (non-legacy) */
5368+
if (!is_colocated_database || is_legacy_colocated_database)
5369+
return;
5370+
5371+
yb_properties = (YbcTableProperties) pg_malloc0(sizeof(YbcTablePropertiesData));
5372+
yb_reloptions = createPQExpBuffer();
5373+
5374+
getYbTablePropertiesAndReloptions(fout, yb_properties,
5375+
yb_reloptions,
5376+
indxinfo->dobj.catId.oid,
5377+
indxinfo->dobj.name,
5378+
tbinfo->relkind);
5379+
5380+
if (yb_properties && yb_properties->is_colocated)
5381+
{
5382+
appendPQExpBufferStr(buffer,
5383+
"\n-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid\n");
5384+
appendPQExpBuffer(buffer,
5385+
"SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('%u'::pg_catalog.oid);\n",
5386+
yb_properties->tablegroup_oid);
5387+
if (strcmp(yb_properties->tablegroup_name, "default") == 0)
5388+
{
5389+
appendPQExpBufferStr(buffer,
5390+
"\n-- For YB colocation backup without tablespace information, must preserve default tablegroup tables\n");
5391+
appendPQExpBuffer(buffer,
5392+
"SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);\n");
5393+
}
5394+
}
5395+
5396+
destroyPQExpBuffer(yb_reloptions);
5397+
freeYbcTablePropertiesIfRequired(yb_properties);
5398+
}
5399+
53485400
/*
53495401
* If the DumpableObject is a member of an extension, add a suitable
53505402
* ALTER EXTENSION ADD command to the creation commands in upgrade_buffer.
@@ -17634,38 +17686,8 @@ dumpIndex(Archive *fout, const IndxInfo *indxinfo)
1763417686
binary_upgrade_set_pg_class_oids(fout, q,
1763517687
indxinfo->dobj.catId.oid, true);
1763617688

17637-
if (is_colocated_database && !is_legacy_colocated_database)
17638-
{
17639-
YbcTableProperties yb_properties;
17640-
17641-
yb_properties = (YbcTableProperties) pg_malloc0(sizeof(YbcTablePropertiesData));
17642-
PQExpBuffer yb_reloptions = createPQExpBuffer();
17643-
17644-
getYbTablePropertiesAndReloptions(fout, yb_properties,
17645-
yb_reloptions,
17646-
indxinfo->dobj.catId.oid,
17647-
indxinfo->dobj.name,
17648-
tbinfo->relkind);
17649-
17650-
if (yb_properties && yb_properties->is_colocated)
17651-
{
17652-
appendPQExpBufferStr(q,
17653-
"\n-- For YB colocation backup, must preserve implicit tablegroup pg_yb_tablegroup oid\n");
17654-
appendPQExpBuffer(q,
17655-
"SELECT pg_catalog.binary_upgrade_set_next_tablegroup_oid('%u'::pg_catalog.oid);\n",
17656-
yb_properties->tablegroup_oid);
17657-
if (strcmp(yb_properties->tablegroup_name, "default") == 0)
17658-
{
17659-
appendPQExpBufferStr(q,
17660-
"\n-- For YB colocation backup without tablespace information, must preserve default tablegroup tables\n");
17661-
appendPQExpBuffer(q,
17662-
"SELECT pg_catalog.binary_upgrade_set_next_tablegroup_default(true);\n");
17663-
}
17664-
}
17665-
destroyPQExpBuffer(yb_reloptions);
17689+
yb_binary_upgrade_preserve_index_tablegroup_oid(fout, q, indxinfo, tbinfo);
1766617690

17667-
freeYbcTablePropertiesIfRequired(yb_properties);
17668-
}
1766917691
/* Plain secondary index */
1767017692
appendPQExpBuffer(q, "%s;\n", indxinfo->indexdef);
1767117693

@@ -17948,6 +17970,13 @@ dumpConstraint(Archive *fout, const ConstraintInfo *coninfo)
1794817970

1794917971
if (dump_index_for_constraint)
1795017972
{
17973+
/*
17974+
* YB: For colocated databases, we need to preserve the tablegroup OID
17975+
* for the index's implicit tablegroup during binary upgrade.
17976+
* This is similar to what's done in dumpIndex for non-constraint indexes.
17977+
*/
17978+
yb_binary_upgrade_preserve_index_tablegroup_oid(fout, q, indxinfo, tbinfo);
17979+
1795117980
if (dopt->include_yb_metadata || (IsYugabyteEnabled && dopt->binary_upgrade))
1795217981
{
1795317982
/*

src/yb/tools/yb-backup/yb-backup-cross-feature-test.cc

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1346,6 +1346,233 @@ TEST_F(
13461346
)#");
13471347
}
13481348

1349+
// Test backup/restore on a colocated database with tables having unique constraints.
1350+
// Verify that unique indexes are properly restored and contain the expected data.
1351+
TEST_F(
1352+
YBBackupTestColocatedTablesWithTablespaces,
1353+
YB_DISABLE_TEST_IN_SANITIZERS(TestBackupColocatedTablesWithUniqueConstraints)) {
1354+
const string& backup_db_name = "backup_db_unique";
1355+
const string& restore_db_name = "restore_db_unique";
1356+
1357+
SetupTablespaces();
1358+
1359+
ASSERT_NO_FATALS(RunPsqlCommand(
1360+
Format("CREATE DATABASE $0 WITH COLOCATION=TRUE", backup_db_name), "CREATE DATABASE"));
1361+
SetDbName(backup_db_name);
1362+
1363+
// Create tables with unique constraints in different tablespaces
1364+
ASSERT_NO_FATALS(
1365+
CreateTable("CREATE TABLE t1 (a INT PRIMARY KEY, b INT UNIQUE, c TEXT) TABLESPACE tsp1"));
1366+
ASSERT_NO_FATALS(
1367+
CreateTable("CREATE TABLE t2 (a INT PRIMARY KEY, b INT, c TEXT UNIQUE) TABLESPACE tsp2"));
1368+
ASSERT_NO_FATALS(CreateTable(
1369+
"CREATE TABLE t3 (a INT PRIMARY KEY, b INT, c TEXT, UNIQUE(b, c)) TABLESPACE tsp3"));
1370+
ASSERT_NO_FATALS(CreateTable("CREATE TABLE t4 (a INT PRIMARY KEY, b INT UNIQUE, c TEXT)"));
1371+
1372+
// Insert data into tables
1373+
for (int i = 0; i < 5; ++i) {
1374+
ASSERT_NO_FATALS(
1375+
InsertOneRow(Format("INSERT INTO t1 VALUES ($0, $1, 'text$2')", i, i + 100, i)));
1376+
}
1377+
1378+
for (int i = 0; i < 5; ++i) {
1379+
ASSERT_NO_FATALS(
1380+
InsertOneRow(Format("INSERT INTO t2 VALUES ($0, $1, 'unique$2')", i, i + 200, i)));
1381+
}
1382+
1383+
for (int i = 0; i < 5; ++i) {
1384+
ASSERT_NO_FATALS(
1385+
InsertOneRow(Format("INSERT INTO t3 VALUES ($0, $1, 'combo$2')", i, i + 300, i)));
1386+
}
1387+
1388+
for (int i = 0; i < 5; ++i) {
1389+
ASSERT_NO_FATALS(
1390+
InsertOneRow(Format("INSERT INTO t4 VALUES ($0, $1, 'default$2')", i, i + 400, i)));
1391+
}
1392+
1393+
const string backup_dir = GetTempDir("backup_unique");
1394+
const auto backup_keyspace = Format("ysql.$0", backup_db_name);
1395+
const auto restore_keyspace = Format("ysql.$0", restore_db_name);
1396+
1397+
// Backup using the --use_tablespaces flag
1398+
ASSERT_OK(RunBackupCommand(
1399+
{"--backup_location", backup_dir, "--keyspace", backup_keyspace, "--use_tablespaces",
1400+
"create"}));
1401+
1402+
ASSERT_OK(RunBackupCommand(
1403+
{"--backup_location", backup_dir, "--keyspace", restore_keyspace, "--use_tablespaces",
1404+
"restore"}));
1405+
1406+
SetDbName(restore_db_name);
1407+
1408+
// Verify data is restored correctly
1409+
RunPsqlCommand(
1410+
"SELECT * FROM t1 ORDER BY a;",
1411+
"a | b | c\n"
1412+
"---+-----+-------\n"
1413+
" 0 | 100 | text0\n"
1414+
" 1 | 101 | text1\n"
1415+
" 2 | 102 | text2\n"
1416+
" 3 | 103 | text3\n"
1417+
" 4 | 104 | text4\n"
1418+
"(5 rows)");
1419+
1420+
RunPsqlCommand(
1421+
"SELECT * FROM t2 ORDER BY a;",
1422+
"a | b | c\n"
1423+
"---+-----+---------\n"
1424+
" 0 | 200 | unique0\n"
1425+
" 1 | 201 | unique1\n"
1426+
" 2 | 202 | unique2\n"
1427+
" 3 | 203 | unique3\n"
1428+
" 4 | 204 | unique4\n"
1429+
"(5 rows)");
1430+
1431+
RunPsqlCommand(
1432+
"SELECT * FROM t3 ORDER BY a;",
1433+
"a | b | c\n"
1434+
"---+-----+--------\n"
1435+
" 0 | 300 | combo0\n"
1436+
" 1 | 301 | combo1\n"
1437+
" 2 | 302 | combo2\n"
1438+
" 3 | 303 | combo3\n"
1439+
" 4 | 304 | combo4\n"
1440+
"(5 rows)");
1441+
1442+
RunPsqlCommand(
1443+
"SELECT * FROM t4 ORDER BY a;",
1444+
"a | b | c\n"
1445+
"---+-----+----------\n"
1446+
" 0 | 400 | default0\n"
1447+
" 1 | 401 | default1\n"
1448+
" 2 | 402 | default2\n"
1449+
" 3 | 403 | default3\n"
1450+
" 4 | 404 | default4\n"
1451+
"(5 rows)");
1452+
1453+
// Verify table structures with unique indexes
1454+
RunPsqlCommand(
1455+
" \\d t1",
1456+
R"#(
1457+
Table "public.t1"
1458+
Column | Type | Collation | Nullable | Default
1459+
--------+---------+-----------+----------+---------
1460+
a | integer | | not null |
1461+
b | integer | | |
1462+
c | text | | |
1463+
Indexes:
1464+
"t1_pkey" PRIMARY KEY, lsm (a ASC), tablespace "tsp1", colocation: true
1465+
"t1_b_key" UNIQUE CONSTRAINT, lsm (b ASC), colocation: true
1466+
Tablespace: "tsp1"
1467+
Colocation: true
1468+
)#");
1469+
1470+
RunPsqlCommand(
1471+
" \\d t2",
1472+
R"#(
1473+
Table "public.t2"
1474+
Column | Type | Collation | Nullable | Default
1475+
--------+---------+-----------+----------+---------
1476+
a | integer | | not null |
1477+
b | integer | | |
1478+
c | text | | |
1479+
Indexes:
1480+
"t2_pkey" PRIMARY KEY, lsm (a ASC), tablespace "tsp2", colocation: true
1481+
"t2_c_key" UNIQUE CONSTRAINT, lsm (c ASC), colocation: true
1482+
Tablespace: "tsp2"
1483+
Colocation: true
1484+
)#");
1485+
1486+
RunPsqlCommand(
1487+
" \\d t3",
1488+
R"#(
1489+
Table "public.t3"
1490+
Column | Type | Collation | Nullable | Default
1491+
--------+---------+-----------+----------+---------
1492+
a | integer | | not null |
1493+
b | integer | | |
1494+
c | text | | |
1495+
Indexes:
1496+
"t3_pkey" PRIMARY KEY, lsm (a ASC), tablespace "tsp3", colocation: true
1497+
"t3_b_c_key" UNIQUE CONSTRAINT, lsm (b ASC, c ASC), colocation: true
1498+
Tablespace: "tsp3"
1499+
Colocation: true
1500+
)#");
1501+
1502+
RunPsqlCommand(
1503+
" \\d t4",
1504+
R"#(
1505+
Table "public.t4"
1506+
Column | Type | Collation | Nullable | Default
1507+
--------+---------+-----------+----------+---------
1508+
a | integer | | not null |
1509+
b | integer | | |
1510+
c | text | | |
1511+
Indexes:
1512+
"t4_pkey" PRIMARY KEY, lsm (a ASC), colocation: true
1513+
"t4_b_key" UNIQUE CONSTRAINT, lsm (b ASC), colocation: true
1514+
Colocation: true
1515+
)#");
1516+
1517+
// Validate unique constraint indexes using yb_index_check
1518+
RunPsqlCommand(
1519+
"SELECT yb_index_check('t1_b_key'::regclass);",
1520+
"yb_index_check\n"
1521+
"----------------\n"
1522+
"\n"
1523+
"(1 row)");
1524+
1525+
RunPsqlCommand(
1526+
"SELECT yb_index_check('t2_c_key'::regclass);",
1527+
"yb_index_check\n"
1528+
"----------------\n"
1529+
"\n"
1530+
"(1 row)");
1531+
1532+
RunPsqlCommand(
1533+
"SELECT yb_index_check('t3_b_c_key'::regclass);",
1534+
"yb_index_check\n"
1535+
"----------------\n"
1536+
"\n"
1537+
"(1 row)");
1538+
1539+
RunPsqlCommand(
1540+
"SELECT yb_index_check('t4_b_key'::regclass);",
1541+
"yb_index_check\n"
1542+
"----------------\n"
1543+
"\n"
1544+
"(1 row)");
1545+
1546+
// Validate indexes by performing index-only scans
1547+
RunPsqlCommand(
1548+
"/*+ IndexOnlyScan(t1 t1_b_key) */ SELECT * FROM t1 WHERE b = 102;",
1549+
"a | b | c\n"
1550+
"---+-----+-------\n"
1551+
" 2 | 102 | text2\n"
1552+
"(1 row)");
1553+
1554+
RunPsqlCommand(
1555+
"/*+ IndexOnlyScan(t2 t2_c_key) */ SELECT * FROM t2 WHERE c = 'unique3';",
1556+
"a | b | c\n"
1557+
"---+-----+---------\n"
1558+
" 3 | 203 | unique3\n"
1559+
"(1 row)");
1560+
1561+
RunPsqlCommand(
1562+
"/*+ IndexOnlyScan(t3 t3_b_c_key) */ SELECT * FROM t3 WHERE b = 301 AND c = 'combo1';",
1563+
"a | b | c\n"
1564+
"---+-----+--------\n"
1565+
" 1 | 301 | combo1\n"
1566+
"(1 row)");
1567+
1568+
RunPsqlCommand(
1569+
"/*+ IndexOnlyScan(t4 t4_b_key) */ SELECT * FROM t4 WHERE b = 403;",
1570+
"a | b | c\n"
1571+
"---+-----+----------\n"
1572+
" 3 | 403 | default3\n"
1573+
"(1 row)");
1574+
}
1575+
13491576
// Test backup/restore of geo-partitioned colocated tables with tablespace information.
13501577
TEST_F(
13511578
YBBackupTestColocatedTablesWithTablespaces,

0 commit comments

Comments
 (0)