Add support for Snowflake's "auth_jwt"

lib/assets/connection_cell/main.js

@@ -316,7 +316,6 @@ export function init(ctx, info) {
     `,
   };
 
-
   const DuckDBForm = {
     name: "DuckDBForm",
 
@@ -328,7 +327,7 @@ export function init(ctx, info) {
       fields: {
         type: Object,
         default: {},
-      }
+      },
     },
 
     template: `
@@ -677,7 +676,7 @@ export function init(ctx, info) {
         event.preventDefault();
       },
 
-      dragLeave(_) { },
+      dragLeave(_) {},
 
       drop(event) {
         event.preventDefault();
@@ -741,6 +740,7 @@ export function init(ctx, info) {
       BaseInput: BaseInput,
       BaseSwitch: BaseSwitch,
       BaseSecret: BaseSecret,
+      BaseSelect: BaseSelect,
     },
 
     props: {
@@ -750,6 +750,16 @@ export function init(ctx, info) {
       },
     },
 
+    computed: {
+      isAuthSnowflake() {
+        return this.fields.auth_type === "auth_snowflake" || "";
+      },
+
+      isAuthJwt() {
+        return this.fields.auth_type === "auth_jwt";
+      },
+    },
+
     template: `
     <div class="row">
       <BaseInput
@@ -788,7 +798,29 @@ export function init(ctx, info) {
         inputClass="input"
         :grow
       />
+      <BaseInput
+        name="warehouse"
+        label="Warehouse"
+        type="text"
+        v-model="fields.warehouse"
+        inputClass="input"
+        :grow
+      />
+    </div>
+    <div class="row">
+      <BaseSelect
+        name="auth_type"
+        label=" Authenticate by "
+        v-model="fields.auth_type"
+        selectClass="input input--xs"
+        :grow
+        :options="[
+          { label: 'Password', value: 'auth_snowflake' },
+          { label: 'JWT Private Key', value: 'auth_jwt' }
+        ]"
+      />
       <BaseSecret
+        v-if="isAuthSnowflake"
         textInputName="password"
         secretInputName="password_secret"
         toggleInputName="use_password_secret"
@@ -799,6 +831,37 @@ export function init(ctx, info) {
         modalTitle="Set password"
         :required
       />
+      <BaseSwitch
+        v-if="isAuthJwt"
+        name="use_encrypted_private_key"
+        label="Encrypted"
+        v-model="fields.use_encrypted_private_key"
+      />
+    </div>
+    <div class="row">
+      <BaseSecret
+        v-if="isAuthJwt"
+        textInputName="private_key"
+        secretInputName="private_key_secret"
+        toggleInputName="use_private_key_secret"
+        label="JWT Private Key"
+        v-model:textInputValue="fields.private_key"
+        v-model:secretInputValue="fields.private_key_secret"
+        v-model:toggleInputValue="fields.use_private_key_secret"
+        modalTitle="Set JWT private key"
+        :required
+      />
+      <BaseSecret
+        v-if="isAuthJwt && fields.use_encrypted_private_key"
+        textInputName="private_key_passphrase"
+        secretInputName="private_key_passphrase_secret"
+        toggleInputName="use_private_key_passphrase_secret"
+        label="JWT Private Key Passphrase (Optional)"
+        v-model:textInputValue="fields.private_key_passphrase"
+        v-model:secretInputValue="fields.private_key_passphrase_secret"
+        v-model:toggleInputValue="fields.use_private_key_passphrase_secret"
+        modalTitle="Set JWT private key passphrase"
+      />
     </div>
     `,
   };

lib/kino_db/connection_cell.ex

@@ -16,6 +16,10 @@ defmodule KinoDB.ConnectionCell do
 
     password = attrs["password"] || ""
     secret_access_key = attrs["secret_access_key"] || ""
+    priv_key = attrs["private_key"] || ""
+    priv_key_passphrase = attrs["private_key_passphrase"] || ""
+
+    auth_type = attrs["auth_type"] || ""
 
     fields = %{
       "variable" => Kino.SmartCell.prefixed_var_name("conn", attrs["variable"]),
@@ -26,6 +30,7 @@ defmodule KinoDB.ConnectionCell do
       "use_ipv6" => Map.get(attrs, "use_ipv6", false),
       "use_ssl" => Map.get(attrs, "use_ssl", false),
       "cacertfile" => attrs["cacertfile"] || "",
+      "auth_type" => auth_type,
       "username" => attrs["username"] || "",
       "password" => password,
       "use_password_secret" => Map.has_key?(attrs, "password_secret") || password == "",
@@ -39,12 +44,21 @@ defmodule KinoDB.ConnectionCell do
       "use_secret_access_key_secret" =>
         Map.has_key?(attrs, "secret_access_key_secret") || secret_access_key == "",
       "secret_access_key_secret" => attrs["secret_access_key_secret"] || "",
+      "private_key" => priv_key,
+      "use_private_key_secret" => Map.has_key?(attrs, "use_private_key_secret") || priv_key == "",
+      "private_key_secret" => attrs["private_key_secret"] || "",
+      "use_encrypted_private_key" => Map.get(attrs, "use_encrypted_private_key", false),
+      "private_key_passphrase" => priv_key_passphrase,
+      "use_private_key_passphrase_secret" =>
+        Map.has_key?(attrs, "use_private_key_passphrase_secret") || priv_key_passphrase == "",
+      "private_key_passphrase_secret" => attrs["private_key_passphrase_secret"] || "",
       "token" => attrs["token"] || "",
       "region" => attrs["region"] || "us-east-1",
       "workgroup" => attrs["workgroup"] || "",
       "output_location" => attrs["output_location"] || "",
       "account" => attrs["account"] || "",
       "schema" => attrs["schema"] || "",
+      "warehouse" => attrs["warehouse"] || "",
       "instance" => attrs["instance"] || ""
     }
 
@@ -138,9 +152,7 @@ defmodule KinoDB.ConnectionCell do
               ~w|access_key_id secret_access_key token region workgroup output_location database|
 
         "snowflake" ->
-          if fields["use_password_secret"],
-            do: ~w|database schema account username password_secret|,
-            else: ~w|database schema account username password|
+          snowflake_fields(fields)
 
         "sqlserver" ->
           if fields["use_password_secret"],
@@ -163,6 +175,33 @@ defmodule KinoDB.ConnectionCell do
     Map.take(fields, @default_keys ++ connection_keys)
   end
 
+  @default_snowflake_keys ~w|account username database schema warehouse auth_type|
+  defp snowflake_fields(fields),
+    do: snowflake_fields(fields["auth_type"], fields, @default_snowflake_keys)
+
+  defp snowflake_fields("auth_jwt", fields, keys) do
+    pk_keys =
+      if fields["use_private_key_secret"],
+        do: ~w|private_key_secret use_encrypted_private_key|,
+        else: ~w|private_key use_encrypted_private_key|
+
+    phrase_keys =
+      if fields["use_encrypted_private_key"] do
+        if fields["use_private_key_passphrase_secret"],
+          do: ~w|private_key_passphrase_secret|,
+          else: ~w|private_key_passphrase|
+      else
+        []
+      end
+
+    keys ++ pk_keys ++ phrase_keys
+  end
+
+  defp snowflake_fields(_, fields, keys) do
+    keys ++
+      if fields["use_password_secret"], do: ~w|password_secret|, else: ~w|password|
+  end
+
   @impl true
   def to_source(attrs) do
     required_keys =
@@ -186,10 +225,7 @@ defmodule KinoDB.ConnectionCell do
               )
 
         "snowflake" ->
