feat: add snowflake keypair authentication (#21322)

docs/docs/databases/snowflake.mdx

@@ -29,3 +29,31 @@ user/role rights during engine creation by default. However, when pressing the 
 button in the Create or Edit Database dialog, user/role credentials are validated by passing
 “validate_default_parameters”: True to the connect() method during engine creation. If the user/role
 is not authorized to access the database, an error is recorded in the Superset logs.
+
+And if you want connect Snowflake with [Key Pair Authentication](https://docs.snowflake.com/en/user-guide/key-pair-auth.html#step-6-configure-the-snowflake-client-to-use-key-pair-authentication).
+Plase make sure you have the key pair and the public key is registered in Snowflake.
+To connect Snowflake with Key Pair Authentication, you need to add the following parameters to "SECURE EXTRA" field.
+
+***Please note that you need to merge multi-line private key content to one line and insert `\n` between each line***
+
+```
+{
+     "auth_method": "keypair",
+     "auth_params": {
+         "privatekey_body": "-----BEGIN ENCRYPTED PRIVATE KEY-----\n...\n...\n-----END ENCRYPTED PRIVATE KEY-----",
+         "privatekey_pass":"Your Private Key Password"
+     }
+ }
+```
+
+If your private key is stored on server, you can replace "privatekey_body" with “privatekey_path” in parameter.
+
+```
+{
+    "auth_method": "keypair",
+    "auth_params": {
+        "privatekey_path":"Your Private Key Path",
+        "privatekey_pass":"Your Private Key Password"
+    }
+}
+```

superset/db_engine_specs/snowflake.py

@@ -15,13 +15,17 @@
 # specific language governing permissions and limitations
 # under the License.
 import json
+import logging
 import re
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Pattern, Tuple, TYPE_CHECKING
 from urllib import parse
 
 from apispec import APISpec
 from apispec.ext.marshmallow import MarshmallowPlugin
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from flask import current_app
 from flask_babel import gettext as __
 from marshmallow import fields, Schema
 from sqlalchemy.engine.url import URL
@@ -46,6 +50,8 @@
     "unexpected '(?P<syntax_error>.+?)'."
 )
 
+logger = logging.getLogger(__name__)
+
 
 class SnowflakeParametersSchema(Schema):
     username = fields.Str(required=True)
@@ -279,3 +285,52 @@ def parameters_json_schema(cls) -> Any:
 
         spec.components.schema(cls.__name__, schema=cls.parameters_schema)
         return spec.to_dict()["components"]["schemas"][cls.__name__]
+
+    @staticmethod
+    def update_params_from_encrypted_extra(
+        database: "Database",
+        params: Dict[str, Any],
+    ) -> None:
+        if not database.encrypted_extra:
+            return
+        try:
+            encrypted_extra = json.loads(database.encrypted_extra)
+        except json.JSONDecodeError as ex:
+            logger.error(ex, exc_info=True)
+            raise ex
+        auth_method = encrypted_extra.get("auth_method", None)
+        auth_params = encrypted_extra.get("auth_params", {})
+        if not auth_method:
+            return
+        connect_args = params.setdefault("connect_args", {})
+        if auth_method == "keypair":
+            privatekey_body = auth_params.get("privatekey_body", None)
+            key = None
+            if privatekey_body:
+                key = privatekey_body.encode()
+            else:
+                with open(auth_params["privatekey_path"], "rb") as key_temp:
+                    key = key_temp.read()
+            p_key = serialization.load_pem_private_key(
+                key,
+                password=auth_params["privatekey_pass"].encode(),
+                backend=default_backend(),
+            )
+            pkb = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+            connect_args["private_key"] = pkb
+        else:
+            allowed_extra_auths = current_app.config[
+                "ALLOWED_EXTRA_AUTHENTICATIONS"
+            ].get("snowflake", {})
+            if auth_method in allowed_extra_auths:
+                snowflake_auth = allowed_extra_auths.get(auth_method)
+            else:
+                raise ValueError(
+                    f"For security reason, custom authentication '{auth_method}' "
+                    f"must be listed in 'ALLOWED_EXTRA_AUTHENTICATIONS' config"
+                )
+            connect_args["auth"] = snowflake_auth(**auth_params)