Support nested transactions for app DB (#33816)

Support nested transactions for app DB
---------

Co-authored-by: metamben <103100869+metamben@users.noreply.github.com>
Co-authored-by: Ngoc Khuat <qn.khuat@gmail.com>

src/metabase/cmd/rotate_encryption_key.clj

 (ns metabase.cmd.rotate-encryption-key
   (:require
    [cheshire.core :as json]
-   [clojure.java.jdbc :as jdbc]
    [metabase.db :as mdb]
    [metabase.db.connection :as mdb.connection]
    [metabase.db.env :as mdb.env]
@@ -19,34 +18,25 @@
     (log/warnf "Database not found. Metabase will create a new database at %s and proceeed encrypting." "2")
     (mdb/setup-db!))
   (log/infof "%s: %s | %s" (trs "Connected to") mdb.env/db-type (mdb.env/db-file))
-  (let [make-encrypt-fn   (fn [maybe-encrypt-fn]
-                            (if to-key
-                              (partial maybe-encrypt-fn (encryption/validate-and-hash-secret-key to-key))
-                              identity))
-        encrypt-str-fn    (make-encrypt-fn encryption/maybe-encrypt)
-        encrypt-bytes-fn  (make-encrypt-fn encryption/maybe-encrypt-bytes)
-        is-h2             (= :h2 (mdb/db-type))
-        value-column      (if is-h2
-                            "\"VALUE\""
-                            :value)
-        setting-where     (if is-h2
-                            "setting.\"KEY\" = ?"
-                            "setting.key = ?")]
+  (let [make-encrypt-fn  (fn [maybe-encrypt-fn]
+                           (if to-key
+                             (partial maybe-encrypt-fn (encryption/validate-and-hash-secret-key to-key))
+                             identity))
+        encrypt-str-fn   (make-encrypt-fn encryption/maybe-encrypt)
+        encrypt-bytes-fn (make-encrypt-fn encryption/maybe-encrypt-bytes)]
     (t2/with-transaction [t-conn {:datasource (mdb.connection/data-source)}]
       (doseq [[id details] (t2/select-pk->fn :details Database)]
         (when (encryption/possibly-encrypted-string? details)
           (throw (ex-info (trs "Can''t decrypt app db with MB_ENCRYPTION_SECRET_KEY") {:database-id id})))
-        (jdbc/update! {:connection t-conn}
-                      :metabase_database
-                      {:details (encrypt-str-fn (json/encode details))}
-                      ["metabase_database.id = ?" id]))
+        (t2/update! :conn t-conn :metabase_database
+                    {:id id}
+                    {:details (encrypt-str-fn (json/encode details))}))
       (doseq [[key value] (t2/select-fn->fn :key :value Setting)]
         (if (= key "settings-last-updated")
           (setting.cache/update-settings-last-updated!)
-          (jdbc/update! {:connection t-conn}
-                        :setting
-                        {value-column (encrypt-str-fn value)}
-                        [setting-where key])))
+          (t2/update! :conn t-conn :setting
+                      {:key key}
+                      {:value (encrypt-str-fn value)})))
       ;; update all secret values according to the new encryption key
       ;; fortunately, we don't need to fetch the latest secret instance per ID, as we would need to in order to update
       ;; a secret value through the regular database save API path; instead, ALL secret values in the app DB (regardless
@@ -54,7 +44,6 @@
       (doseq [[id value] (t2/select-pk->fn :value Secret)]
         (when (encryption/possibly-encrypted-string? value)
           (throw (ex-info (trs "Can''t decrypt secret value with MB_ENCRYPTION_SECRET_KEY") {:secret-id id})))
-        (jdbc/update! {:connection t-conn}
-                      :secret
-                      {value-column (encrypt-bytes-fn value)}
-                      ["secret.id = ?" id])))))
+        (t2/update! :conn t-conn :secret
+                    {:id id}
+                    {:value (encrypt-bytes-fn value)})))))

src/metabase/db/connection.clj

@@ -138,8 +138,54 @@
   [_connectable f]
   (t2.conn/do-with-connection *application-db* f))
 
