diff --git a/Babylon/commands/namespace/get_all_states.py b/Babylon/commands/namespace/get_all_states.py index 55a309f4..2c12029a 100644 --- a/Babylon/commands/namespace/get_all_states.py +++ b/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() diff --git a/Babylon/utils/environment.py b/Babylon/utils/environment.py index 4fc0377d..80df274b 100644 --- a/Babylon/utils/environment.py +++ b/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 { diff --git a/tests/e2e/test_e2e.sh b/tests/e2e/test_e2e.sh index bc427a4f..1f8a3c3a 100755 --- a/tests/e2e/test_e2e.sh +++ b/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 diff --git a/tests/integration/test_api_endpoints.sh b/tests/integration/test_api_endpoints.sh index 57372755..c9a8c50f 100755 --- a/tests/integration/test_api_endpoints.sh +++ b/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 diff --git a/tests/unit/test_macro.py b/tests/unit/test_macro.py index 8714b1ce..19f27254 100644 --- a/tests/unit/test_macro.py +++ b/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",))