Make remote state reference handling more robust

tests/util/test_util_hooks.py

@@ -60,7 +60,7 @@ class TestGetStateItem:
     def test_get_state_item_from_output_success(self, mock_remote, mock_output):
         mock_output.return_value = '{"key": "value"}'
         result = hooks.get_state_item(
-            "working_dir", {}, "terraform_bin", "state", "item"
+            "working_dir", {}, "terraform_bin", "state", "item", None
         )
         assert result == '{"key": "value"}'
         mock_output.assert_called_once()
@@ -73,7 +73,7 @@ def test_get_state_item_from_output_success(self, mock_remote, mock_output):
     def test_get_state_item_from_remote_success(self, mock_remote, mock_output):
         mock_remote.return_value = '{"key": "value"}'
         result = hooks.get_state_item(
-            "working_dir", {}, "terraform_bin", "state", "item"
+            "working_dir", {}, "terraform_bin", "state", "item", None
         )
         assert result == '{"key": "value"}'
         mock_output.assert_called_once()
@@ -291,7 +291,7 @@ def test_populate_environment_with_terraform_remote_vars(
 
         local_env = {}
         hooks._populate_environment_with_terraform_remote_vars(
-            local_env, "working_dir", "terraform_path", False
+            local_env, "working_dir", "terraform_path", False, None
         )
         assert "TF_REMOTE_LOCAL_KEY" in local_env.keys()
         assert "TF_REMOTE_LOCAL_ANOTHER_KEY" in local_env.keys()

tfworker/backends/base.py

@@ -91,6 +91,19 @@ def remotes(self) -> list:
         """
         pass  # pragma: no cover
 
+    @abstractmethod
+    def get_state(self, remote: str) -> JSONType:
+        """
+        Get the state file from the backend, for a remote.
+
+        Args:
+            deployment (str): The deployment name
+
+        Returns:
+            JSONType: The state from the backend
+        """
+        pass
+
 
 def validate_backend_empty(state: JSONType) -> bool:
     """

tfworker/backends/gcs.py

@@ -8,6 +8,7 @@
 from google.cloud import storage
 from google.cloud.exceptions import Conflict, NotFound
 from tfworker.exceptions import BackendError
+from tfworker.types import JSONType
 
 from .base import BaseBackend, validate_backend_empty
 
@@ -192,3 +193,6 @@ def data_hcl(self, remotes: list) -> str:
             remote_data_config.append("  }")
             remote_data_config.append("}")
         return "\n".join(remote_data_config)
+
+    def get_state(self, remote: str) -> JSONType:
+        raise NotImplementedError

tfworker/backends/s3.py

@@ -10,6 +10,7 @@
 import click
 import tfworker.util.log as log
 from tfworker.exceptions import BackendError
+from tfworker.types import JSONType
 
 from .base import BaseBackend, validate_backend_empty
 
@@ -478,3 +479,34 @@ def _list_bucket_definitions(self) -> set:
         log.trace(f"bucket files: {bucket_files}")
 
         return bucket_files
+
+    def get_state(self, remote: str) -> JSONType:
+        """
+        Retrieve the state item from the backend
+
+        Args:
+            remote (str): The remote name
+            item (str): The item name
+
+        Returns:
+            JSONType: A JSON representation of the state item
+        """
+
+        # check if remote is a valid remote
+        if remote not in self.remotes:
+            raise BackendError(f"remote {remote} not found in backend")
+
+        # get the terraform state file from S3
+        state_key = (
+            f"{self._app_state.root_options.backend_prefix}/{remote}/terraform.tfstate"
+        )
+        log.debug(f"getting state file: {state_key}")
+        state_obj = self._s3_client.get_object(
+            Bucket=self._app_state.root_options.backend_bucket, Key=state_key
+        )
+
+        # load the state file as JSON
+        with closing(state_obj["Body"]) as body:
+            state = json.load(body)
+
+        return state

tfworker/commands/terraform.py

@@ -353,6 +353,7 @@ def _exec_hook(
                 extra_vars=definition.get_template_vars(
                     self.app_state.loaded_config.global_vars.template_vars
                 ),
+                backend=self.app_state.backend,
             )
         except HookError as e:
             log.error(f"hook execution error on definition {definition.name}: \n{e}")