From 34b2cb034e048aaeed3064614c19e606c15a37bc Mon Sep 17 00:00:00 2001
From: apollorion <joey@apollorion.com>
Date: Tue, 4 Jun 2024 09:46:37 -0400
Subject: [PATCH] feat: grab contexts via variable extraction

---
 spacemk/exporters/terraform.py | 54 ++++++++++++++++++++++++++++------
 1 file changed, 45 insertions(+), 9 deletions(-)

diff --git a/spacemk/exporters/terraform.py b/spacemk/exporters/terraform.py
index e33add1..e625e78 100644
--- a/spacemk/exporters/terraform.py
+++ b/spacemk/exporters/terraform.py
@@ -312,6 +312,37 @@ def find_variable(data: dict, variable_id: str) -> dict:
 
             return None
 
+        def find_variable_set(data: dict, line: str, valid_sets: list[str]) -> dict | None:
+            # Ensure we're even validating a variable
+            split_line = line.split("=")
+            if len(split_line) != 2:
+                return None
+
+            # Check if the line is a variable set variable
+            set_variables = data.get("variable_set_variables")
+            for set_variable in set_variables:
+                key = set_variable.get("attributes.key")
+                if key == split_line[0]:
+                    if set_variable.get("relationships.varset.data.id") in valid_sets:
+                        return set_variable
+
+            return None
+
+        def _get_attached_variable_sets(data_: dict, workspace_id_: str, organization_id_: str) -> list[str]:
+            valid_sets_ = []
+            sets_ = data.get("variable_sets")
+            for set_ in sets_:
+                workspaces_ = set_.get("relationships.workspaces.data")
+                if workspaces_ is not None:
+                    for workspace_ in workspaces_:
+                        if workspace_.get("id") == workspace_id_:
+                            valid_sets_.append(set_.get("id"))
+
+                if set_.get("relationships.organization.data.id") == organization_id_:
+                    valid_sets_.append(set_.get("id"))
+
+            return list(set(valid_sets_))
+
         if not is_command_available(["docker", "ps"], execute=True):
             logging.warning("Docker is not available. Skipping enriching workspace variables data.")
             return data
@@ -360,6 +391,7 @@ def find_variable(data: dict, variable_id: str) -> dict:
                 agent_container_id = agent_container.id
 
                 for workspace_id, workspace_variables in workspaces.items():
+                    valid_sets = _get_attached_variable_sets(data, workspace_id, organization_id)
                     current_workspace_id = workspace_id
                     current_configuration_version_id = find_workspace(data, workspace_id).get(
                         "relationships.current-configuration-version.data.id"
@@ -458,6 +490,12 @@ def find_variable(data: dict, variable_id: str) -> dict:
                                     if workspace and not workspace.get("attributes.vcs-repo.branch"):
                                         workspace["attributes.vcs-repo.branch"] = branch_name
 
+                                # Check if this line corresponds to a variable set attached to this workspace
+                                # If it is, we're going to assume this is the correct value for the variable set
+                                variable_set = find_variable_set(data, line, valid_sets)
+                                if variable_set is not None:
+                                    variable_set["attributes.value"] = line.split("=")[1]
+
                     self._restore_workspace_exec_mode(organization_id, workspace_id, workspace_data_backup)
                     restored_agent = True
 
@@ -576,11 +614,11 @@ def _extract_data(self) -> list[dict]:
             data["providers"].extend(self._extract_providers_data(organization))
             data["tasks"].extend(self._extract_tasks_data(organization))
             data["teams"].extend(self._extract_teams_data(organization))
-            # data["variable_sets"].extend(self._extract_variable_sets_data(organization))
+            data["variable_sets"].extend(self._extract_variable_sets_data(organization))
             data["workspaces"].extend(self._extract_workspaces_data(organization))
 
-        # for variable_set in data.variable_sets:
-        #     data["variable_set_variables"].extend(self._extract_variable_set_variables_data(variable_set))
+        for variable_set in data.variable_sets:
+            data["variable_set_variables"].extend(self._extract_variable_set_variables_data(variable_set))
 
         for workspace in data.workspaces:
             data["workspace_variables"].extend(self._extract_workspace_variables_data(workspace))
@@ -1326,12 +1364,10 @@ def _map_data(self, src_data: dict) -> dict:
         data = benedict(
             {
                 "spaces": self._map_spaces_data(src_data),  # Must be first due to dependency
-                "contexts": [],
-                "context_variables": [],
-                # "contexts": self._map_contexts_data(src_data),
-                # "context_variables": self._map_context_variables_data(
-                #     src_data
-                # ),  # Must be after contexts due to dependency
+                "contexts": self._map_contexts_data(src_data),
+                "context_variables": self._map_context_variables_data(
+                    src_data
+                ),  # Must be after contexts due to dependency
                 "modules": self._map_modules_data(src_data),
                 "stacks": self._map_stacks_data(src_data),
                 "stack_variables": self._map_stack_variables_data(src_data),  # Must be after stacks due to dependency