Auth rotate only overwrite connection (#2142)

Author: sfc-gh-astus

src/snowflake/cli/_plugins/auth/keypair/commands.py

@@ -100,8 +100,6 @@ def setup(
 
 @app.command("rotate", requires_connection=True)
 def rotate(
-    new_connection: bool = _new_connection_option,
-    connection_name: str = _connection_name_option,
     key_length: int = _key_length_option,
     output_path: Path = _output_path_option,
     private_key_passphrase: SecretType = _private_key_passphrase_option,
@@ -111,7 +109,6 @@ def rotate(
     Rotates the key for the connection. Generates the key pair, sets the public key for the user in Snowflake and creates or updates the connection.
     """
     AuthManager().rotate(
-        connection_name=connection_name,
         key_length=key_length,
         output_path=SecurePath(output_path),
         private_key_passphrase=private_key_passphrase,

src/snowflake/cli/_plugins/auth/keypair/manager.py

@@ -48,39 +48,31 @@ def setup(
         if not connection_name:
             connection_name = cli_context.connection_context.connection_name
 
+        key_name = AuthManager._get_free_key_name(output_path, connection_name)  # type: ignore[arg-type]
         self._generate_key_pair_and_set_public_key(
             user=cli_context.connection.user,
             key_length=key_length,
             output_path=output_path,
-            connection_name=connection_name,  # type: ignore[arg-type]
+            key_name=key_name,  # type: ignore[arg-type]
             private_key_passphrase=private_key_passphrase,
         )
 
         self._create_or_update_connection(
             current_connection=cli_context.connection_context.connection_name,
             connection_name=connection_name,  # type: ignore[arg-type]
             private_key_path=self._get_private_key_path(
-                output_path=output_path, key_name=connection_name  # type: ignore[arg-type]
+                output_path=output_path, key_name=key_name  # type: ignore[arg-type]
             ),
         )
 
     def rotate(
         self,
-        connection_name: str,
         key_length: int,
         output_path: SecurePath,
         private_key_passphrase: SecretType,
     ):
-        # When the user provide new connection name
-        if connection_name and connection_exists(connection_name):
-            raise ClickException(
-                f"Connection with name {connection_name} already exists."
-            )
-
         cli_context = get_cli_context()
-        # When the use not provide connection name, so we overwrite the current connection
-        if not connection_name:
-            connection_name = cli_context.connection_context.connection_name
+        connection_name = cli_context.connection_context.connection_name
 
         self._ensure_connection_has_private_key(
             cli_context.connection_context.connection_name
@@ -98,10 +90,11 @@ def rotate(
                 public_key_2,
             )
 
+        key_name = AuthManager._get_free_key_name(output_path, connection_name)
         public_key = self._generate_keys_and_return_public_key(
             key_length=key_length,
             output_path=output_path,
-            key_name=connection_name,
+            key_name=key_name,
             private_key_passphrase=private_key_passphrase,
         )
         self.set_public_key(
@@ -111,7 +104,7 @@ def rotate(
             current_connection=cli_context.connection_context.connection_name,
             connection_name=connection_name,
             private_key_path=self._get_private_key_path(
-                output_path=output_path, key_name=connection_name
+                output_path=output_path, key_name=key_name
             ),
         )
 
@@ -120,7 +113,7 @@ def _generate_key_pair_and_set_public_key(
         user: str,
         key_length: int,
         output_path: SecurePath,
-        connection_name: str,
+        key_name: str,
         private_key_passphrase: SecretType,
     ):
         public_key_exists, public_key_2_exists = self._get_public_keys()
@@ -136,7 +129,7 @@ def _generate_key_pair_and_set_public_key(
         public_key = self._generate_keys_and_return_public_key(
             key_length=key_length,
             output_path=output_path,
-            key_name=connection_name,  # type: ignore[arg-type]
+            key_name=key_name,  # type: ignore[arg-type]
             private_key_passphrase=private_key_passphrase,
         )
         self.set_public_key(user, PublicKeyProperty.RSA_PUBLIC_KEY, public_key)
@@ -192,19 +185,19 @@ def extend_connection_add(
         output_path: SecurePath,
         private_key_passphrase: SecretType,
     ) -> Dict:
+        key_name = AuthManager._get_free_key_name(output_path, connection_name)
+
         self._generate_key_pair_and_set_public_key(
             user=connection_options["user"],
             key_length=key_length,
             output_path=output_path,
-            connection_name=connection_name,
+            key_name=key_name,
             private_key_passphrase=private_key_passphrase,
         )
 
         connection_options["authenticator"] = "SNOWFLAKE_JWT"
         connection_options["private_key_file"] = str(
-            self._get_private_key_path(
-                output_path=output_path, key_name=connection_name
-            ).path
+            self._get_private_key_path(output_path=output_path, key_name=key_name).path
         )
         if connection_options.get("password"):
             del connection_options["password"]
@@ -293,6 +286,30 @@ def _generate_keys_and_return_public_key(
 
         return public_pem.decode("utf-8")
 
+    @staticmethod
+    def _get_free_key_name(output_path: SecurePath, key_name: str) -> str:
+        new_private_key = f"{key_name}.p8"
+        new_public_key = f"{key_name}.pub"
+        new_key_name = key_name
+        counter = 1
+
+        while (
+            (output_path / new_private_key).exists()
+            and (output_path / new_public_key).exists()
+            and counter <= 100
+        ):
+            new_key_name = f"{key_name}_{counter}"
+            new_private_key = f"{new_key_name}.p8"
+            new_public_key = f"{new_key_name}.pub"
+            counter += 1
+
+        if counter == 100:
+            raise ClickException(
+                "Too many key pairs with the same name in the output directory."
+            )
+
+        return new_key_name
+
     @staticmethod
     def _get_private_key_path(output_path: SecurePath, key_name: str) -> SecurePath:
         return (output_path / f"{key_name}.p8").resolve()

tests/auth/__snapshots__/test_auth.ambr

@@ -42,22 +42,8 @@
 
   '''
 # ---
-# name: test_rotate_connection_already_exists
-  '''
-  Create a new connection? [Y/n]: 
-  Enter connection name: default
-  Enter key length [2048]: 
-  Enter private key passphrase []: 
-  +- Error ----------------------------------------------------------------------+
-  | Connection with name default already exists.                                 |
-  +------------------------------------------------------------------------------+
-  
-  '''
-# ---
 # name: test_rotate_create_output_directory_with_proper_privileges
   '''
-  Create a new connection? [Y/n]: Y
-  Enter connection name: keypairconnection
   Enter key length [2048]: 4096
   Enter private key passphrase []: 
   Rotate completed.
@@ -66,29 +52,23 @@
 # ---
 # name: test_rotate_no_prompts
   '''
-  Create a new connection? [Y/n]: 
-  Enter connection name: keypairconnection
   Set the `PRIVATE_KEY_PASSPHRASE` environment variable before using the connection.
   Rotate completed.
 
   '''
 # ---
 # name: test_rotate_no_public_key_set
   '''
-  Create a new connection? [Y/n]: 
-  Enter connection name: keypairconnection
   Enter key length [2048]: 
   Enter private key passphrase []: 
   +- Error ----------------------------------------------------------------------+
-  | Connection with name keypairconnection already exists.                       |
+  | No public key found. Use the setup command first.                            |
   +------------------------------------------------------------------------------+
 
   '''
 # ---
 # name: test_rotate_only_public_key_set
   '''
-  Create a new connection? [Y/n]: Y
-  Enter connection name: keypairconnection
   Enter key length [2048]: 4096
   Enter private key passphrase []: 
   Rotate completed.
@@ -97,8 +77,6 @@
 # ---
 # name: test_rotate_other_public_key_set_options[KEY-KEY]
   '''
-  Create a new connection? [Y/n]: Y
-  Enter connection name: keypairconnection
   Enter key length [2048]: 4096
   Enter private key passphrase []: 
   Rotate completed.
@@ -107,27 +85,14 @@
 # ---
 # name: test_rotate_other_public_key_set_options[None-KEY]
   '''
-  Create a new connection? [Y/n]: Y
-  Enter connection name: keypairconnection
   Enter key length [2048]: 4096
   Enter private key passphrase []: 
   Rotate completed.
 
   '''
 # ---
-# name: test_rotate_overwrite_connection
-  '''
-  Create a new connection? [Y/n]: n
-  Enter key length [2048]: 
-  Enter private key passphrase []: 
-  Rotate completed.
-  
-  '''
-# ---
 # name: test_rotate_with_password
   '''
-  Create a new connection? [Y/n]: Y
-  Enter connection name: keypairconnection
   Enter key length [2048]: 4096
   Enter private key passphrase []: 
   Set the `PRIVATE_KEY_PASSPHRASE` environment variable before using the connection.