[PLAT-13802]: Add a field to allow bootstrap on whole universe when o…

…nly few tables requires bootstrap

Summary:
Added a boolean field in Bootstarp params to allow bootstrap on the whole universe when only a few tables require bootstrap.

`Note`: We need to remove the stream of tables that are already part of database during backup/restore while editing xcluster config but they are added back after restore.

Test Plan:
Added UT.

Manual Test:

Create a replication with 5 tables in a ysql database and later remove them. 2 tables and insert some data in them so that they require bootstrap.
So, while editing tables and adding again 2 new tables in the replication, passed only the new 2 tables in bootstrap require and passed allowBootstrap as true which allow YBA to perform backup/restore on the whole database, and finally set replication with the requested tables.

Reviewers: hzare, #yba-api-review, sanketh, cwang

Reviewed By: hzare, #yba-api-review

Subscribers: dkumar, cwang, yugaware

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

managed/src/main/java/com/yugabyte/yw/commissioner/tasks/EditXClusterConfig.java

@@ -25,6 +25,7 @@
 import javax.inject.Inject;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.util.CollectionUtils;
+import org.yb.CommonTypes;
 import org.yb.master.MasterDdlOuterClass;
 import org.yb.master.MasterDdlOuterClass.ListTablesResponsePB.TableInfo;
 
