Enable Key-pair authentication for Snowflake

CHANGELOG.md

@@ -55,6 +55,7 @@
     columns instead of discarding them.
 - [Added DB_Table.Offset for SQLServer][12206]
 - [Added DB_Table.Offset for Snowflake, Postgres, SQLite][12251]
+- [Support for key-pair authentication in Snowflake connector.][12247]
 
 [11926]: https://github.com/enso-org/enso/pull/11926
 [12031]: https://github.com/enso-org/enso/pull/12031
@@ -63,6 +64,7 @@
 [12231]: https://github.com/enso-org/enso/pull/12231
 [12206]: https://github.com/enso-org/enso/pull/12206
 [12251]: https://github.com/enso-org/enso/pull/12251
+[12247]: https://github.com/enso-org/enso/pull/12247
 
 #### Enso Language & Runtime
 

build.sbt

@@ -317,6 +317,7 @@ lazy val enso = (project in file("."))
     `engine-runner`,
     `engine-runner-common`,
     `enso-test-java-helpers`,
+    `snowflake-test-java-helpers`,
     `exploratory-benchmark-java-helpers`,
     `fansi-wrapper`,
     filewatcher,
@@ -2828,6 +2829,7 @@ lazy val runtime = (project in file("engine/runtime"))
     (Runtime / compile) := (Runtime / compile)
       .dependsOn(`std-base` / Compile / packageBin)
       .dependsOn(`enso-test-java-helpers` / Compile / packageBin)
+      .dependsOn(`snowflake-test-java-helpers` / Compile / packageBin)
       .dependsOn(`benchmark-java-helpers` / Compile / packageBin)
       .dependsOn(`exploratory-benchmark-java-helpers` / Compile / packageBin)
       .dependsOn(`std-image` / Compile / packageBin)
@@ -4781,6 +4783,16 @@ lazy val `enso-test-java-helpers` = project
   .dependsOn(`std-base` % "provided")
   .dependsOn(`std-table` % "provided")
 
