[Console] Status for RFS Backfill

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

@@ -231,12 +231,18 @@
 
 
 @backfill_group.command(name="status")
+@click.option('--deep-check', is_flag=True, help='Perform a deep status check of the backfill')
 @click.pass_obj
-def status_backfill_cmd(ctx):
-    exitcode, message = logic_backfill.status(ctx.env.backfill)
+def status_backfill_cmd(ctx, deep_check):
+    logger.info(f"Called `console backfill status`, with {deep_check=}")
+    exitcode, message = logic_backfill.status(ctx.env.backfill, deep_check=deep_check)
     if exitcode != ExitCode.SUCCESS:
         raise click.ClickException(message)
-    click.echo(message)
+    if len(message) == 2:
+        click.echo(message[0])
+        click.echo(message[1])
+    else:
+        click.echo(message)
 
 
 # ##################### REPLAY ###################

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/logic/backfill.py

@@ -119,16 +119,16 @@
     return ExitCode.FAILURE, "Backfill scale failed." + "\n" + result.display()
 
 
-def status(backfill: Backfill, *args, **kwargs) -> Tuple[ExitCode, str]:
-    logger.info("Getting backfill status")
+def status(backfill: Backfill, deep_check, *args, **kwargs) -> Tuple[ExitCode, str]:
+    logger.info(f"Getting backfill status with {deep_check=}")
     try:
-        status = backfill.get_status(*args, **kwargs)
+        status = backfill.get_status(deep_check, *args, **kwargs)
     except NotImplementedError:
         logger.error(f"Status is not implemented for backfill {type(backfill).__name__}")
         return ExitCode.FAILURE, f"Status is not implemented for backfill: {type(backfill).__name__}"
     except Exception as e:
         logger.error(f"Failed to get status of backfill: {e}")
         return ExitCode.FAILURE, f"Failure when getting status of backfill: {type(e).__name__} {e}"
-    if status:
+    if status.success:
         return ExitCode.SUCCESS, status.value
-    return ExitCode.FAILURE, "Backfill status retrieval failed." + "\n" + status
+    return ExitCode.FAILURE, "Backfill status retrieval failed." + "\n" + status.value

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

@@ -19,7 +19,7 @@
 }
 
 
-BackfillStatus = Enum("BackfillStatus", ["NOT_STARTED", "RUNNING", "STOPPED", "FAILED"])
+BackfillStatus = Enum("BackfillStatus", ["NOT_STARTED", "STARTING", "RUNNING", "STOPPED", "FAILED"])
 
 
 class Backfill(ABC):
@@ -50,7 +50,7 @@ def stop(self, *args, **kwargs) -> CommandResult:
         pass
 
     @abstractmethod
-    def get_status(self, *args, **kwargs) -> BackfillStatus:
+    def get_status(self, *args, **kwargs) -> CommandResult:
         """Return a status"""
         pass
 

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

@@ -1,7 +1,7 @@
 from console_link.models.osi_utils import (create_pipeline_from_env, start_pipeline, stop_pipeline,
                                            OpenSearchIngestionMigrationProps)
 from console_link.models.cluster import Cluster
-from console_link.models.backfill_base import Backfill, BackfillStatus
+from console_link.models.backfill_base import Backfill
 from console_link.models.command_result import CommandResult
 from typing import Dict
 from cerberus import Validator
@@ -92,7 +92,7 @@ def stop(self, pipeline_name=None):
             pipeline_name = self.osi_props.pipeline_name
         stop_pipeline(osi_client=self.osi_client, pipeline_name=pipeline_name)
 
-    def get_status(self, *args, **kwargs) -> BackfillStatus:
+    def get_status(self, *args, **kwargs) -> CommandResult:
         raise NotImplementedError()
 
     def scale(self, units: int, *args, **kwargs) -> CommandResult:

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

@@ -61,7 +61,7 @@
     def stop(self, *args, **kwargs) -> CommandResult:
         raise NotImplementedError()
 
-    def get_status(self, *args, **kwargs) -> BackfillStatus:
+    def get_status(self, *args, **kwargs) -> CommandResult:
         raise NotImplementedError()
 
     def scale(self, units: int, *args, **kwargs) -> CommandResult:
@@ -96,3 +96,40 @@
     def scale(self, units: int, *args, **kwargs) -> CommandResult:
         logger.info(f"Scaling RFS backfill by setting desired count to {units} instances")
         return self.ecs_client.set_desired_count(units)