-          if(Map.has_key?(attrs, "password_secret"),
-            do: ~w|account username password_secret|,
-            else: ~w|account username password|
-          )
+          snowflake_source_required_keys(attrs)
 
         "sqlserver" ->
           ~w|hostname port|
@@ -215,6 +251,20 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
+  defp snowflake_source_required_keys(attrs),
+    do: snowflake_source_required_keys(attrs["auth_type"], attrs, ~w|account username auth_type|)
+
+  defp snowflake_source_required_keys("auth_jwt", attrs, keys) do
+    keys ++
+      if Map.has_key?(attrs, "private_key_secret"),
+        do: ~w|private_key_secret|,
+        else: ~w|private_key|
+  end
+
+  defp snowflake_source_required_keys(_, attrs, keys) do
+    keys ++ if Map.has_key?(attrs, "password_secret"), do: ~w|password_secret|, else: ~w|password|
+  end
+
   defp all_fields_filled?(attrs, keys) do
     not Enum.any?(keys, fn key -> attrs[key] in [nil, ""] end)
   end
@@ -243,10 +293,15 @@ defmodule KinoDB.ConnectionCell do
   defp to_quoted(%{"type" => "snowflake"} = attrs) do
     var = quoted_var(attrs["variable"])
 
+    snowflake_opts =
+      attrs
+      |> trim_opts()
+      |> Map.new()
+      |> snowflake_options()
+
     quote do
       :ok = Adbc.download_driver!(:snowflake)
