Mtor/feat show state

Babylon/commands/namespace/get_all_states.py

@@ -1,6 +1,6 @@
 from logging import getLogger
 
-from click import command
+from click import Choice, argument, command, echo, style
 
 from Babylon.utils.environment import Environment
 from Babylon.utils.response import CommandResponse
@@ -10,13 +10,37 @@
 
 
 @command()
-def get_states() -> CommandResponse:
-    """Display all states in your local machine"""
-    states_dir = env.state_dir
-    if not env.state_dir.exists():
-        logger.error(f"directory {env.state_dir} not found")
-        return CommandResponse.fail()
-    states_files = sorted(states_dir.glob("state.*.yaml"))
-    for f in states_files:
-        print(f" {f.name}")
-    return CommandResponse.success()
+@argument("target", type=Choice(["local", "remote"], case_sensitive=False))
+def get_states(target: str) -> CommandResponse:
+    """Display states from local machine or Azure remote storage."""
+
+    results_found = False
+    if target == "local":
+        echo(style("\n 📂 Local States", bold=True, fg="cyan"))
+        states_dir = env.state_dir
+
+        if not states_dir.exists():
+            logger.error(f"  [bold red]✘[/bold red] Directory not found: [dim]{states_dir}[/dim]")
+        else:
+            local_files = sorted(states_dir.glob("state.*.yaml"))
+            if not local_files:
+                logger.warning("  [yellow]⚠[/yellow] No local state files found")
+            else:
+                for f in local_files:
+                    echo(style("  • ", fg="green") + f.name)
+                    results_found = True
+
+    elif target == "remote":
+        echo(style("\n ☁️  Remote States", bold=True, fg="cyan"))
+        try:
+            remote_files = env.list_remote_states()
+            if not remote_files:
+                logger.warning("  [yellow]⚠[/yellow] No remote states found on Azure")
+            else:
+                for name in sorted(remote_files):
+                    echo(style("  • ", fg="green") + name)
+                    results_found = True
+        except Exception as e:
+            logger.error(f"  [bold red]✘[/bold red] Failed to reach Azure storage: {e}")
+
+    return CommandResponse.success() if results_found else CommandResponse.fail()

Babylon/utils/environment.py

@@ -46,6 +46,9 @@ def __call__(cls, *args, **kwargs):
 
 
 class Environment(metaclass=SingletonMeta):
+    # Azure Blob Storage configuration
+    STATE_CONTAINER = "babylon-states"
+
     def __init__(self):
         self.remote = False
         self.pwd = Path.cwd()
@@ -68,6 +71,10 @@ def __init__(self):
         self.working_dir = WorkingDir(working_dir_path=self.pwd)
         self.variable_files: list[Path] = []
 
+    def _get_state_blob_client(self, blob_name: str):
+        """Get a blob client for state management"""
+        return self.blob_client.get_blob_client(container=self.STATE_CONTAINER, blob=blob_name)
+
     def get_variables(self):
         merged_data, duplicate_keys = self.merge_yaml_files(self.variable_files)
         if len(duplicate_keys) > 0:
@@ -179,14 +186,25 @@ def store_state_in_local(self, state: dict):
 
     def store_state_in_cloud(self, state: dict):
         state_file = f"state.{self.context_id}.{self.environ_id}.{self.state_id}.yaml"
-        state_container = self.blob_client.get_container_client(container="babylon-states")
+        state_container = self.blob_client.get_container_client(container=self.STATE_CONTAINER)
         if not state_container.exists():
             state_container.create_container()
-        state_blob = self.blob_client.get_blob_client(container="babylon-states", blob=state_file)
+        state_blob = self._get_state_blob_client(state_file)
         if state_blob.exists():
             state_blob.delete_blob()
         state_blob.upload_blob(data=dump(state).encode("utf-8"))
 
+    def list_remote_states(self) -> list[str]:
+        """List state file names present in the Azure blob container."""
+        try:
+            self.set_blob_client()
+            container_client = self.blob_client.get_container_client(container=self.STATE_CONTAINER)
+            blobs = container_client.list_blobs(name_starts_with="state.")
+            return [b.name for b in blobs if b.name.endswith(".yaml")]
+        except Exception as e:
+            logger.error(f"  [bold red]✘[/bold red] Failed to list remote states: {e}")
+            return []
+
     def get_state_from_local(self):
         state_file = self.state_dir / f"state.{self.context_id}.{self.environ_id}.{self.state_id}.yaml"
         if not state_file.exists():
@@ -212,7 +230,7 @@ def get_state_from_local(self):
 
     def get_state_from_cloud(self) -> dict:
         s = f"state.{self.context_id}.{self.environ_id}.{self.state_id}.yaml"
-        state_blob = self.blob_client.get_blob_client(container="babylon-states", blob=s)
+        state_blob = self._get_state_blob_client(s)
         exists = state_blob.exists()
         if not exists:
             return {

tests/e2e/test_e2e.sh

@@ -14,7 +14,7 @@ export TENANT="sphinx"
 export STATE="teststate"
 
 babylon namespace use -c ${CONTEXT} -t ${TENANT} -s $STATE
-babylon namespace get-states
+babylon namespace get-states local
 babylon namespace get-contexts
 
 # Get version

tests/integration/test_api_endpoints.sh

@@ -16,7 +16,7 @@ export TENANT="sphinx"
 export STATE="teststate"
 
 babylon namespace use -c ${CONTEXT} -t ${TENANT} -s $STATE
-babylon namespace get-states
+babylon namespace get-states local
 babylon namespace get-contexts
 
 # Get version

tests/unit/test_macro.py

@@ -82,7 +82,6 @@ def test_resolve_inclusion_exclusion_include_duplicates():
         False,
     )
 
-
 def test_resolve_inclusion_exclusion_invalid_exclude():
     with pytest.raises(Abort):
         resolve_inclusion_exclusion(include=(), exclude=("invalid",))