Fix snowflake ssh tunnel (#39283)

* Fix snowflake ssh tunnel

* Put test where it will run in CI

* Override incorporate-ssh-tunnel

* Revert snowflake in util.ssh-test

* Move tunnel connection tests to sql_jdbc

* Add mock-server fixture

* Exclude drivers from tunnel test

* Fix up reconnection tests

* Exclude more drivers

* Address PR feedback

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

@@ -662,3 +662,10 @@
 (defmethod driver.sql/default-database-role :snowflake
   [_ database]
   (-> database :details :role))
+
+(defmethod driver/incorporate-ssh-tunnel-details :snowflake
+  [_driver {:keys [account host port] :as db-details}]
+  (let [details (cond-> db-details
+                  (not host) (assoc :host (str account ".snowflakecomputing.com"))
+                  (not port) (assoc :port 443))]
+    (driver/incorporate-ssh-tunnel-details :sql-jdbc details)))

src/metabase/util/ssh.clj

@@ -131,15 +131,18 @@
 ;; TODO Seems like this definitely belongs in [[metabase.driver.sql-jdbc.connection]] or something like that.
 (defmethod driver/incorporate-ssh-tunnel-details :sql-jdbc
   [_driver db-details]
-  (cond (not (use-ssh-tunnel? db-details))
-        ;; no ssh tunnel in use
-        db-details
-        (ssh-tunnel-open? db-details)
-        ;; tunnel in use, and is open
-        db-details
-        :else
-        ;; tunnel in use, and is not open
-        (include-ssh-tunnel! db-details)))
+  (cond
+    ;; no ssh tunnel in use
+    (not (use-ssh-tunnel? db-details))
+    db-details
+
+    ;; tunnel in use, and is open
+    (ssh-tunnel-open? db-details)
+    db-details
+
+    ;; tunnel in use, and is not open
+    :else
+    (include-ssh-tunnel! db-details)))
 
 (defn close-tunnel!
   "Close a running tunnel session"

test/metabase/driver/sql_jdbc/connection_test.clj

@@ -12,18 +12,27 @@
    [metabase.driver.sql-jdbc.test-util :as sql-jdbc.tu]
    [metabase.driver.util :as driver.u]
    [metabase.models :refer [Database Secret]]
+   [metabase.query-processor :as qp]
+   [metabase.query-processor.test-util :as qp.test-util]
    [metabase.sync :as sync]
    [metabase.test :as mt]
    [metabase.test.data :as data]
+   [metabase.test.data.interface :as tx]
    [metabase.test.fixtures :as fixtures]
+   [metabase.test.util :as tu]
    [metabase.util :as u]
+   [metabase.util.ssh :as ssh]
+   [metabase.util.ssh-test :as ssh-test]
    [next.jdbc :as next.jdbc]
    [toucan2.core :as t2]
-   [toucan2.tools.with-temp :as t2.with-temp]))
+   [toucan2.tools.with-temp :as t2.with-temp])
+  (:import
+   (org.h2.tools Server)))
 
 (set! *warn-on-reflection* true)
 
 (use-fixtures :once (fixtures/initialize :db))
+(use-fixtures :once ssh-test/do-with-mock-servers)
 
 (deftest can-connect-with-details?-test
   (testing "Should not be able to connect without setting h2/*allow-testing-h2-connections*"
@@ -212,3 +221,145 @@
     (mt/with-temp-env-var-value! [mb-jdbc-data-warehouse-unreturned-connection-timeout-seconds "20"]
       (is (= 20
              (sql-jdbc.conn/jdbc-data-warehouse-unreturned-connection-timeout-seconds))))))
