From 5868e28e7e617811b0975f243516ff0520c6a6ec Mon Sep 17 00:00:00 2001
From: Bruno Antonellini <antonellinibruno@gmail.com>
Date: Tue, 10 Sep 2024 15:27:55 -0300
Subject: [PATCH] DCV-2893 dbt-coves setup: add new --template-url (#497)

* DCV-2893 dbt-coves setup: add new --setup-template-url flag

* Rename to --template-url
---
 dbt_coves/config/config.py    | 2 ++
 dbt_coves/tasks/setup/main.py | 6 +++++-
 dbt_coves/utils/flags.py      | 6 +++---
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/dbt_coves/config/config.py b/dbt_coves/config/config.py
index 05c4687e..c280c69c 100644
--- a/dbt_coves/config/config.py
+++ b/dbt_coves/config/config.py
@@ -132,6 +132,7 @@ class LoadModel(BaseModel):
 class SetupModel(BaseModel):
     no_prompt: Optional[bool] = False
     quiet: Optional[bool] = False
+    template_url: Optional[str] = "https://github.com/datacoves/setup_template.git"
 
 
 class RunDbtModel(BaseModel):
@@ -241,6 +242,7 @@ class DbtCovesConfig:
         "load.airbyte.secrets_key",
         "setup.no_prompt",
         "setup.quiet",
+        "setup.template_url",
         "dbt.command",
         "dbt.project_dir",
         "dbt.virtualenv",
diff --git a/dbt_coves/tasks/setup/main.py b/dbt_coves/tasks/setup/main.py
index 646e2911..38810f4f 100644
--- a/dbt_coves/tasks/setup/main.py
+++ b/dbt_coves/tasks/setup/main.py
@@ -53,6 +53,10 @@ def register_parser(cls, sub_parsers, base_subparser):
             help="Skip rendering results",
             default=False,
         )
+        ext_subparser.add_argument(
+            "--template-url",
+            help="URL to the setup template repository",
+        )
         ext_subparser.set_defaults(cls=cls, which="setup")
         cls.arg_parser = ext_subparser
         return ext_subparser
@@ -136,7 +140,7 @@ def setup_datacoves(self):
         for service in services:
             self.copier_context[service] = True
         copier.run_auto(
-            src_path="git@github.com:datacoves/setup_template.git",
+            src_path=self.get_config_value("template_url"),
             dst_path=self.repo_path,
             data=self.copier_context,
             quiet=self.get_config_value("quiet"),
diff --git a/dbt_coves/utils/flags.py b/dbt_coves/utils/flags.py
index b6535870..928b4e3b 100644
--- a/dbt_coves/utils/flags.py
+++ b/dbt_coves/utils/flags.py
@@ -133,9 +133,7 @@ def __init__(self, cli_parser: ArgumentParser) -> None:
             "template": "https://github.com/datacoves/cookiecutter-dbt.git",
             "current-dir": False,
         }
-        self.setup = {
-            "no_prompt": False,
-        }
+        self.setup = {"no_prompt": False, "template_url": None}
         self.dbt = {"command": None, "project_dir": None, "virtualenv": None, "cleanup": False}
         self.data_sync = {"redshift": {"tables": []}, "snowflake": {"tables": []}}
         self.blue_green = {
@@ -395,6 +393,8 @@ def parse_args(self, cli_args: List[str] = list()) -> None:
                     self.setup["no_prompt"] = self.args.no_prompt
                 if self.args.quiet:
                     self.setup["quiet"] = self.args.quiet
+                if self.args.template_url:
+                    self.setup["template_url"] = self.args.template_url
 
             # run dbt
             if self.args.cls.__name__ == "RunDbtTask":