-(methodical/defmethod t2.conn/do-with-transaction :around java.sql.Connection
-  [connection options f]
-  ;; Do not deadlock if using a Connection in a different thread inside of a transaction -- see
-  ;; https://github.com/seancorfield/next-jdbc/issues/244.
-  (next-method connection (assoc options :nested-transaction-rule :ignore) f))
+(def ^:private ^:dynamic *transaction-depth* 0)
+
+(defn- do-transaction [^java.sql.Connection connection f]
+  (letfn [(thunk []
+            (let [savepoint (.setSavepoint connection)]
+              (try
+                (let [result (f connection)]
+                  (when (= *transaction-depth* 1)
+                    ;; top-level transaction, commit
+                    (.commit connection))
+                  result)
+                (catch Throwable e
+                  (.rollback connection savepoint)
+                  (throw e)))))]
+    ;; optimization: don't set and unset autocommit if it's already false
+    (if (.getAutoCommit connection)
+      (try
+        (.setAutoCommit connection false)
+        (thunk)
+        (finally
+          (.setAutoCommit connection true)))
+      (thunk))))
+
+(methodical/defmethod t2.conn/do-with-transaction java.sql.Connection
+  "Support nested transactions without introducing a lock like `next.jdbc` does, as that can cause deadlocks -- see
+  https://github.com/seancorfield/next-jdbc/issues/244. Use `Savepoint`s because MySQL only supports nested
+  transactions when done this way.
+
+  See also https://metaboat.slack.com/archives/CKZEMT1MJ/p1694103570500929
+
+  Note that these \"nested transactions\" are not the real thing (e.g., as in Oracle):
+    - there is only one commit, meaning that every transaction in a tree of transactions can see the changes
+      other transactions have made,
+    - in the presence of unsynchronized concurrent threads running nested transactions, the effects of rollback
+      are not well defined - a rollback will undo all work done by other transactions in the same tree that
+      started later."
+  [^java.sql.Connection connection {:keys [nested-transaction-rule] :or {nested-transaction-rule :allow} :as options} f]
+  (assert (#{:allow :ignore :prohibit} nested-transaction-rule))
+  (cond
+   (and (pos? *transaction-depth*)
+        (= nested-transaction-rule :ignore))
+   (f connection)
+
+   (and (pos? *transaction-depth*)
+        (= nested-transaction-rule :prohibit))
+   (throw (ex-info "Attempted to create nested transaction with :nested-transaction-rule set to :prohibit"
+                   {:options options}))
+
+   :else
+   (binding [*transaction-depth* (inc *transaction-depth*)]
+     (do-transaction connection f))))

test/metabase/cmd/rotate_encryption_key_test.clj

@@ -75,7 +75,7 @@
 
                       ;; while we're at it, disable the setting cache entirely; we are effectively creating a new app DB
                       ;; so the cache itself is invalid and can only mask the real issues
-                      setting/*disable-cache*         true?
+                      setting/*disable-cache*         true
                       mdb.connection/*application-db* (mdb.connection/application-db driver/*driver* data-source)]
               (when-not (= driver/*driver* :h2)
                 (tx/create-db! driver/*driver* {:database-name db-name}))

test/metabase/db/connection_test.clj

+(ns metabase.db.connection-test
+  (:require
+   [clojure.test :refer :all]
+   [metabase.test :as mt]
+   [metabase.test.fixtures :as fixtures]
+   [toucan2.core :as t2]
+   [toucan2.tools.with-temp :as t2.with-temp])
+  (:import
+   (java.util.concurrent Semaphore)))
+
+(set! *warn-on-reflection* true)
+
+(use-fixtures
+  :once
+  (fixtures/initialize :db))
+
+(deftest nested-transaction-test
+  (let [user-1                    (mt/random-email)
+        user-2                    (mt/random-email)
+        user-exists?              (fn [email]
+                                    (t2/exists? :model/User :email email))
+        create-user!              (fn [email]
+                                    (t2/insert! :model/User (assoc (t2.with-temp/with-temp-defaults :model/User) :email email)))
+        transaction-exception     (Exception. "(Abort the current transaction)")
+        is-transaction-exception? (fn is-transaction-exception? [e]
+                                    (or (identical? e transaction-exception)
+                                        (some-> (ex-cause e) is-transaction-exception?)))
+        do-in-transaction         (fn [thunk]
+                                    (try
+                                      (t2/with-transaction []
+                                        (thunk))
+                                      (catch Throwable e
+                                        (when-not (is-transaction-exception? e)
+                                          (throw e)))))]
+    (testing "inside transaction"
+      (do-in-transaction
+       (fn []
+         (create-user! user-1)
+         (is (user-exists? user-1))
+         (testing "inside nested transaction"
+           ;; do this on a separate thread to make sure we're not deadlocking, see
+           ;; https://github.com/seancorfield/next-jdbc/issues/244
+           (let [futur (future
+                         (do-in-transaction
+                          (fn []
+                            (create-user! user-2)
+                            (is (user-exists? user-2))
+                            (throw transaction-exception))))]
+             (is (not= ::timed-out
+                       (deref futur 1000 ::timed-out)))))
+         (testing "nested transaction aborted"
+           (is (user-exists? user-1))
+           (is (not (user-exists? user-2)))
+           (throw transaction-exception)))))
+
+    (testing "top-level transaction aborted"
+      (is (not (user-exists? user-1)))
+      (is (not (user-exists? user-2))))
+
+    (testing "make sure we set autocommit back after the transaction"
+      (t2/with-connection [^java.sql.Connection conn]
+        (t2/with-transaction [_t-conn conn]
+          ;; dummy op
+          (is (false? (.getAutoCommit conn))))
+        (is (true? (.getAutoCommit conn)))))
+
+    (testing "throw error when trying to create nested transaction when nested-transaction-rule=:prohibit"
+      (t2/with-connection [conn]
+        (t2/with-transaction [t-conn conn]
+          (is (thrown-with-msg?
+               clojure.lang.ExceptionInfo
+               #"Attempted to create nested transaction with :nested-transaction-rule set to :prohibit"
+               (t2/with-transaction [_ t-conn {:nested-transaction-rule :prohibit}]))))))
+
+    (testing "reuse transaction when creating nested transaction with nested-transaction-rule=:ignore"
+      (is (not (user-exists? user-1)))
+      (try
+        ;; the top-level transaction cleans up everything
+        (t2/with-transaction []
+          ;; This transaction doesn't modify the DB. It catches the exception
+          ;; from the nested transaction and sees its change because the nested
+          ;; transaction doesn't set any new savepoint.
+          (t2/with-transaction [t-conn]
+            (try
+              (t2/with-transaction [_ t-conn {:nested-transaction-rule :ignore}]
+                ;; Create a user...
+                (create-user! user-1)
+                (is (user-exists? user-1))
+                ;; and fail.
+                (throw transaction-exception))
+              (catch Exception e
+                (when-not (is-transaction-exception? e)
+                  (throw e)))))
+          ;; this user has not been rolled back because of :ignore
+          (is (user-exists? user-1))
+          (throw transaction-exception))
+        (catch Exception e
+          (when-not (is-transaction-exception? e)
+            (throw e))))
+      (is (not (user-exists? user-1))))
+
+    (testing "nested transaction anomalies -- this is not desired behavior"
+      (testing "commit and rollback"
+        (try
+          (t2/with-transaction []
+            (let [finished1 (Semaphore. 0)
+                  finished2 (Semaphore. 0)
+                  futur1 (future
+                           (do-in-transaction
+                            (fn []
+                              (create-user! user-1)
+                              (is (user-exists? user-1))
+                              (.release finished1)
+                              (.acquire finished2)
+                              (throw transaction-exception))))
+                  futur2 (future
+                           (.acquire finished1)
+                           (do-in-transaction
+                            (fn []
+                              ;; can see uncommited change
+                              (is (user-exists? user-1))
+                              (create-user! user-2)
+                              (is (user-exists? user-2))))
+                           (.release finished2))]
+              @futur2
+              @futur1)
+            (is (not (user-exists? user-1)))
+            ;; "committed" change has been rolled back
+            (is (not (user-exists? user-2)))
+            (throw transaction-exception))
+          (catch Exception e
+            (when-not (is-transaction-exception? e)
+              (throw e))))))))