+
+(defn- init-h2-tcp-server [port]
+  (let [args   ["-tcp" "-tcpPort", (str port), "-tcpAllowOthers" "-tcpDaemon"]
+        server (Server/createTcpServer (into-array args))]
+    (doto server (.start))))
+
+(deftest test-ssh-tunnel-connection
+  ;; sqlite cannot be behind a tunnel, h2 is tested below, unsure why others fail
+  (mt/test-drivers (disj (sql-jdbc.tu/sql-jdbc-drivers) :sqlite :h2 :oracle :vertica :presto-jdbc :bigquery-cloud-sdk :redshift :athena)
+    (testing "ssh tunnel is established"
+      (let [tunnel-db-details (assoc (:details (mt/db))
+                                     :tunnel-enabled true
+                                     :tunnel-host "localhost"
+                                     :tunnel-auth-option "password"
+                                     :tunnel-port ssh-test/ssh-mock-server-with-password-port
+                                     :tunnel-user ssh-test/ssh-username
+                                     :tunnel-pass ssh-test/ssh-password)]
+        (t2.with-temp/with-temp [Database tunneled-db {:engine (tx/driver), :details tunnel-db-details}]
+          (mt/with-db tunneled-db
+            (sync/sync-database! (mt/db))
+            (is (= [["Polo Lounge"]]
+                   (mt/rows (mt/run-mbql-query venues {:filter [:= $id 60] :fields [$name]}))))))))))
+
+(deftest test-ssh-tunnel-reconnection
+  ;; for now, run against Postgres and mysql, although in theory it could run against many different kinds
+  (mt/test-drivers #{:postgres :mysql}
+    (testing "ssh tunnel is reestablished if it becomes closed, so subsequent queries still succeed"
+      (let [tunnel-db-details (assoc (:details (mt/db))
+                                     :tunnel-enabled true
+                                     :tunnel-host "localhost"
+                                     :tunnel-auth-option "password"
+                                     :tunnel-port ssh-test/ssh-mock-server-with-password-port
+                                     :tunnel-user ssh-test/ssh-username
+                                     :tunnel-pass ssh-test/ssh-password)]
+        (t2.with-temp/with-temp [Database tunneled-db {:engine (tx/driver), :details tunnel-db-details}]
+          (mt/with-db tunneled-db
+            (sync/sync-database! (mt/db))
+            (letfn [(check-row []
+                      (is (= [["Polo Lounge"]]
+                             (mt/rows (mt/run-mbql-query venues {:filter [:= $id 60] :fields [$name]})))))]
+              ;; check that some data can be queried
+              (check-row)
+              ;; kill the ssh tunnel; fortunately, we have an existing function that can do that
+              (ssh/close-tunnel! (sql-jdbc.conn/db->pooled-connection-spec (mt/db)))
+              ;; check the query again; the tunnel should have been reestablished
+              (check-row))))))))
+
+(deftest test-ssh-tunnel-connection-h2
+  (testing (str "We need a customized version of this test for H2, because H2 requires bringing up its TCP server to tunnel into. "
+                "It will bring up a new H2 TCP server, pointing to an existing DB file (stored in source control, called 'tiny-db', "
+                "with a single table called 'my_tbl' and a GUEST user with password 'guest'); it will then use an SSH tunnel over "
+                "localhost to connect to this H2 server's TCP port to execute native queries against that table.")
+    (mt/with-driver :h2
+      (testing "ssh tunnel is established"
+        (let [h2-port (tu/find-free-port)
+              server  (init-h2-tcp-server h2-port)
+              ;; Use ACCESS_MODE_DATA=r to avoid updating the DB file
+              uri     (format "tcp://localhost:%d/./test_resources/ssh/tiny-db;USER=GUEST;PASSWORD=guest;ACCESS_MODE_DATA=r" h2-port)
+              h2-db   {:port               h2-port
+                       :host               "localhost"
+                       :db                 uri
+                       :tunnel-enabled     true
+                       :tunnel-host        "localhost"
+                       :tunnel-auth-option "password"
+                       :tunnel-port        ssh-test/ssh-mock-server-with-password-port
+                       :tunnel-user        ssh-test/ssh-username
+                       :tunnel-pass        ssh-test/ssh-password}]
+          (try
+            (t2.with-temp/with-temp [Database db {:engine :h2, :details h2-db}]
+              (mt/with-db db
+                (sync/sync-database! db)
+                (is (= {:cols [{:base_type    :type/Text
+                                  :effective_type :type/Text
+                                  :display_name "COL1"
+                                  :field_ref    [:field "COL1" {:base-type :type/Text}]
+                                  :name         "COL1"
+                                  :source       :native}
+                                 {:base_type    :type/Decimal
+                                  :effective_type :type/Decimal
+                                  :display_name "COL2"
+                                  :field_ref    [:field "COL2" {:base-type :type/Decimal}]
+                                  :name         "COL2"
+                                  :source       :native}]
+                          :rows [["First Row"  19.10M]
+                                 ["Second Row" 100.40M]
+                                 ["Third Row"  91884.10M]]}
+                         (-> {:query "SELECT col1, col2 FROM my_tbl;"}
+                             (mt/native-query)
+                             (qp/process-query)
+                             (qp.test-util/rows-and-cols))))))
+            (finally (.stop ^Server server))))))))
+
+(deftest test-ssh-tunnel-reconnection-h2
+  (testing (str "We need a customized version of this test for H2, because H2 requires bringing up its TCP server to tunnel into. "
+                "It will bring up a new H2 TCP server, pointing to an existing DB file (stored in source control, called 'tiny-db', "
+                "with a single table called 'my_tbl' and a GUEST user with password 'guest'); it will then use an SSH tunnel over "
+                "localhost to connect to this H2 server's TCP port to execute native queries against that table.")
+    (mt/with-driver :h2
+      (testing "ssh tunnel is reestablished if it becomes closed, so subsequent queries still succeed (H2 version)"
+        (let [h2-port (tu/find-free-port)
+              server  (init-h2-tcp-server h2-port)
+              ;; Use ACCESS_MODE_DATA=r to avoid updating the DB file
+              uri     (format "tcp://localhost:%d/./test_resources/ssh/tiny-db;USER=GUEST;PASSWORD=guest;ACCESS_MODE_DATA=r" h2-port)
+              h2-db   {:port               h2-port
+                       :host               "localhost"
+                       :db                 uri
+                       :tunnel-enabled     true
+                       :tunnel-host        "localhost"
+                       :tunnel-auth-option "password"
+                       :tunnel-port        ssh-test/ssh-mock-server-with-password-port
+                       :tunnel-user        ssh-test/ssh-username
+                       :tunnel-pass        ssh-test/ssh-password}]
+          (try
+            (t2.with-temp/with-temp [Database db {:engine :h2, :details h2-db}]
+              (mt/with-db db
+                (sync/sync-database! db)
+                (letfn [(check-data [] (is (= {:cols [{:base_type    :type/Text
+                                                       :effective_type :type/Text
+                                                       :display_name "COL1"
+                                                       :field_ref    [:field "COL1" {:base-type :type/Text}]
+                                                       :name         "COL1"
+                                                       :source       :native}
+                                                      {:base_type    :type/Decimal
+                                                       :effective_type :type/Decimal
+                                                       :display_name "COL2"
+                                                       :field_ref    [:field "COL2" {:base-type :type/Decimal}]
+                                                       :name         "COL2"
+                                                       :source       :native}]
+                                               :rows [["First Row"  19.10M]
+                                                      ["Second Row" 100.40M]
+                                                      ["Third Row"  91884.10M]]}
+                                              (-> {:query "SELECT col1, col2 FROM my_tbl;"}
+                                                  (mt/native-query)
+                                                  (qp/process-query)
+                                                  (qp.test-util/rows-and-cols)))))]
+                  ;; check that some data can be queried
+                  (check-data)
+                  ;; kill the ssh tunnel; fortunately, we have an existing function that can do that
+                  (ssh/close-tunnel! (sql-jdbc.conn/db->pooled-connection-spec db))
+                  ;; check the query again; the tunnel should have been reestablished
+                  (check-data))))
+            (finally (.stop ^Server server))))))))