@@ -235,13 +236,31 @@ protected void addSubtasksToAddTablesToXClusterConfig(
 
     // Add the subtasks to set up replication for tables that need bootstrapping if any.
     if (!dbToTablesInfoMapNeedBootstrap.isEmpty()) {
-      // YSQL tables replication with bootstrapping can only be set up with DB granularity. The
-      // following subtasks remove tables in replication, so the replication can be set up again
-      // for all the tables in the DB including the new tables.
       Set<String> tableIdsDeleteReplication = new HashSet<>();
       dbToTablesInfoMapNeedBootstrap.forEach(
           (namespaceName, tablesInfo) -> {
             Set<String> tableIdsNeedBootstrap = getTableIds(tablesInfo);
+            if (taskParams().getBootstrapParams().allowBootstrap && !tablesInfo.isEmpty()) {
+              if (tablesInfo.get(0).getTableType() == CommonTypes.TableType.PGSQL_TABLE_TYPE) {
+                List<TableInfo> sourceTablesInfo =
+                    getTableInfoListByNamespaceName(
+                        ybService,
+                        Universe.getOrBadRequest(xClusterConfig.getSourceUniverseUUID()),
+                        CommonTypes.TableType.PGSQL_TABLE_TYPE,
+                        namespaceName);
+                tableIdsNeedBootstrap.addAll(getTableIds(sourceTablesInfo));
+                tableIdsNeedBootstrap.addAll(getTableIds(tablesInfo));
+              } else {
+                groupByNamespaceName(requestedTableInfoList).get(namespaceName).stream()
+                    .map(tableInfo -> getTableId(tableInfo))
+                    .forEach(tableIdsNeedBootstrap::add);
+              }
+            }
+            // YSQL tables replication with bootstrapping can only be set up with DB granularity.
+            // The
+            // following subtasks remove tables in replication, so the replication can be set up
+            // again
+            // for all the tables in the DB including the new tables.
             Set<String> tableIdsNeedBootstrapInReplication =
                 xClusterConfig.getTableIdsWithReplicationSetup(
                     tableIdsNeedBootstrap, true /* done */);

managed/src/main/java/com/yugabyte/yw/commissioner/tasks/XClusterConfigTaskBase.java

@@ -959,7 +959,8 @@ protected boolean syncReplicationSetUpStateForTables(
 
     if (bootstrapParams != null && bootstrapParams.tables != null) {
       // Ensure tables in bootstrapParams is a subset of requestedTableIds.
-      if (!requestedTableIds.containsAll(bootstrapParams.tables)) {
+      if (!bootstrapParams.allowBootstrap
+          && !requestedTableIds.containsAll(bootstrapParams.tables)) {
         throw new IllegalArgumentException(
             String.format(
                 "The set of tables in bootstrapParams (%s) is not a subset of "
@@ -977,13 +978,16 @@ protected boolean syncReplicationSetUpStateForTables(
                   HashMap::new,
                   (map, entry) -> map.put(entry.getKey(), entry.getValue()),
                   HashMap::putAll);
-      bootstrapParams.tables =
-          getTableIdsWithoutTablesOnTargetInReplication(
-              ybService,
-              requestedTableInfoList,
-              sourceTableIdTargetTableIdWithBootstrapMap,
-              targetUniverse,
-              currentReplicationGroupName);
+
+      if (!bootstrapParams.allowBootstrap) {
+        bootstrapParams.tables =
+            getTableIdsWithoutTablesOnTargetInReplication(
+                ybService,
+                requestedTableInfoList,
+                sourceTableIdTargetTableIdWithBootstrapMap,
+                targetUniverse,
+                currentReplicationGroupName);
+      }
 
       // If table type is YSQL and bootstrap is requested, all tables in that keyspace are selected.
       if (tableType == CommonTypes.TableType.PGSQL_TABLE_TYPE) {
@@ -1009,8 +1013,9 @@ protected boolean syncReplicationSetUpStateForTables(
                                             .equals(namespaceId))
                             .map(tableInfo -> tableInfo.getId().toStringUtf8())
                             .collect(Collectors.toSet());
-                    if (tableIdsInNamespace.size()
-                        != selectedTableIdsInNamespaceToBootstrap.size()) {
+                    if (!bootstrapParams.allowBootstrap
+                        && tableIdsInNamespace.size()
+                            != selectedTableIdsInNamespaceToBootstrap.size()) {
                       throw new IllegalArgumentException(
                           String.format(
                               "For YSQL tables, all the tables in a keyspace must be selected: "
@@ -1385,6 +1390,44 @@ public static List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo> getTableI
         .collect(Collectors.toList());
   }
 
+  /**
+   * This method returns all the tablesInfo list present in the namespace on a universe.
+   *
+   * @param ybService The service to get a YB client from
+   * @param universe The universe to get the table schema information from
+   * @param tableType The table type to filter the tables
+   * @param namespaceName The namespace name to get the table schema information from
+   * @return A list of {@link MasterDdlOuterClass.ListTablesResponsePB.TableInfo} containing table
+   *     info of the tables in the namespace
+   */
+  public static List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo>
+      getTableInfoListByNamespaceName(
+          YBClientService ybService, Universe universe, TableType tableType, String namespaceName) {
+    return getTableInfoList(ybService, universe).stream()
+        .filter(tableInfo -> tableInfo.getNamespace().getName().equals(namespaceName))
+        .filter(tableInfo -> tableInfo.getTableType().equals(tableType))
+        .collect(Collectors.toList());
+  }
+
+  /**
+   * This method returns all the tablesInfo list present in the namespace on a universe.
+   *
+   * @param ybService The service to get a YB client from
+   * @param universe The universe to get the table schema information from
+   * @param tableType The table type to filter the tables
+   * @param namespaceId The namespace Id to get the table schema information from
+   * @return A list of {@link MasterDdlOuterClass.ListTablesResponsePB.TableInfo} containing table
+   *     info of the tables in the namespace
+   */
+  public static List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo>
+      getTableInfoListByNamespaceId(
+          YBClientService ybService, Universe universe, TableType tableType, String namespaceId) {
+    return getTableInfoList(ybService, universe).stream()
+        .filter(tableInfo -> tableInfo.getNamespace().getId().toStringUtf8().equals(namespaceId))
+        .filter(tableInfo -> tableInfo.getTableType().equals(tableType))
+        .collect(Collectors.toList());
+  }
+
   public static List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo> getTableInfoList(
       YBClientService ybService, Universe universe, Collection<String> tableIds) {
     return getTableInfoList(ybService, universe).stream()

managed/src/main/java/com/yugabyte/yw/controllers/XClusterConfigController.java

@@ -76,6 +76,7 @@
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.collections4.CollectionUtils;
 import org.yb.CommonTypes;
 import org.yb.CommonTypes.TableType;
 import org.yb.master.MasterDdlOuterClass;
@@ -218,6 +219,14 @@ public Result create(UUID customerUUID, Http.Request request) {
         XClusterConfigTaskBase.getSourceTableIdTargetTableIdMap(
             requestedTableInfoList, targetTableInfoList);
 
+    if (createFormData.bootstrapParams != null
+        && createFormData.bootstrapParams.allowBootstrap
+        && !requestedTableInfoList.isEmpty()) {
+      createFormData.bootstrapParams.tables =
+          getAllBootstrapRequiredTableForXClusterRequestedTable(
+              ybService, requestedTableInfoList, createFormData.tables, sourceUniverse);
+    }
+
     xClusterBootstrappingPreChecks(
         requestedTableInfoList,
         sourceTableInfoList,
@@ -604,6 +613,18 @@ static XClusterConfigTaskParams getSetTablesTaskParams(
           XClusterConfigTaskBase.getSourceTableIdTargetTableIdMap(
               requestedTableInfoList, targetTableInfoList);
 
+      Set<String> tablesInReplication =
+          new HashSet<>(CollectionUtils.union(tableIds, tableIdsToAdd));
+      List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo> allTablesInReplicationInfoList =
+          XClusterConfigTaskBase.getTableInfoList(ybService, sourceUniverse, tablesInReplication);
+      if (bootstrapParams != null
+          && !allTablesInReplicationInfoList.isEmpty()
+          && bootstrapParams.allowBootstrap) {
+        bootstrapParams.tables =
+            getAllBootstrapRequiredTableForXClusterRequestedTable(
+                ybService, allTablesInReplicationInfoList, tablesInReplication, sourceUniverse);
+      }
+
       // We send null as sourceTableIdTargetTableIdMap because add table does not create tables
       // on the target universe through bootstrapping, and the user is responsible to create the
       // same table on the target universe.
@@ -1514,7 +1535,8 @@ public static void xClusterBootstrappingPreChecks(
     Set<String> requestedTableIds = XClusterConfigTaskBase.getTableIds(requestedTableInfoList);
     if (bootstrapParams != null && bootstrapParams.tables != null) {
       // Ensure tables in bootstrapParams is a subset of requestedTableIds.
-      if (!requestedTableIds.containsAll(bootstrapParams.tables)) {
+      if (!bootstrapParams.allowBootstrap
+          && !requestedTableIds.containsAll(bootstrapParams.tables)) {
         throw new IllegalArgumentException(
             String.format(
                 "The set of tables in bootstrapParams (%s) is not a subset of "
@@ -1587,8 +1609,9 @@ public static void xClusterBootstrappingPreChecks(
                                             .equals(namespaceId))
                             .map(tableInfo -> tableInfo.getId().toStringUtf8())
                             .collect(Collectors.toSet());
-                    if (tableIdsInNamespace.size()
-                        != selectedTableIdsInNamespaceToBootstrap.size()) {
+                    if (!bootstrapParams.allowBootstrap
+                        && tableIdsInNamespace.size()
+                            != selectedTableIdsInNamespaceToBootstrap.size()) {
                       throw new IllegalArgumentException(
                           String.format(
                               "For YSQL tables, all the tables in a keyspace must be selected: "
@@ -1600,4 +1623,40 @@ public static void xClusterBootstrappingPreChecks(
       }
     }
   }
+
+  /**
+   * This method retrieves all the tables required for bootstrapping in a cross-cluster setup. It
+   * first checks the type of the tables in the replication info list. If the table type is
+   * PGSQL_TABLE_TYPE, it groups all tables by their namespace ID and adds all table IDs in each
+   * namespace to the bootstrapping list. For YCQL tables, we add all tables in the replication info
+   * list to the bootstrapping list.
+   *
+   * @param ybService The YB client service.
+   * @param tablesInReplicationInfoList A list of all tables in the replication info list.
+   * @param tablesInReplication A set of table IDs that are already in replication.
+   * @param sourceUniverse The source universe of the cross-cluster setup.
+   * @return A set of table IDs that are required for bootstrapping.
+   */
+  public static Set<String> getAllBootstrapRequiredTableForXClusterRequestedTable(
+      YBClientService ybService,
+      List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo> tablesInReplicationInfoList,
+      Set<String> tablesInReplication,
+      Universe sourceUniverse) {
+
+    Set<String> tableIdsForBootstrap = new HashSet<>(tablesInReplication);
+    CommonTypes.TableType tableType = tablesInReplicationInfoList.get(0).getTableType();
+    if (tableType == CommonTypes.TableType.PGSQL_TABLE_TYPE) {
+      XClusterConfigTaskBase.groupByNamespaceId(tablesInReplicationInfoList)
+          .forEach(
+              (namespaceId, tablesInfoList) -> {
+                List<MasterDdlOuterClass.ListTablesResponsePB.TableInfo> namespaceTables =
+                    XClusterConfigTaskBase.getTableInfoListByNamespaceId(
+                        ybService, sourceUniverse, tableType, namespaceId);
+                namespaceTables.stream()
+                    .map(tableInfo -> XClusterConfigTaskBase.getTableId(tableInfo))
+                    .forEach(tableIdsForBootstrap::add);
+              });
+    }
+    return tableIdsForBootstrap;
+  }
 }

managed/src/main/java/com/yugabyte/yw/forms/XClusterConfigCreateFormData.java

@@ -2,6 +2,7 @@
 
 import com.yugabyte.yw.models.XClusterConfig;
 import com.yugabyte.yw.models.XClusterConfig.ConfigType;
+import com.yugabyte.yw.models.common.YbaApi;
 import io.swagger.annotations.ApiModel;
 import io.swagger.annotations.ApiModelProperty;
 import java.util.Set;
@@ -68,6 +69,14 @@ public static class BootstrapParams {
     @ApiModelProperty(value = "Parameters used to do Backup/restore", required = true)
     public BootstarpBackupParams backupRequestParams;
 
+    @ApiModelProperty(
+        value =
+            "WARNING: This is a preview API that could change. Allow backup on whole database when"
+                + " only set of tables require bootstrap",
+        required = false)
+    @YbaApi(visibility = YbaApi.YbaApiVisibility.PREVIEW, sinceYBAVersion = "2.23.0.0")
+    public boolean allowBootstrap = false;
+
     @ApiModel(description = "Backup parameters for bootstrapping")
     @ToString
     public static class BootstarpBackupParams {

managed/src/main/resources/swagger-strict.json

@@ -2830,6 +2830,10 @@
     "BootstrapParams" : {
       "description" : "Bootstrap parameters",
       "properties" : {
+        "allowBootstrap" : {
+          "description" : "WARNING: This is a preview API that could change. Allow backup on whole database when only set of tables require bootstrap",
+          "type" : "boolean"
+        },
         "backupRequestParams" : {
           "$ref" : "#/definitions/BootstarpBackupParams",
           "description" : "Parameters used to do Backup/restore"

managed/src/main/resources/swagger.json

@@ -2846,6 +2846,10 @@
     "BootstrapParams" : {
       "description" : "Bootstrap parameters",
       "properties" : {
+        "allowBootstrap" : {
+          "description" : "WARNING: This is a preview API that could change. Allow backup on whole database when only set of tables require bootstrap",
+          "type" : "boolean"
+        },
         "backupRequestParams" : {
           "$ref" : "#/definitions/BootstarpBackupParams",
           "description" : "Parameters used to do Backup/restore"