Large snapshot generator test

CreateSnapshot/src/main/java/org/opensearch/migrations/CreateSnapshot.java

@@ -84,6 +84,12 @@ public static class Args {
                 description = "The role ARN the cluster will assume to write a snapshot to S3")
         public String s3RoleArn;
 
+        @Parameter(
+                names = {"--s3-endpoint" },
+                required = false,
+                description = "The S3 endpoint URL to use for the S3 bucket, like: s3.us-west-2.amazonaws.com")
+        public String s3Endpoint;
+
         @Parameter(
                 names = {"--index-allowlist"},
                 required = false,
@@ -163,6 +169,7 @@ public void run() {
                 arguments.indexAllowlist,
                 arguments.maxSnapshotRateMBPerNode,
                 arguments.s3RoleArn,
+                arguments.s3Endpoint,
                 context
             );
         }

RFS/src/main/java/org/opensearch/migrations/bulkload/common/S3SnapshotCreator.java

@@ -14,6 +14,7 @@ public class S3SnapshotCreator extends SnapshotCreator {
     private final String s3Region;
     private final Integer maxSnapshotRateMBPerNode;
     private final String snapshotRoleArn;
+    private final String s3Endpoint;
 
     public S3SnapshotCreator(
         String snapshotName,
@@ -24,7 +25,7 @@ public S3SnapshotCreator(
         List<String> indexAllowlist,
         IRfsContexts.ICreateSnapshotContext context
     ) {
-        this(snapshotName, snapshotRepoName, client, s3Uri, s3Region, indexAllowlist, null, null, context);
+        this(snapshotName, snapshotRepoName, client, s3Uri, s3Region, indexAllowlist, null, null, null, context);
     }
 
     public S3SnapshotCreator(
@@ -36,13 +37,15 @@ public S3SnapshotCreator(
         List<String> indexAllowlist,
         Integer maxSnapshotRateMBPerNode,
         String snapshotRoleArn,
+        String s3Endpoint,
         IRfsContexts.ICreateSnapshotContext context
     ) {
         super(snapshotName, snapshotRepoName, indexAllowlist, client, context);
         this.s3Uri = s3Uri;
         this.s3Region = s3Region;
         this.maxSnapshotRateMBPerNode = maxSnapshotRateMBPerNode;
         this.snapshotRoleArn = snapshotRoleArn;
+        this.s3Endpoint = s3Endpoint;
     }
 
     @Override
@@ -60,6 +63,10 @@ public ObjectNode getRequestBodyForRegisterRepo() {
         if (maxSnapshotRateMBPerNode != null) {
             settings.put("max_snapshot_bytes_per_sec", maxSnapshotRateMBPerNode + "mb");
         }
+
+        if (s3Endpoint != null) {
+            settings.put("endpoint", s3Endpoint);
+        }
 
         ObjectNode body = mapper.createObjectNode();
         body.put("type", "s3");

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/environment.py

@@ -96,7 +96,8 @@ def __init__(self, config_file: str):
 
         if 'snapshot' in self.config:
             self.snapshot: Snapshot = get_snapshot(self.config["snapshot"],
-                                                   source_cluster=self.source_cluster)
+                                                   source_cluster=self.source_cluster,
+                                                   target_cluster=self.target_cluster)
             logger.info(f"Snapshot initialized: {self.snapshot}")
         else:
             logger.info("No snapshot provided")

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/middleware/clusters.py

@@ -34,7 +34,7 @@ def connection_check(cluster: Cluster) -> ConnectionResult:
     caught_exception = None
     r = None
     try:
-        r = cluster.call_api(cluster_details_path, timeout=3)
+        r = cluster.call_api(cluster_details_path, timeout=20)
     except Exception as e:
         caught_exception = e
         logging.debug(f"Unable to access cluster: {cluster} with exception: {e}")

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/models/cluster.py

@@ -143,7 +143,7 @@ def _generate_auth_object(self) -> requests.auth.AuthBase | None:
         raise NotImplementedError(f"Auth type {self.auth_type} not implemented")
 
     def call_api(self, path, method: HttpMethod = HttpMethod.GET, data=None, headers=None,
-                 timeout=None, session=None, raise_error=True, **kwargs) -> requests.Response:
+                 timeout=300, session=None, raise_error=True, **kwargs) -> requests.Response:
         """
         Calls an API on the cluster.
         """

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/models/command_runner.py

@@ -16,7 +16,12 @@ def __init__(self, command_root: str, command_args: Dict[str, Any], sensitive_fi
                  run_as_detatched: bool = False, log_file: Optional[str] = None):
         self.command_args = command_args
         self.command = [command_root]
+        if "__positional__" in command_args:
+            self.command.extend(command_args["__positional__"])
+
         for key, value in command_args.items():
+            if key == "__positional__":
+                continue
             self.command.append(key)
             if value is not FlagOnlyArgument:
                 if type(value) is not str:

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/models/factories.py

@@ -46,11 +46,17 @@ def __init__(self, supplied_backfill: str):
         super().__init__("Unsupported backfill type", supplied_backfill)
 
 
-def get_snapshot(config: Dict, source_cluster: Cluster):
+def get_snapshot(config: Dict, source_cluster: Cluster = None, target_cluster: Cluster = None):
+    # Use target_cluster as a fallback when source_cluster is not available
+    cluster_to_use = source_cluster if source_cluster is not None else target_cluster
+
+    if cluster_to_use is None:
+        raise ValueError("Either source_cluster or target_cluster must be provided for snapshot operations")
+
     if 'fs' in config:
-        return FileSystemSnapshot(config, source_cluster)
+        return FileSystemSnapshot(config, cluster_to_use)
     elif 's3' in config:
-        return S3Snapshot(config, source_cluster)
+        return S3Snapshot(config, cluster_to_use)
     logger.error(f"An unsupported snapshot type was provided: {config.keys()}")
     if len(config.keys()) > 1:
         raise UnsupportedSnapshotError(', '.join(config.keys()))