test/metabase/util/ssh_test.clj

@@ -2,36 +2,26 @@
   (:require
    [clojure.java.io :as io]
    [clojure.test :refer :all]
-   [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
-   [metabase.models.database :refer [Database]]
-   [metabase.query-processor :as qp]
-   [metabase.query-processor.test-util :as qp.test-util]
-   [metabase.sync :as sync]
-   [metabase.test :as mt]
-   [metabase.test.data.interface :as tx]
-   [metabase.test.util :as tu]
    [metabase.util :as u]
    [metabase.util.log :as log]
-   [metabase.util.ssh :as ssh]
-   [toucan2.tools.with-temp :as t2.with-temp])
+   [metabase.util.ssh :as ssh])
   (:import
    (java.io BufferedReader InputStreamReader PrintWriter)
    (java.net InetSocketAddress ServerSocket Socket)
    (org.apache.sshd.server SshServer)
-   (org.apache.sshd.server.forward AcceptAllForwardingFilter)
-   (org.h2.tools Server)))
+   (org.apache.sshd.server.forward AcceptAllForwardingFilter)))
 
 (set! *warn-on-reflection* true)
 
-(def ^:private ssh-username "jsmith")
-(def ^:private ssh-password "supersecret")
+(def ssh-username "jsmith")
+(def ssh-password "supersecret")
+(def ssh-mock-server-with-password-port "mock port" 12221)
 (def ^:private ssh-publickey "test_resources/ssh/ssh_test.pub")
 (def ^:private ssh-key "test_resources/ssh/ssh_test")
 (def ^:private ssh-key-invalid "test_resources/ssh/ssh_test_invalid")
 (def ^:private ssh-publickey-passphrase "test_resources/ssh/ssh_test_passphrase.pub")
 (def ^:private ssh-key-with-passphrase "test_resources/ssh/ssh_test_passphrase")
 (def ^:private ssh-key-passphrase "Password1234")
-(def ^:private ssh-mock-server-with-password-port 12221)
 (def ^:private ssh-mock-server-with-publickey-port 12222)   ; ED25519 pubkey
 (def ^:private ssh-mock-server-with-publickey-passphrase-port 12223) ; RSA pubkey
 
@@ -107,7 +97,9 @@
       (log/error e "Error starting servers")
       (throw (ex-info "Error starting mock server" {} e)))))
 
