From 50ff737e9ed909f27d46247149de771266c8b2a0 Mon Sep 17 00:00:00 2001
From: Anubhav Srivastava <asrivastava@yugabyte.com>
Date: Tue, 27 Aug 2024 09:13:03 -0700
Subject: [PATCH] [#23741] docdb: Fix cloning of colocated databases with only
 parent table

Summary:
A colocated database might have only the parent table if all user tables were dropped.
Import snapshot will fail to find the corresponding running parent table if we try to clone right
now. This diff excludes the parent table when generating backup entries if there are no user
tables.

Test Plan: `./yb_build.sh release --cxx-test integration-tests_minicluster-snapshot-test --gtest_filter PgCloneColocationTest.NoColocatedChildTables`

Reviewers: mhaddad

Reviewed By: mhaddad

Subscribers: ybase

Differential Revision: https://phorge.dev.yugabyte.com/D37695
---
 .../minicluster-snapshot-test.cc              | 18 ++++++++++++++
 src/yb/master/catalog_manager_ext.cc          | 24 ++++++++++++-------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/yb/integration-tests/minicluster-snapshot-test.cc b/src/yb/integration-tests/minicluster-snapshot-test.cc
index bf885f2ec3c4..9151fe89c16d 100644
--- a/src/yb/integration-tests/minicluster-snapshot-test.cc
+++ b/src/yb/integration-tests/minicluster-snapshot-test.cc
@@ -1032,5 +1032,23 @@ TEST_P(PgCloneTestWithColocatedDBParam, YB_DISABLE_TEST_IN_SANITIZERS(CloneOfClo
   ASSERT_RESULT(ConnectToDB(kTargetNamespaceName2));
 }
 
+TEST_F(PgCloneColocationTest, YB_DISABLE_TEST_IN_SANITIZERS(NoColocatedChildTables)) {
+  ASSERT_OK(source_conn_->Execute("CREATE TABLE t2(k int, v1 int) WITH (COLOCATION = false)"));
+  ASSERT_OK(source_conn_->Execute("DROP TABLE t1"));
+  auto no_child_tables_time = ASSERT_RESULT(GetCurrentTime()).ToInt64();
+  ASSERT_OK(source_conn_->Execute("DROP TABLE t2"));
+
+  // Clone to a time when there are no colocated child tables.
+  ASSERT_OK(source_conn_->ExecuteFormat(
+      "CREATE DATABASE $0 TEMPLATE $1 AS OF $2", kTargetNamespaceName1, kSourceNamespaceName,
+      no_child_tables_time));
+  ASSERT_RESULT(ConnectToDB(kTargetNamespaceName1));
+
+  // Clone to a time when there are no tables.
+  ASSERT_OK(source_conn_->ExecuteFormat(
+      "CREATE DATABASE $0 TEMPLATE $1", kTargetNamespaceName2, kSourceNamespaceName));
+  ASSERT_RESULT(ConnectToDB(kTargetNamespaceName2));
+}
+
 }  // namespace master
 }  // namespace yb
diff --git a/src/yb/master/catalog_manager_ext.cc b/src/yb/master/catalog_manager_ext.cc
index 3131b2c0d1fe..f4cf3100fa91 100644
--- a/src/yb/master/catalog_manager_ext.cc
+++ b/src/yb/master/catalog_manager_ext.cc
@@ -1209,25 +1209,30 @@ Result<RepeatedPtrField<BackupRowEntryPB>> CatalogManager::GetBackupEntriesAsOfT
   // partitions' start keys.
   std::map<TableId, TableWithTabletsEntries> tables_to_tablets;
   std::optional<std::string> colocation_parent_table_id;
+  bool found_colocated_user_table = false;
   docdb::DocRowwiseIterator tables_iter = docdb::DocRowwiseIterator(
       projection, doc_read_cntxt, TransactionOperationContext(), doc_db,
       docdb::ReadOperationData::FromSingleReadTime(read_time), db_pending_op);
   RETURN_NOT_OK(EnumerateSysCatalog(
       &tables_iter, doc_read_cntxt.schema(), SysRowEntryType::TABLE,
-      [&source_ns_id, &tables_to_tablets, &colocation_parent_table_id](
+      [&source_ns_id, &tables_to_tablets, &colocation_parent_table_id, &found_colocated_user_table](
           const Slice& id, const Slice& data) -> Status {
         auto pb = VERIFY_RESULT(pb_util::ParseFromSlice<SysTablesEntryPB>(data));
         if (pb.namespace_id() == source_ns_id && pb.state() == SysTablesEntryPB::RUNNING &&
             pb.hide_state() == SysTablesEntryPB_HideState_VISIBLE &&
             !pb.schema().table_properties().is_ysql_catalog_table()) {
           VLOG_WITH_FUNC(1) << "Found SysTablesEntryPB: " << pb.ShortDebugString();
-          if (IsColocatedDbParentTableId(id.ToBuffer()) ||
-              IsColocatedDbTablegroupParentTableId(id.ToBuffer())) {
-            colocation_parent_table_id = id.ToBuffer();
+          const auto id_str = id.ToBuffer();
+          if (pb.colocated()) {
+            if (IsColocationParentTableId(id_str)) {
+              colocation_parent_table_id = id_str;
+            } else {
+              found_colocated_user_table = true;
+            }
           }
           // Tables and tablets will be added to backup entries at the end.
           tables_to_tablets.insert(std::make_pair(
-              id.ToBuffer(), TableWithTabletsEntries(pb, SysTabletsEntriesWithIds())));
+              id_str, TableWithTabletsEntries(pb, SysTabletsEntriesWithIds())));
         }
         return Status::OK();
       }));
@@ -1261,11 +1266,14 @@ Result<RepeatedPtrField<BackupRowEntryPB>> CatalogManager::GetBackupEntriesAsOfT
   for (auto& sys_table_entry : tables_to_tablets) {
     sys_table_entry.second.OrderTabletsByPartitions();
   }
-  // Populate the backup_entries with SysTablesEntry and SysTabletsEntry
+  // Populate the backup_entries with SysTablesEntry and SysTabletsEntry.
   // Start with the colocation_parent_table_id if the database is colocated.
   if (colocation_parent_table_id) {
-    tables_to_tablets[colocation_parent_table_id.value()].AddToBackupEntries(
-        colocation_parent_table_id.value(), &backup_entries);
+    // Only create the colocated parent table if there are colocated user tables.
+    if (found_colocated_user_table) {
+      tables_to_tablets[colocation_parent_table_id.value()].AddToBackupEntries(
+          colocation_parent_table_id.value(), &backup_entries);
+    }
     tables_to_tablets.erase(colocation_parent_table_id.value());
   }
   for (auto& sys_table_entry : tables_to_tablets) {