feat: [SNOW-2002186] save password based connection if subsequent key… (#2149)

* feat: [SNOW-2002186] save password based connection if subsequent key pair setup fails

* feat: [SNOW-2002186] display different messages when key pair cannot be set

* feat: [SNOW-2002186] code review fixes

Author: sfc-gh-jwilkowski

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

@@ -5,6 +5,7 @@
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from snowflake.cli._plugins.object.manager import ObjectManager
+from snowflake.cli.api import exceptions
 from snowflake.cli.api.cli_global_context import (
     _CliGlobalContextAccess,
     get_cli_context,
@@ -119,9 +120,7 @@ def _generate_key_pair_and_set_public_key(
         public_key_exists, public_key_2_exists = self._get_public_keys()
 
         if public_key_exists or public_key_2_exists:
-            raise ClickException(
-                "The public key is set already. Use the rotate command instead."
-            )
+            raise exceptions.CouldNotSetKeyPairError()
 
         if not output_path.exists():
             output_path.mkdir(parents=True)

src/snowflake/cli/_plugins/connection/commands.py

@@ -16,8 +16,9 @@
 
 import logging
 import os.path
+from copy import deepcopy
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 
 import typer
 from click import (  # type: ignore
@@ -35,6 +36,7 @@
     strip_if_value_present,
 )
 from snowflake.cli._plugins.object.manager import ObjectManager
+from snowflake.cli.api import exceptions
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.flags import (
     PLAIN_PASSWORD_MSG,
@@ -288,7 +290,11 @@ def add(
         raise UsageError(f"Connection {connection_name} already exists")
 
     if not no_interactive:
-        _extend_add_with_key_pair(connection_name, connection_options)
+        connection_options, keypair_error = _extend_add_with_key_pair(
+            connection_name, connection_options
+        )
+    else:
+        keypair_error = ""
 
     connections_file = add_connection_to_proper_file(
         connection_name,
@@ -297,6 +303,12 @@ def add(
     if set_as_default:
         set_config_value(path=["default_connection_name"], value=connection_name)
 
+    if keypair_error:
+        return MessageResult(
+            f"Wrote new password-based connection {connection_name} to {connections_file}, "
+            f"however there were some issues during key pair setup. Review the following error and check 'snow auth keypair' "
+            f"commands to setup key pair authentication:\n * {keypair_error}"
+        )
     return MessageResult(
         f"Wrote new connection {connection_name} to {connections_file}"
     )
@@ -412,43 +424,57 @@ def _decrypt(passphrase: str | None):
         raise ClickException(str(err))
 
 
-def _extend_add_with_key_pair(connection_name: str, connection_options: Dict):
-    if (
+def _extend_add_with_key_pair(
+    connection_name: str, connection_options: Dict
+) -> Tuple[Dict, str]:
+    if not _should_extend_with_key_pair(connection_options):
+        return connection_options, ""
+
+    configure_key_pair = typer.confirm(
+        "Do you want to configure key pair authentication?",
+        default=False,
+    )
+    if not configure_key_pair:
+        return connection_options, ""
+
+    key_length = typer.prompt(
+        "Key length",
+        default=2048,
+        show_default=True,
+    )
+
+    output_path = typer.prompt(
+        "Output path",
+        default=KEY_PAIR_DEFAULT_PATH,
+        show_default=True,
+        value_proc=lambda value: SecurePath(value),
+    )
+    private_key_passphrase = typer.prompt(
+        "Private key passphrase",
+        default="",
+        hide_input=True,
+        show_default=False,
+        value_proc=lambda value: SecretType(value),
+    )
+    connection = connect_to_snowflake(temporary_connection=True, **connection_options)
+    try:
+        connection_options = AuthManager(connection=connection).extend_connection_add(
+            connection_name=connection_name,
+            connection_options=deepcopy(connection_options),
+            key_length=key_length,
+            output_path=output_path,
+            private_key_passphrase=private_key_passphrase,
+        )
+    except exceptions.CouldNotSetKeyPairError:
+        return connection_options, "The public key is set already."
+    except Exception as e:
+        return connection_options, str(e)
+    return connection_options, ""
+
+
+def _should_extend_with_key_pair(connection_options: Dict) -> bool:
+    return (
         connection_options.get("password") is not None
         and connection_options.get("private_key_file") is None
         and connection_options.get("private_key_path") is None
-    ):
-        configure_key_pair = typer.confirm(
-            "Do you want to configure key pair authentication?",
-            default=False,
-        )
-        if configure_key_pair:
-            key_length = typer.prompt(
-                "Key length",
-                default=2048,
-                show_default=True,
-            )
-
-            output_path = typer.prompt(
-                "Output path",
-                default=KEY_PAIR_DEFAULT_PATH,
-                show_default=True,
-                value_proc=lambda value: SecurePath(value),
-            )
-            private_key_passphrase = typer.prompt(
-                "Private key passphrase",
-                default="",
-                hide_input=True,
-                show_default=False,
-                value_proc=lambda value: SecretType(value),
-            )
-            connection = connect_to_snowflake(
-                temporary_connection=True, **connection_options
-            )
-            AuthManager(connection=connection).extend_connection_add(
-                connection_name=connection_name,
-                connection_options=connection_options,
-                key_length=key_length,
-                output_path=output_path,
-                private_key_passphrase=private_key_passphrase,
-            )
+    )

src/snowflake/cli/api/exceptions.py

@@ -236,3 +236,10 @@ def __init__(self, show_obj_query: str):
         super().__init__(
             f"Received multiple rows from result of SQL statement: {show_obj_query}. Usage of 'show_specific_object' may not be properly scoped."
         )
+
+
+class CouldNotSetKeyPairError(ClickException):
+    def __init__(self):
+        super().__init__(
+            "The public key is set already. Use the rotate command instead."
+        )

tests/test_connection.py

@@ -1546,3 +1546,83 @@ def test_connection_add_no_key_pair_setup_if_no_interactive(
     assert (
         result.output.strip() == f"Wrote new connection conn to {test_snowcli_config}"
     )
+
+
+@mock.patch("snowflake.cli._plugins.auth.keypair.manager.AuthManager.execute_query")
+@mock.patch("snowflake.cli._plugins.object.manager.ObjectManager.execute_query")
+@mock.patch("snowflake.connector.connect")
+def test_connection_add_with_key_pair_saves_password_if_keypair_is_set(
+    mock_connect,
+    mock_object_execute_query,
+    mock_auth_execute_query,
+    runner,
+    tmp_path,
+    mock_cursor,
+    test_snowcli_config,
+):
+    mock_connect.return_value.user = "user"
+    mock_object_execute_query.return_value = mock_cursor(
+        rows=[
+            {"property": "RSA_PUBLIC_KEY", "value": None},
+            {"property": "RSA_PUBLIC_KEY_2", "value": "public key..."},
+        ],
+        columns=[],
+    )
+
+    result = runner.invoke(
+        [
+            "connection",
+            "add",
+        ],
+        input="conn\n"  # connection name: zz
+        "test\n"  # account:
+        "user\n"  # user:
+        "123\n"  # password:
+        "\n"  # role:
+        "\n"  # warehouse:
+        "\n"  # database:
+        "\n"  # schema:
+        "\n"  # host:
+        "\n"  # port:
+        "\n"  # region:
+        "\n"  # authenticator:
+        "\n"  # private key file:
+        "\n"  # token file path:
+        "y\n"  #
+        "\n"  # key_length
+        f"{tmp_path}\n"  # output_path
+        "123\n",  # passphrase
+    )
+
+    private_key_path = tmp_path / "conn.p8"
+    public_key_path = tmp_path / "conn.pub"
+    assert result.exit_code == 0, result.output
+    assert result.output == dedent(
+        f"""\
+        Enter connection name: conn
+        Enter account: test
+        Enter user: user
+        Enter password: 
+        Enter role: 
+        Enter warehouse: 
+        Enter database: 
+        Enter schema: 
+        Enter host: 
+        Enter port: 
+        Enter region: 
+        Enter authenticator: 
+        Enter private key file: 
+        Enter token file path: 
+        Do you want to configure key pair authentication? [y/N]: y
+        Key length [2048]: 
+        Output path [~/.ssh]: {tmp_path}
+        Private key passphrase: 
+        Wrote new password-based connection conn to {test_snowcli_config}, however there were some issues during key pair setup. Review the following error and check 'snow auth keypair' commands to setup key pair authentication:
+         * The public key is set already.
+        """
+    )
+    assert not private_key_path.exists()
+    assert not public_key_path.exists()
+    with open(test_snowcli_config, "r") as f:
+        connections = tomlkit.load(f)
+        assert connections["connections"]["conn"]["password"] == "123"