Merge pull request #4971 from metabase/support-additional-connection-string-options-for-mong

Support additional connection string options for MongoDB

src/metabase/driver/mongo.clj

@@ -206,7 +206,10 @@
                                                             {:name         "ssl"
                                                              :display-name "Use a secure connection (SSL)?"
                                                              :type         :boolean
-                                                             :default      false}]))
+                                                             :default      false}
+                                                            {:name         "additional-options"
+                                                             :display-name "Additional Mongo connection string options"
+                                                             :placeholder  "readPreference=nearest&replicaSet=test"}]))
           :execute-query                     (u/drop-first-arg qp/execute-query)
           :features                          (constantly #{:basic-aggregations :dynamic-schema :nested-fields})
           :field-values-lazy-seq             (u/drop-first-arg field-values-lazy-seq)

src/metabase/driver/mongo/util.clj

@@ -2,12 +2,14 @@
   "`*mongo-connection*`, `with-mongo-connection`, and other functions shared between several Mongo driver namespaces."
   (:require [clojure.tools.logging :as log]
             [metabase
+             [config :as config]
              [driver :as driver]
              [util :as u]]
             [metabase.util.ssh :as ssh]
             [monger
              [core :as mg]
-             [credentials :as mcred]]))
+             [credentials :as mcred]])
+  (:import [com.mongodb MongoClientOptions MongoClientOptions$Builder MongoClientURI]))
 
 (def ^:const ^:private connection-timeout-ms
   "Number of milliseconds to wait when attempting to establish a Mongo connection.
@@ -24,12 +26,45 @@
    Bound by top-level `with-mongo-connection` so it may be reused within its body."
   nil)
 
+;; the code below is done to support "additional connection options" the way some of the JDBC drivers do.
+;; For example, some people might want to specify a`readPreference` of `nearest`. The normal Java way of
+;; doing this would be to do
+;;
+;;     (.readPreference builder (ReadPreference/nearest))
+;;
+;; But the user will enter something like `readPreference=nearest`. Luckily, the Mongo Java lib can parse
+;; these options for us and return a `MongoClientOptions` like we'd prefer. Code below:
+
+(defn- client-options-for-url-params
+  "Return an instance of `MongoClientOptions` from a URL-PARAMS string, e.g.
+
+     (client-options-for-url-params \"readPreference=nearest\")
+      ;; -> #MongoClientOptions{readPreference=nearest, ...}"
+  ^MongoClientOptions [^String url-params]
+  (when (seq url-params)
+    ;; just make a fake connection string to tack the URL params on to. We can use that to have the Mongo lib
+    ;; take care of parsing the params and converting them to Java-style `MongoConnectionOptions`
+    (.getOptions (MongoClientURI. (str "mongodb://localhost/?" url-params)))))
+
+(defn- client-options->builder
+  "Return a `MongoClientOptions.Builder` for a `MongoClientOptions` CLIENT-OPTIONS.
+   If CLIENT-OPTIONS is `nil`, return a new 'default' builder."
+  ^MongoClientOptions$Builder [^MongoClientOptions client-options]
+  ;; We do it tnis way because (MongoClientOptions$Builder. nil) throws a NullPointerException
+  (if client-options
+    (MongoClientOptions$Builder. client-options)
+    (MongoClientOptions$Builder.)))
+
 (defn- build-connection-options
   "Build connection options for Mongo.
-   We have to use `MongoClientOptions.Builder` directly to configure our Mongo connection
-   since Monger's wrapper method doesn't support `.serverSelectionTimeout` or `.sslEnabled`."
-  [& {:keys [ssl?]}]
-  (-> (com.mongodb.MongoClientOptions$Builder.)
+   We have to use `MongoClientOptions.Builder` directly to configure our Mongo connection since Monger's wrapper method doesn't
+   support `.serverSelectionTimeout` or `.sslEnabled`. ADDITIONAL-OPTIONS, a String like `readPreference=nearest`, can be specified
+   as well; when passed, these are parsed into a `MongoClientOptions` that serves as a starting point for the changes made below."
+  ^MongoClientOptions [& {:keys [ssl? additional-options]
+                          :or   {ssl? false}}]
+  (-> (client-options-for-url-params additional-options)
+      client-options->builder
+      (.description (str "Metabase " (config/mb-version-info :tag)))
       (.connectTimeout connection-timeout-ms)
       (.serverSelectionTimeout connection-timeout-ms)
       (.sslEnabled ssl?)
@@ -42,17 +77,23 @@
                                             [server-address options]
                                             [server-address options credentials]))
 
+(defn- database->details
+  "Make sure DATABASE is in a standard db details format. This is done so we can accept several different types of
+   values for DATABASE, such as plain strings or the usual MB details map."
+  [database]
+  (cond
+    (string? database)            {:dbname database}
+    (:dbname (:details database)) (:details database) ; entire Database obj
+    (:dbname database)            database            ; connection details map only
+    :else                         (throw (Exception. (str "with-mongo-connection failed: bad connection details:" (:details database))))))
+
 (defn -with-mongo-connection
   "Run F with a new connection (bound to `*mongo-connection*`) to DATABASE.
    Don't use this directly; use `with-mongo-connection`."
   [f database]
-  (let [details (cond
-                  (string? database)            {:dbname database}
-                  (:dbname (:details database)) (:details database) ; entire Database obj
-                  (:dbname database)            database            ; connection details map only
-                  :else                         (throw (Exception. (str "with-mongo-connection failed: bad connection details:" (:details database)))))]
+  (let [details (database->details database)]
     (ssh/with-ssh-tunnel [details-with-tunnel details]
-      (let [{:keys [dbname host port user pass ssl authdb tunnel-host tunnel-user tunnel-pass]
+      (let [{:keys [dbname host port user pass ssl authdb tunnel-host tunnel-user tunnel-pass additional-options]
              :or   {port 27017, pass "", ssl false}} details-with-tunnel
             user             (when (seq user) ; ignore empty :user and :pass strings
                                user)
@@ -64,7 +105,7 @@
             server-address   (mg/server-address host port)
             credentials      (when user
                                (mcred/create user authdb pass))
-            connect          (partial mg/connect server-address (build-connection-options :ssl? ssl))
+            connect          (partial mg/connect server-address (build-connection-options :ssl? ssl, :additional-options additional-options))
             conn             (if credentials
                                (connect credentials)
                                (connect))

test/metabase/driver/mongo/util_test.clj

+(ns metabase.driver.mongo.util-test
+  (:require [expectations :refer :all]
+            metabase.driver.mongo.util
+            [metabase.test.util :as tu])
+  (:import com.mongodb.ReadPreference))
+
+(tu/resolve-private-vars metabase.driver.mongo.util build-connection-options)
+
+;; test that people can specify additional connection options like `?readPreference=nearest`
+(expect
+  (ReadPreference/nearest)
+  (.getReadPreference (build-connection-options :additional-options "readPreference=nearest")))
+
+(expect
+  (ReadPreference/secondaryPreferred)
+  (.getReadPreference (build-connection-options :additional-options "readPreference=secondaryPreferred")))
+
+;; make sure we can specify multiple options
+(expect
+  "test"
+  (.getRequiredReplicaSetName (build-connection-options :additional-options "readPreference=secondary&replicaSet=test")))
+
+(expect
+  (ReadPreference/secondary)
+  (.getReadPreference (build-connection-options :additional-options "readPreference=secondary&replicaSet=test")))
+
+;; make sure that invalid additional options throw an Exception
+(expect
+  IllegalArgumentException
+  (build-connection-options :additional-options "readPreference=ternary"))