Add connection string parsing to Snowflake (#43923)

* Add connection string parsing

* Add test case with role suffixed key

modules/drivers/snowflake/src/metabase/driver/snowflake.clj

@@ -43,7 +43,8 @@
    (java.io File)
    (java.sql Connection DatabaseMetaData ResultSet Types)
    (java.time LocalDate LocalDateTime LocalTime OffsetDateTime OffsetTime ZonedDateTime)
-   (net.snowflake.client.jdbc SnowflakeSQLException)))
+   (java.util Properties)
+   (net.snowflake.client.jdbc SnowflakeConnectString SnowflakeSQLException)))
 
 (set! *warn-on-reflection* true)
 
@@ -125,6 +126,15 @@
   (when raw-name
     (str "\"" (str/replace raw-name "\"" "\"\"") "\"")))
 
+(defn connection-str->parameters
+  "Get map of parameters from Snowflake `conn-str`, where keys are uppercase string parameter names and values
+  are strings. Returns nil when string is invalid."
+  [conn-str]
+  (let [^SnowflakeConnectString conn-str* (SnowflakeConnectString/parse conn-str (Properties.))]
+    (if-not (.isValid conn-str*)
+      (log/warn "Invalid connection string.")
+      (.getParameters conn-str*))))
+
 (defn- maybe-add-role-to-spec-url
   "Maybe add role to `spec`'s `:connection-uri`. This is necessary for rsa auth to work, because at the time of writing
   Snowflake jdbc driver ignores `:role` connection property when `:connection-uri` (presumably containing
@@ -132,7 +142,7 @@
   [spec details]
   (if (and (string? (not-empty (:connection-uri spec)))
            (string? (not-empty (:role details)))
-           (not (str/includes? (:connection-uri spec) "role=")))
+           (not (contains? (connection-str->parameters (:connection-uri spec)) "ROLE")))
     (let [role-opts-str (sql-jdbc.common/additional-opts->string :url {:role (codec/url-encode (:role details))})]
       (-> spec
           ;; It is advised to use either connection property or url parameter, not both. Eg. in the following link.

modules/drivers/snowflake/test/metabase/driver/snowflake_test.clj

@@ -35,6 +35,7 @@
    [metabase.test.data.sql :as sql.tx]
    [metabase.test.data.sql.ddl :as ddl]
    [metabase.util :as u]
+   [ring.util.codec :as codec]
    [toucan2.core :as t2]
    [toucan2.tools.with-temp :as t2.with-temp]))
 
@@ -427,9 +428,10 @@
        ;; Following is required when private key and role are used together. See
        ;; the [[metabase.driver.snowflake/maybe-add-role-to-spec-url]] for the details.
        (testing "Role is added to connection url, if url is present (#43600)"
-         (is (str/includes?
-              (:connection-uri (sql-jdbc.conn/connection-details->spec :snowflake (assoc details :role "SOME_ROLE")))
-              "role=SOME_ROLE")))))))
+         (let [details-with-role (assoc details :role "SOME_ROLE")
+               conn-str (:connection-uri (sql-jdbc.conn/connection-details->spec :snowflake details-with-role))
+               parsed-params (driver.snowflake/connection-str->parameters conn-str)]
+           (is (= "SOME_ROLE" (get parsed-params "ROLE")))))))))
 
 (deftest ^:parallel replacement-snippet-date-param-test
   (mt/test-driver :snowflake
@@ -653,3 +655,34 @@
                     ["2023-11-01T00:00:00+11:00" 1]]
                    (mt/with-temporary-setting-values [report-timezone "Australia/Sydney"]
                      (mt/rows (qp/process-query query)))))))))))
+
+(deftest ^:parallel connection-str->parameters-test
+  (testing "Returns nil for invalid connection string"
+    (are [conn-str] (= nil (driver.snowflake/connection-str->parameters conn-str))
+      nil "" "asdf" "snowflake:jdbc://x"))
+  (testing "Returns `\"ACCOUNT\"` for valid strings of no parameters"
+    (are [conn-str] (= {"ACCOUNT" "x"} (driver.snowflake/connection-str->parameters conn-str))
+      "jdbc:snowflake://x.snowflakecomputing.com"
+      "jdbc:snowflake://x.snowflakecomputing.com/"
+      "jdbc:snowflake://x.snowflakecomputing.com/?"))
+  (testing "Returns decoded parameters"
+    (let [role "!@#$%^&*()"]
+      (is (= {"ACCOUNT" "x"
+              "ROLE" role}
+             (driver.snowflake/connection-str->parameters (str "jdbc:snowflake://x.snowflakecomputing.com/"
+                                                               "?role=" (codec/url-encode role)))))))
+  (testing "Returns multiple url parameters"
+    (let [role "!@#$%^&*()"]
+      (is (= {"ACCOUNT" "x"
+              "ROLE" role
+              "FOO" "bar"}
+             (driver.snowflake/connection-str->parameters (str "jdbc:snowflake://x.snowflakecomputing.com/"
+                                                               "?role=" (codec/url-encode role)
+                                                               "&foo=bar"))))))
+  (testing (str "Returns nothing for role suffixed keys "
+                "(https://github.com/metabase/metabase/pull/43602#discussion_r1628043704)")
+    (let [role "!@#$%^&*()"
+          params (driver.snowflake/connection-str->parameters (str "jdbc:snowflake://x.snowflakecomputing.com/"
+                                                                   "?asdfrole=" (codec/url-encode role)))]
+      (is (not (contains? params "ROLE")))
+      (is (contains? params "ASDFROLE")))))