+
+    def get_status(self, deep_check, *args, **kwargs) -> CommandResult:
+        logger.info(f"Getting status of RFS backfill, with {deep_check=}")
+        instance_statuses = self.ecs_client.get_instance_statuses()
+        if not instance_statuses:
+            return CommandResult(False, "Failed to get instance statuses")
+
+        status_string = str(instance_statuses)
+        if deep_check:
+            try:
+                shard_status = self._get_detailed_status()
+            except Exception as e:
+                logger.error(f"Failed to get detailed status: {e}")
+            if shard_status:
+                status_string += f"\n{shard_status}"
+
+        if instance_statuses.running > 0:
+            return CommandResult(True, (BackfillStatus.RUNNING, status_string))
+        elif instance_statuses.pending > 0:
+            return CommandResult(True, (BackfillStatus.STARTING, status_string))
+        return CommandResult(True, (BackfillStatus.STOPPED, status_string))
+
+    def _get_detailed_status(self):
+        status_query = {"query": {
+            "bool": {
+                "must": [{"exists": {"field": "expiration"}}],
+                "must_not": [{"exists": {"field": "completedAt"}}]
+            }
+        }}
+        response = self.target_cluster.call_api("/.migrations_working_state/_search", json_body=status_query)
+        r_body = response.json()
+        logger.debug(f"Raw response: {r_body}")
+        if "hits" in r_body:
+            logger.info(f"Hits on detailed status query: {r_body['hits']}")
+            logger.info(f"Sample of remaining shards: {[hit['_id'] for hit in r_body['hits']['hits']]}")
+            return f"Remaining shards: {r_body['hits']['total']['value']}"
+        return response.json()

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

@@ -85,7 +85,7 @@
         elif 'sigv4' in config:
             self.auth_type = AuthMethod.SIGV4
 
-    def call_api(self, path, method: HttpMethod = HttpMethod.GET, timeout=None) -> requests.Response:
+    def call_api(self, path, method: HttpMethod = HttpMethod.GET, timeout=None, json_body=None) -> requests.Response:
         """
         Calls an API on the cluster.
         """
@@ -100,13 +100,21 @@
         else:
             raise NotImplementedError(f"Auth type {self.auth_type} not implemented")
 
+        if json_body is not None:
+            # headers = {'Content-type': 'application/json'}
+            data = json_body
+        else:
+            # headers = None
+            data = None
+
         logger.info(f"Making api call to {self.endpoint}{path}")
         r = requests.request(
             method.name,
             f"{self.endpoint}{path}",
             verify=(not self.allow_insecure),
             auth=auth,
-            timeout=timeout
+            timeout=timeout,
+            json=data
         )
         logger.debug(f"Cluster API call request: {r.request}")
         r.raise_for_status()

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

@@ -1,4 +1,5 @@
 import logging
+from typing import NamedTuple, Optional
 
 import boto3
 
@@ -9,6 +10,15 @@
 logger = logging.getLogger(__name__)
 
 
+class InstanceStatues(NamedTuple):
+    running: int = 0
+    pending: int = 0
+    desired: int = 0
+
+    def __str__(self):
+        return f"Running={self.running}\nPending={self.pending}\nDesired={self.desired}"
+
+
 class ECSService:
     def __init__(self, cluster_name, service_name, aws_region=None):
         self.cluster_name = cluster_name
@@ -25,7 +35,6 @@
             service=self.service_name,
             desiredCount=desired_count
         )
-        logger.debug(f"Response from update_service: {response}")
 
         try:
             raise_for_aws_api_error(response)
@@ -38,3 +47,22 @@
         desired_count = response["service"]["desiredCount"]
         return CommandResult(True, f"Service {self.service_name} set to {desired_count} desired count."
                              f" Currently {running_count} running and {pending_count} pending.")
+
+    def get_instance_statuses(self) -> Optional[InstanceStatues]:
+        logger.info(f"Getting instance statuses for service {self.service_name}")
+        response = self.client.describe_services(
+            cluster=self.cluster_name,
+            services=[self.service_name]
+        )
+        try:
+            raise_for_aws_api_error(response)
+        except AWSAPIError as e:
+            logger.error(f"Error getting instance statuses: {e}")
+            return None
+
+        service = response["services"][0]
+        return InstanceStatues(
+            running=service["runningCount"],
+            pending=service["pendingCount"],
+            desired=service["desiredCount"]
+        )

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

@@ -107,8 +107,7 @@ def create(self, *args, **kwargs) -> CommandResult:
             logger.error(f"Failed to create snapshot: {str(e)}")
             return CommandResult(success=False, value=f"Failed to create snapshot: {str(e)}")
 
-    def status(self, *args, **kwargs) -> CommandResult:
-        deep_check = kwargs.get('deep_check', False)
+    def status(self, *args, deep_check=False, **kwargs) -> CommandResult:
         if deep_check:
             return get_snapshot_status_full(self.source_cluster, self.snapshot_name)
         return get_snapshot_status(self.source_cluster, self.snapshot_name)