-      uri = unquote(build_snowflake_uri(attrs))
-      {:ok, db} = Kino.start_child({Adbc.Database, driver: :snowflake, uri: uri})
+      {:ok, db} = Kino.start_child({Adbc.Database, unquote(snowflake_opts)})
       {:ok, unquote(var)} = Kino.start_child({Adbc.Connection, database: db})
     end
   end
@@ -363,6 +418,64 @@ defmodule KinoDB.ConnectionCell do
     end
   end
 
+  defp quoted_private_key(%{"private_key" => pk}), do: pk
+
+  defp quoted_private_key(%{"private_key_secret" => ""}), do: ""
+
+  defp quoted_private_key(%{"private_key_secret" => secret}) do
+    quote do
+      System.fetch_env!(unquote("LB_#{secret}"))
+    end
+  end
+
+  defp quoted_private_key_passphrase(%{"private_key_passphrase" => phrase}), do: phrase
+
+  defp quoted_private_key_passphrase(%{"private_key_passphrase_secret" => ""}), do: ""
+
+  defp quoted_private_key_passphrase(%{"private_key_passphrase_secret" => secret}) do
+    quote do
+      System.fetch_env!(unquote("LB_#{secret}"))
+    end
+  end
+
+  defp snowflake_options(attrs) do
+    shared_opts =
+      [
+        driver: :snowflake,
+        username: attrs["username"],
+        "adbc.snowflake.sql.account": attrs["account"],
+        "adbc.snowflake.sql.db": attrs["database"],
+        "adbc.snowflake.sql.schema": attrs["schema"],
+        "adbc.snowflake.sql.warehouse": attrs["warehouse"]
+      ]
+
+    snowflake_auth_opts(attrs["auth_type"], attrs, shared_opts)
+  end
+
+  defp snowflake_auth_opts("auth_jwt", attrs, opts) do
+    opts ++
+      [
+        "adbc.snowflake.sql.auth_type": "auth_jwt",
+        "adbc.snowflake.sql.client_option.jwt_private_key_pkcs8_value": quoted_private_key(attrs)
+      ] ++
+      if attrs["use_encrypted_private_key"] do
+        [
+          "adbc.snowflake.sql.client_option.jwt_private_key_pkcs8_password":
+            quoted_private_key_passphrase(attrs)
+        ]
+      else
+        []
+      end
+  end
+
+  defp snowflake_auth_opts(_, attrs, opts) do
+    opts ++
+      [
+        "adbc.snowflake.sql.auth_type": "auth_snowflake",
+        password: quoted_pass(attrs)
+      ]
+  end
+
   defp shared_options(attrs) do
     opts = [
       hostname: attrs["hostname"],
@@ -586,27 +699,4 @@ defmodule KinoDB.ConnectionCell do
          do: true,
          else: (_ -> false)
   end
-
-  defp build_snowflake_uri(attrs), do: build_snowflake_uri(attrs, quoted_pass(attrs))
-
-  defp build_snowflake_uri(attrs, password) when is_binary(password) do
-    "#{attrs["username"]}:#{password}@#{attrs["account"]}"
-    |> build_database_and_schema(attrs)
-  end
-
-  defp build_snowflake_uri(%{"username" => username, "account" => account} = attrs, password) do
-    rest = build_database_and_schema("@#{account}", attrs)
-
-    quote do
-      unquote("#{username}:") <> unquote(password) <> unquote(rest)
-    end
-  end
-
-  defp build_database_and_schema(uri, %{"database" => ""}), do: uri
-
-  defp build_database_and_schema(uri, %{"database" => database, "schema" => ""}),
-    do: "#{uri}/#{database}"
-
-  defp build_database_and_schema(uri, %{"database" => database, "schema" => schema}),
-    do: "#{uri}/#{database}/#{schema}"
 end