From d69c3051bc6fe7493c1e51ffc63d97208fdf3e50 Mon Sep 17 00:00:00 2001
From: Tomasz Urbaszek <tomasz.urbaszek@snowflake.com>
Date: Tue, 11 Jun 2024 16:36:09 +0200
Subject: [PATCH] Add --default to add command to mark connection as default
 (#1186)

---
 RELEASE-NOTES.md                              |  3 ++-
 .../cli/plugins/connection/commands.py        | 11 +++++++++
 tests/__snapshots__/test_connection.ambr      | 10 ++++++++
 tests/__snapshots__/test_help_messages.ambr   |  2 ++
 tests/test_connection.py                      | 23 +++++++++++++++++++
 5 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 08dbe9315..a124e04ce 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -26,7 +26,8 @@
   * `snow notebook execute` enabling head-less execution of a notebook.
   * `snow notebook create` proving an option to create a Snowflake Notebook from a file on stage.
 * Added templating support for project definition file.
-  * Template variables can now be used anywhere in the the project definition file.
+  * Template variables can now be used anywhere in the project definition file.
+* Added `--default` flag to `snow connection add` commands allowing users to mark the new connection as default.
 
 ## Fixes and improvements
 * Fixed error handling for malformatted `config.toml`
diff --git a/src/snowflake/cli/plugins/connection/commands.py b/src/snowflake/cli/plugins/connection/commands.py
index faa8185b4..9230e9e45 100644
--- a/src/snowflake/cli/plugins/connection/commands.py
+++ b/src/snowflake/cli/plugins/connection/commands.py
@@ -218,6 +218,12 @@ def add(
         prompt="Path to private key file",
         help="Path to file containing private key",
     ),
+    set_as_default: bool = typer.Option(
+        False,
+        "--default",
+        is_flag=True,
+        help="If provided the connection will be configured as default connection.",
+    ),
     **options,
 ) -> CommandResult:
     """Adds a connection to configuration file."""
@@ -241,6 +247,11 @@ def add(
             private_key_path=private_key_path,
         ),
     )
+    if set_as_default:
+        set_config_value(
+            section=None, key="default_connection_name", value=connection_name
+        )
+
     return MessageResult(
         f"Wrote new connection {connection_name} to {CONFIG_MANAGER.file_path}"
     )
diff --git a/tests/__snapshots__/test_connection.ambr b/tests/__snapshots__/test_connection.ambr
index cc7708ef5..0378a7001 100644
--- a/tests/__snapshots__/test_connection.ambr
+++ b/tests/__snapshots__/test_connection.ambr
@@ -26,6 +26,16 @@
   
   '''
 # ---
+# name: test_new_connection_can_be_added_as_default
+  '''
+  default_connection_name = "default-conn"
+  [connections.default-conn]
+  account = "account1"
+  user = "user1"
+  password = "password1"
+  
+  '''
+# ---
 # name: test_new_connection_with_jwt_auth
   '''
   [connections.conn2]
diff --git a/tests/__snapshots__/test_help_messages.ambr b/tests/__snapshots__/test_help_messages.ambr
index 2f6103989..f5a20ba5e 100644
--- a/tests/__snapshots__/test_help_messages.ambr
+++ b/tests/__snapshots__/test_help_messages.ambr
@@ -901,6 +901,8 @@
   │                                        [default: optional]                   │
   │ --private-key            -k      TEXT  Path to file containing private key   │
   │                                        [default: optional]                   │
+  │ --default                              If provided the connection will be    │
+  │                                        configured as default connection.     │
   │ --help                                 Show this message and exit.           │
   ╰──────────────────────────────────────────────────────────────────────────────╯
   ╭─ Global configuration ───────────────────────────────────────────────────────╮
diff --git a/tests/test_connection.py b/tests/test_connection.py
index cf52e9536..65ee6bc7a 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -48,6 +48,29 @@ def test_new_connection_can_be_added(runner, snapshot):
     assert content == snapshot
 
 
+def test_new_connection_can_be_added_as_default(runner, snapshot):
+    with NamedTemporaryFile("w+", suffix=".toml") as tmp_file:
+        result = runner.invoke_with_config_file(
+            tmp_file.name,
+            [
+                "connection",
+                "add",
+                "--connection-name",
+                "default-conn",
+                "--username",
+                "user1",
+                "--password",
+                "password1",
+                "--account",
+                "account1",
+                "--default",
+            ],
+        )
+        content = tmp_file.read()
+    assert result.exit_code == 0, result.output
+    assert content == snapshot
+
+
 def test_new_connection_with_jwt_auth(runner, snapshot):
     with NamedTemporaryFile("w+", suffix=".toml") as tmp_file:
         result = runner.invoke_with_config_file(