-(defn- do-with-mock-servers [thunk]
+(defn do-with-mock-servers
+  "Fixture to start and stop mock ssh servers."
+  [thunk]
   (try
     (stop-mock-servers!)
     (start-mock-servers!)
@@ -222,83 +214,3 @@
           (u/deref-with-timeout server-thread 12000)
           (with-open [in-client (BufferedReader. (InputStreamReader. (.getInputStream socket)))]
             (is (= "hello from the ssh tunnel" (.readLine in-client)))))))))
-
-(defn- init-h2-tcp-server [port]
-  (let [args   ["-tcp" "-tcpPort", (str port), "-tcpAllowOthers" "-tcpDaemon"]
-        server (Server/createTcpServer (into-array args))]
-    (doto server (.start))))
-
-(deftest test-ssh-tunnel-reconnection
-  ;; for now, run against Postgres, although in theory it could run against many different kinds
-  (mt/test-drivers #{:postgres :mysql}
-    (testing "ssh tunnel is reestablished if it becomes closed, so subsequent queries still succeed"
-      (let [tunnel-db-details (assoc (:details (mt/db))
-                                     :tunnel-enabled true
-                                     :tunnel-host "localhost"
-                                     :tunnel-auth-option "password"
-                                     :tunnel-port ssh-mock-server-with-password-port
-                                     :tunnel-user ssh-username
-                                     :tunnel-pass ssh-password)]
-        (t2.with-temp/with-temp [Database tunneled-db {:engine (tx/driver), :details tunnel-db-details}]
-          (mt/with-db tunneled-db
-            (sync/sync-database! (mt/db))
-            (letfn [(check-row []
-                      (is (= [["Polo Lounge"]]
-                             (mt/rows (mt/run-mbql-query venues {:filter [:= $id 60] :fields [$name]})))))]
-              ;; check that some data can be queried
-              (check-row)
-              ;; kill the ssh tunnel; fortunately, we have an existing function that can do that
-              (ssh/close-tunnel! (sql-jdbc.conn/db->pooled-connection-spec (mt/db)))
-              ;; check the query again; the tunnel should have been reestablished
-              (check-row))))))))
-
-(deftest test-ssh-tunnel-reconnection-h2
-  (testing (str "We need a customized version of this test for H2. It will bring up a new H2 TCP server, pointing to "
-                "an existing DB file (stored in source control, called 'tiny-db', with a single table called 'my_tbl' "
-                "and a GUEST user with password 'guest'); it will then use an SSH tunnel over localhost to connect to "
-                "this H2 server's TCP port to execute native queries against that table.")
-    (mt/with-driver :h2
-      (testing "ssh tunnel is reestablished if it becomes closed, so subsequent queries still succeed (H2 version)"
-        (let [h2-port (tu/find-free-port)
-              server  (init-h2-tcp-server h2-port)
-              ;; Use ACCESS_MODE_DATA=r to avoid updating the DB file
-              uri     (format "tcp://localhost:%d/./test_resources/ssh/tiny-db;USER=GUEST;PASSWORD=guest;ACCESS_MODE_DATA=r" h2-port)
-              h2-db   {:port               h2-port
-                       :host               "localhost"
-                       :db                 uri
-                       :tunnel-enabled     true
-                       :tunnel-host        "localhost"
-                       :tunnel-auth-option "password"
-                       :tunnel-port        ssh-mock-server-with-password-port
-                       :tunnel-user        ssh-username
-                       :tunnel-pass        ssh-password}]
-          (try
-            (t2.with-temp/with-temp [Database db {:engine :h2, :details h2-db}]
-              (mt/with-db db
-                (sync/sync-database! db)
-                (letfn [(check-data [] (is (= {:cols [{:base_type    :type/Text
-                                                       :effective_type :type/Text
-                                                       :display_name "COL1"
-                                                       :field_ref    [:field "COL1" {:base-type :type/Text}]
-                                                       :name         "COL1"
-                                                       :source       :native}
-                                                      {:base_type    :type/Decimal
-                                                       :effective_type :type/Decimal
-                                                       :display_name "COL2"
-                                                       :field_ref    [:field "COL2" {:base-type :type/Decimal}]
-                                                       :name         "COL2"
-                                                       :source       :native}]
-                                               :rows [["First Row"  19.10M]
-                                                      ["Second Row" 100.40M]
-                                                      ["Third Row"  91884.10M]]}
-                                              (-> {:query "SELECT col1, col2 FROM my_tbl;"}
-                                                  (mt/native-query)
-                                                  (qp/process-query)
-                                                  (qp.test-util/rows-and-cols)))))]
-                  ;; check that some data can be queried
-                  (check-data)
-                  ;; kill the ssh tunnel; fortunately, we have an existing function that can do that
-                  (ssh/close-tunnel! (sql-jdbc.conn/db->pooled-connection-spec db))
-                  ;; check the query again; the tunnel should have been reestablished
-                  (check-data))))
-            (finally (.stop ^Server server))))))))