From d6f0416bdb26cb0462372c75d0f540a8cc8dc518 Mon Sep 17 00:00:00 2001
From: apollorion <joey@apollorion.com>
Date: Tue, 6 Aug 2024 12:35:39 -0400
Subject: [PATCH] feat: support gitlab modules

---
 config.yml.example              |  5 +++++
 spacemk/commands/generate.py    |  3 +++
 spacemk/exporters/terraform.py  | 17 ++++++++++++++---
 spacemk/templates/base.tf.jinja |  9 +++++++++
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/config.yml.example b/config.yml.example
index dc1491c..7a461bd 100644
--- a/config.yml.example
+++ b/config.yml.example
@@ -28,6 +28,11 @@ generator:
 #  # Default: N/A, must be provided (however, not always used. If it is used and not provided an exception will be raised during code generation)
 #  custom_runner_image: public.ecr.aws/spacelift/runner-terraform
 
+#  # The default branch for all modules. Only used if the module does not have a branch specified.
+#  # Default: empty string
+#  modules:
+#    default_branch: main
+
   extra_vars:
     foo: bar # "{{ extra_vars.foo }}" in a template will be replaced by "bar"
 
diff --git a/spacemk/commands/generate.py b/spacemk/commands/generate.py
index 936f6a8..7624617 100644
--- a/spacemk/commands/generate.py
+++ b/spacemk/commands/generate.py
@@ -24,6 +24,9 @@ def default(value, _default):
         },
         "custom_runner_image": default(config.get("generator.custom_runner_image"),
                                        "SPACELIFT_DEFAULT_INVALID"),
+        "modules": {
+            "default_branch": default(config.get("generator.modules.default_branch"), ""),
+        }
     }
 
     generator = Generator()
diff --git a/spacemk/exporters/terraform.py b/spacemk/exporters/terraform.py
index 209b30a..3d55aa2 100644
--- a/spacemk/exporters/terraform.py
+++ b/spacemk/exporters/terraform.py
@@ -37,6 +37,8 @@ def __init__(self, config: dict):
             }
         }
 
+        self.is_gitlab = False
+
     def _build_stack_slug(self, workspace: dict) -> str:
         return slugify(workspace.get("attributes.name"))
 
@@ -1114,9 +1116,15 @@ def _map_contexts_data(self, src_data: dict) -> dict:
     def _map_modules_data(self, src_data: dict) -> dict:
         logging.info("Start mapping modules data")
 
+        vcs_provider = "gitlab" if self.is_gitlab else "github_custom"
+
         data = []
         for module in src_data.get("modules"):
-            if module.get("attributes.vcs-repo.identifier"):
+            if self.is_gitlab and module.get("attributes.vcs-repo.identifier"):
+                segments = module.get("attributes.vcs-repo.identifier").split("/")
+                vcs_namespace = "/".join(segments[:-1])
+                vcs_repository = segments[-1]
+            elif module.get("attributes.vcs-repo.identifier"):
                 segments = module.get("attributes.vcs-repo.identifier").split("/")
                 vcs_namespace = segments[0]
                 vcs_repository = segments[1]
@@ -1140,7 +1148,7 @@ def _map_modules_data(self, src_data: dict) -> dict:
                     "vcs": {
                         "branch": module.get("attributes.vcs-repo.branch"),
                         "namespace": vcs_namespace,
-                        "provider": "github_custom",  # KLUDGE: TFC/TFE does not provide that information
+                        "provider": vcs_provider,
                         "repository": vcs_repository,
                     },
                 }
@@ -1277,6 +1285,7 @@ def find_workspace_variable_with_invalid_name(data: dict, workspace_id: str, typ
                 raise ValueError(f"Unknown VCS provider name ({provider})")
 
             if provider == "gitlab" and workspace.get("attributes.vcs-repo.identifier"):
+                self.is_gitlab = True
                 segments = workspace.get("attributes.vcs-repo.identifier").split("/")
                 vcs_namespace = "/".join(segments[:-1])
                 vcs_repository = segments[-1]
@@ -1343,8 +1352,10 @@ def _map_data(self, src_data: dict) -> dict:
                 # "context_variables": self._map_context_variables_data(
                 #     src_data
                 # ),  # Must be after contexts due to dependency
-                "modules": self._map_modules_data(src_data),
+
+                # Stacks must be before modules so we can determine is_gitlab so we set VCS properly
                 "stacks": self._map_stacks_data(src_data),
+                "modules": self._map_modules_data(src_data),
                 "stack_variables": self._map_stack_variables_data(src_data),  # Must be after stacks due to dependency
             }
         )
diff --git a/spacemk/templates/base.tf.jinja b/spacemk/templates/base.tf.jinja
index f8a1831..c0a45d1 100644
--- a/spacemk/templates/base.tf.jinja
+++ b/spacemk/templates/base.tf.jinja
@@ -258,7 +258,12 @@ resource "spacelift_environment_variable" "{{ variable._relationships.space._mig
 {% for module in modules %}
 {% if module.status == "setup_complete" and module.visibility == "private" %}
 resource "spacelift_module" "{{ module._relationships.space._migration_id }}_{{ module._migration_id }}" {
+
+  {%  if  module.vcs.branch == ""%}
+  {{ argument("branch", generation_config.modules.default_branch, required=True) }}
+  {% else %}
   {{ argument("branch", module.vcs.branch, required=True) }}
+  {% endif %}
   {{ argument("name", module.name) }}
   {{ argument("repository", module.vcs.repository, required=True) }}
   {{ argument("space_id", "spacelift_space." ~ module._relationships.space._migration_id ~ ".id", serialize=False) }}
@@ -271,6 +276,10 @@ resource "spacelift_module" "{{ module._relationships.space._migration_id }}_{{
     {{ argument("namespace", module.vcs.namespace) }}
   }
   {% endif %}
+  {% elif module.vcs.provider %}
+  {{ module.vcs.provider }} {
+  {{ argument("namespace", module.vcs.namespace) }}
+  }
   {% endif %}
 }
 {% block module_extra scoped %}{% endblock %}