+lazy val `snowflake-test-java-helpers` = project
+  .in(file("test/Snowflake_Tests/polyglot-sources/snowflake-test-java-helpers"))
+  .settings(
+    frgaalJavaCompilerSetting,
+    autoScalaLibrary := false,
+    Compile / packageBin / artifactPath :=
+      file("test/Snowflake_Tests/polyglot/java/snowflake-test-helpers.jar")
+  )
+  .dependsOn(`std-snowflake` % "provided")
+
 lazy val `exploratory-benchmark-java-helpers` = project
   .in(
     file(
@@ -5484,6 +5496,7 @@ pkgStdLibInternal := Def.inputTask {
       (`std-table` / Compile / packageBin).value
     case "TestHelpers" =>
       (`enso-test-java-helpers` / Compile / packageBin).value
+      (`snowflake-test-java-helpers` / Compile / packageBin).value
       (`exploratory-benchmark-java-helpers` / Compile / packageBin).value
       (`benchmark-java-helpers` / Compile / packageBin).value
     case "AWS" =>
@@ -5497,6 +5510,7 @@ pkgStdLibInternal := Def.inputTask {
     case _ if buildAllCmd =>
       (`std-base` / Compile / packageBin).value
       (`enso-test-java-helpers` / Compile / packageBin).value
+      (`snowflake-test-java-helpers` / Compile / packageBin).value
       (`exploratory-benchmark-java-helpers` / Compile / packageBin).value
       (`benchmark-java-helpers` / Compile / packageBin).value
       (`std-table` / Compile / packageBin).value

distribution/lib/Standard/Database/0.0.0-dev/src/Internal/JDBC_Connection.enso

@@ -22,6 +22,7 @@ polyglot java import java.sql.DatabaseMetaData
 polyglot java import java.sql.PreparedStatement
 polyglot java import java.sql.SQLException
 polyglot java import java.sql.SQLTimeoutException
+polyglot java import org.enso.base.enso_cloud.HideableValue
 polyglot java import org.enso.database.dryrun.OperationSynchronizer
 polyglot java import org.enso.database.JDBCProxy
 polyglot java import org.graalvm.collections.Pair as Java_Pair
@@ -275,7 +276,11 @@ type JDBC_Connection
 create : Text -> Vector -> JDBC_Connection
 create url properties = handle_sql_errors <|
     java_props = properties.map pair->
-        Java_Pair.create pair.first (as_hideable_value pair.second)
+        # Some parameters may be passed by the dialect as a `HideableValue` directly, so they do not need to be converted.
+        value = pair.second
+        java_value = if value.is_a HideableValue then value else
+            as_hideable_value value
+        Java_Pair.create pair.first java_value
     java_connection = JDBCProxy.getConnection url java_props
 
     resource = Managed_Resource.register java_connection close_connection

distribution/lib/Standard/Snowflake/0.0.0-dev/src/Connection/Key_Pair_Credentials.enso

@@ -0,0 +1,156 @@
+from Standard.Base import all
+import Standard.Base.Errors.Common.Missing_Argument
+import Standard.Base.Errors.Illegal_State.Illegal_State
+import Standard.Base.Runtime.Context
+from Standard.Base.Runtime import assert
+
+import project.Internal.Key_Pair_Generator
+
+## Key-pair authentication for Snowflake.
+type Key_Pair_Credentials
+    ## Credentials for key-pair based authentication.
+
+       Arguments:
+       - username: The username to use for the connection.
+       - private_key: The private key to use for the connection.
+         It can either be a path to the private key file or an `Enso_Secret`.
+         The file must be located on the local machine.
+       - passphrase: The passphrase for the private key file. Only applicable if
+         the private key is a file.
+    Key_Pair username:Text=(Missing_Argument.throw "username") private_key:Enso_Secret|File|Enso_File=(Missing_Argument.throw "private_key") passphrase:Text|Nothing=Nothing
+
+    ## GROUP Output
+       ICON key
+       Generate a key pair that can be used for authentication.
+
+       This function will generate a new key pair and store it in the provided
+       Cloud location. The private key will be stored as an `Enso_Secret` that
+       can be used when establishing a Database connection. The public key will
+       be saved as a text file next to that secret.
+
+       To finalize the setup, the public key must be associated with your account.
+       To do so, you need to run a query like:
+           ALTER USER <your_username> SET RSA_PUBLIC_KEY=<public_key_content>;
+
+       See https://docs.snowflake.com/en/user-guide/key-pair-auth#assign-the-public-key-to-a-snowflake-user for more information on the setup.
+
+       Arguments:
+       - location: A directory where the secret and the public key will be stored.
+         Defaults to the user's home directory.
+       - name: The name of the secret that will store the private key. The
+         public key will be stored under the same name with the `.pub` extension.
+       - if_exists: Specifies what to do if the key pair with the given name
+         already exists. Defaults to warning the user and re-using the existing
+         key pair.
+    generate_key_pair (location : Enso_File = Enso_File.home) (name : Text) (if_exists:On_Existing_Key_Pair=..Use_Existing) -> Generated_Key_Pair =
+        public_key_file = location / (name + ".pub")
+        public_key_exists = public_key_file.exists
+        existing_secret = Enso_Secret.get name location
+        secret_exists = existing_secret.is_error.not
+
+        ## Check consistency - if one part exist but the other does not, that is an error.
+        if public_key_exists && secret_exists.not then
+            Error.throw (Illegal_State.Error "The public key "+public_key_file.path+" already exists, but the corresponding secret does not. Please clean up any leftover keys or use a different name/location.")
+        if secret_exists && public_key_exists.not then
+            Error.throw (Illegal_State.Error "The secret "+name+" already exists, but the corresponding public key does not. Please clean up any leftover keys or use a different name/location.")
+
+        assert (public_key_exists == secret_exists)
+        already_exists = secret_exists
+
+        # Early return - if key pair exists and we should error.
+        if already_exists && (if_exists == On_Existing_Key_Pair.Error) then
+            Error.throw (Illegal_State.Error "The key pair with the name "+name+" already exists at the given location.")
+
+        # Early return - if key pair exists and we should use it - we just return it.
+        if already_exists && (if_exists == On_Existing_Key_Pair.Use_Existing) then Generated_Key_Pair.new existing_secret public_key_file else
+            ## Finally, we need to generate the key pair - because it didn't exist or we wanted to overwrite it.
+            assert ((if_exists == On_Existing_Key_Pair.Overwrite) || already_exists.not)
+            Context.Output.if_enabled disabled_message="A new key-pair can only be generated in Output mode. Press the Write button ▶ to perform the operation." panic=False <|
+                key_pair = Key_Pair_Generator.generate
+
+                # Create/update the secret
+                secret = case secret_exists of
+                    True -> existing_secret.update_value key_pair.raw_private_key
+                    False -> Enso_Secret.create name key_pair.raw_private_key location
+
+                # Create/update the public key
+                on_existing_file = if if_exists == On_Existing_Key_Pair.Overwrite then Existing_File_Behavior.Overwrite else Existing_File_Behavior.Error
+                key_pair.formatted_public_key.write public_key_file on_existing_file=on_existing_file
+
+                Generated_Key_Pair.new secret public_key_file
+
+## The result of `Key_Pair_Credentials.generate_key_pair`.
+
+   It contains references to the generated secret and public key and can be
+   directly used when passing the private key to the credentials.
+type Generated_Key_Pair
+    private Value secret:Enso_Secret internal_file:Enso_File ~public_key_content:Text
+
+    private new secret:Enso_Secret public_key_file:Enso_File =
+        lazy_read_file =
+            public_key_file.read_text Encoding.utf_8 ..Report_Error
+        Generated_Key_Pair.Value secret public_key_file lazy_read_file
+
+    ## Returns the file containing the public key.
+    public_key_file self -> Enso_File =
+        self.internal_file
+
+    ## PRIVATE
+    to_text self -> Text =
+        "Generated_Key_Pair "+self.secret.name
+
+    ## PRIVATE
+    to_display_text self -> Text =
+        self.to_text
+
+    ## Returns the query that should be run to associate the public key with the user.
+
+       Arguments:
+       - username: Name of the user to associate the public key with. If not
+         provided, a placeholder will be used.
+    alter_user_query self (username:Text|Nothing = Nothing) -> Text =
+        generate_alter_user_query username self.public_key_content
+
+    ## PRIVATE
+    to_js_object self =
+        text = 'To finish the setup you need to add the public key to your Snowflake account.\n'
+            + 'To do so, you can run the following command in the admin console:'
+        query = self.alter_user_query
+        [text, query]
+
+## PRIVATE
+private parse_public_key public_key_content:Text -> Text =
+    base = public_key_content.trim
+    if base.starts_with Key_Pair_Generator.public_key_prefix . not then
+        Error.throw (Illegal_State.Error "The public key does not start with the expected prefix: "+Key_Pair_Generator.public_key_prefix)
+    if base.ends_with Key_Pair_Generator.public_key_suffix . not then
+        Error.throw (Illegal_State.Error "The public key does not end with the expected suffix: "+Key_Pair_Generator.public_key_suffix)
+    trimmed = base
+        . drop (..First Key_Pair_Generator.public_key_prefix.length)
+        . drop (..Last Key_Pair_Generator.public_key_suffix.length)
+    trimmed.lines.join ''
+
+## PRIVATE
+generate_alter_user_query username:Text|Nothing public_key_file_content:Text -> Text =
+    username_part = case username of
+        Nothing -> "<your_username>"
+        text : Text ->
+            escaped = text.replace '"' '""'
+            '"'+escaped+'"'
+    normalized_public_key = parse_public_key public_key_file_content
+    "ALTER USER "+username_part+" SET RSA_PUBLIC_KEY='"+normalized_public_key+"';"
+
+## PRIVATE
+   Allows to pass `Generated_Key_Pair` directly into credentials, as an `Enso_Secret`.
+Enso_Secret.from that:Generated_Key_Pair = that.secret
+
+## Specifies what to do if the key pair with the given name already exists.
+type On_Existing_Key_Pair
+    ## Skips generation and returns the already existing key pair.
+    Use_Existing
+
+    ## Generates a new key pair, overwriting the existing one.
+    Overwrite
+
+    ## Throws an error if the key pair already exists.
+    Error

distribution/lib/Standard/Snowflake/0.0.0-dev/src/Connection/Snowflake_Details.enso

@@ -2,15 +2,21 @@ from Standard.Base import all
 import Standard.Base.Data.Numbers.Number_Parse_Error
 import Standard.Base.Errors.Common.Missing_Argument
 import Standard.Base.Errors.Common.Type_Error
+import Standard.Base.Errors.Illegal_Argument.Illegal_Argument
 import Standard.Base.Errors.Illegal_State.Illegal_State
 import Standard.Base.Metadata.Widget.Text_Input
+from Standard.Base.Enso_Cloud.Enso_Secret import as_hideable_value
 
 import Standard.Database.Connection.Connection_Options.Connection_Options
 import Standard.Database.Connection.Credentials.Credentials
+import Standard.Database.Errors.SQL_Error
 import Standard.Database.Internal.Data_Link_Setup.Data_Link_Setup
 
+import project.Connection.Key_Pair_Credentials.Key_Pair_Credentials
+import project.Connection.Widgets
 import project.Snowflake_Connection.Snowflake_Connection
 
+polyglot java import org.enso.base.enso_cloud.InterpretAsPrivateKey
 polyglot java import net.snowflake.client.jdbc.SnowflakeDriver
 
 type Snowflake_Details
@@ -23,8 +29,8 @@ type Snowflake_Details
        - schema: The name of the schema to connect to.
        - warehouse: The name of the warehouse to use.
     @account (Text_Input display=..Always)
-    @credentials Credentials.default_widget
-    Snowflake (account:Text=(Missing_Argument.throw "account")) (credentials:Credentials=(Missing_Argument.throw "credentials")) database:Text="SNOWFLAKE" schema:Text="PUBLIC" warehouse:Text=""
+    @credentials Widgets.password_or_keypair_widget
+    Snowflake (account:Text=(Missing_Argument.throw "account")) (credentials:Credentials|Key_Pair_Credentials=(Missing_Argument.throw "credentials")) database:Text="SNOWFLAKE" schema:Text="PUBLIC" warehouse:Text=""
 
     ## PRIVATE
        Attempt to resolve the constructor.
@@ -47,7 +53,8 @@ type Snowflake_Details
         make_new database schema warehouse =
             Snowflake_Details.Snowflake self.account self.credentials (database.if_nothing self.database) (schema.if_nothing self.schema) (warehouse.if_nothing self.warehouse) . connect options allow_data_links
 
-        Snowflake_Connection.create self.jdbc_url properties make_new data_link_setup
+        _enhance_connection_errors self <|
+            Snowflake_Connection.create self.jdbc_url properties make_new data_link_setup
 
     ## PRIVATE
        Provides the jdbc url for the connection.
@@ -64,7 +71,22 @@ type Snowflake_Details
         ## should be prepended to properties. That would make it fallback to plain `json`.
 
         account = [Pair.new 'account' self.account]
-        credentials = [Pair.new 'user' self.credentials.username, Pair.new 'password' self.credentials.password]
+        credentials = case self.credentials of
+            Credentials.Username_And_Password username password -> [Pair.new 'user' username, Pair.new 'password' password]
+            Key_Pair_Credentials.Key_Pair username private_key passphrase ->
+                key_part = case private_key of
+                    local_file : File ->
+                        passphrase_part = if passphrase != Nothing then [Pair.new 'private_key_file_pwd' passphrase] else []
+                        [Pair.new 'private_key_file' local_file.absolute.path] + passphrase_part
+                    _ : Enso_File ->
+                        Error.throw (Illegal_Argument.Error "Currently only local files can be used as private keys. In Cloud, consider using a Secret instead of a file.")
+                    secret : Enso_Secret ->
+                        if passphrase != Nothing then
+                            Error.throw (Illegal_Argument.Error "Passphrase is not applicable when using a secret as a private key.")
+                        secret_as_private_key = InterpretAsPrivateKey.new (as_hideable_value secret)
+                        [Pair.new 'privateKey' secret_as_private_key]
+                [Pair.new 'user' username] + key_part
+
         database = [Pair.new 'db' self.database]
         schema = [Pair.new 'schema' self.schema]
         warehouse = if self.warehouse=="" then [] else [Pair.new 'warehouse' self.warehouse]
@@ -83,3 +105,21 @@ private create_data_link_structure details:Snowflake_Details data_link_location:
         warehouse_part = if details.warehouse.not_empty then [["warehouse", details.warehouse]] else []
         credential_part = [["credentials", credentials_json]]
         header + connection_part + schema_part + warehouse_part + credential_part
+
+private _enhance_connection_errors (details:Snowflake_Details) ~action =
+    result = action
+    result.catch SQL_Error error->
+        is_using_key_pair_auth = details.credentials.is_a Key_Pair_Credentials
+        message = error.java_exception.getMessage
+        is_likely_missing_public_key = is_using_key_pair_auth && message.contains "JWT token is invalid"
+        if is_likely_missing_public_key then
+            Error.throw (Illegal_State.Error "Authentication failed. Make sure you have associated the public key with your account and used correct username. The original error was: "+message cause=error)
+
+        is_using_cloud = case details.credentials of
+            Key_Pair_Credentials.Key_Pair _ private_key _ -> private_key.is_a Enso_Secret
+            _ -> False
+        is_invalid_private_key_format = message.contains "Private key provided is invalid or not supported"
+        if is_using_cloud.not && is_invalid_private_key_format then
+            Error.throw (Illegal_State.Error "The provided private key file is not supported by the Snowflake driver. We recommend generating the key using `generate_key_pair` method and storing it as a Secret in Enso Cloud. The original error was: "+message cause=error)
+
+        result

distribution/lib/Standard/Snowflake/0.0.0-dev/src/Connection/Widgets.enso

@@ -0,0 +1,16 @@
+from Standard.Base import all
+import Standard.Base.Metadata.Display
+import Standard.Base.Metadata.Widget
+from Standard.Base.Metadata.Choice import Option
+from Standard.Base.Metadata.Widget import Single_Choice
+
+import Standard.Database.Connection.Credentials.Credentials
+
+import project.Connection.Key_Pair_Credentials.Key_Pair_Credentials
+
+## PRIVATE
+password_or_keypair_widget display:Display=..When_Modified -> Widget =
+    basic_fqn = Credentials.to Meta.Type . qualified_name
+    key_pair_fqn = Key_Pair_Credentials.to Meta.Type . qualified_name
+    values = [Option "Username and Password" basic_fqn+".Username_And_Password", Option "Key Pair" key_pair_fqn+".Key_Pair"]
+    Single_Choice values=values display=display

distribution/lib/Standard/Snowflake/0.0.0-dev/src/Internal/Key_Pair_Generator.enso

@@ -0,0 +1,27 @@
+private
+
+from Standard.Base import all
+
+polyglot java import java.security.KeyPairGenerator
+polyglot java import java.util.Base64
+
+public_key_prefix = "-----BEGIN PUBLIC KEY-----"
+public_key_suffix = "-----END PUBLIC KEY-----"
+
+type Key_Pair
+    ## A pair of RSA keys stored as base64 encoded strings.
+    Value raw_private_key:Text raw_public_key:Text
+
+    formatted_public_key self -> Text =
+        crlf = '\r\n'
+        public_key_prefix + crlf + self.raw_public_key + crlf + public_key_suffix + crlf
+
+generate =
+    key_pair_generator = KeyPairGenerator.getInstance "RSA"
+    key_pair_generator.initialize 2048
+    key_pair = key_pair_generator.generateKeyPair
+
+    base64_encoder = Base64.getMimeEncoder
+    private_key = base64_encoder.encodeToString key_pair.getPrivate.getEncoded
+    public_key = base64_encoder.encodeToString key_pair.getPublic.getEncoded
+    Key_Pair.Value private_key public_key