message: feat: add snowflake keypair authentication
diff: # 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
  @@ -22,6 +23,9 @@ 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 @@ SYNTAX_ERROR_REGEX = re.compile(
  "unexpected '(?P<syntax_error>.+?)'."
  )
  +logger = logging.getLogger(__name__)
  +
  class SnowflakeParametersSchema(Schema):
  username = fields.Str(required=True)
  @@ -279,3 +285,52 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
  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)
  