From 040b336ef98efba1c3326725a8520e3cf66319a0 Mon Sep 17 00:00:00 2001
From: Andrey Antukh <niwi@niwi.nz>
Date: Tue, 6 Feb 2024 19:20:25 +0100
Subject: [PATCH] :sparkles: Add helper for restoring team after migration to
 comp-v2

---
 backend/src/app/features/components_v2.clj | 23 ++++++------
 backend/src/app/srepl/components_v2.clj    | 42 ++++++++++++++++++++++
 2 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/backend/src/app/features/components_v2.clj b/backend/src/app/features/components_v2.clj
index 8b3cd1da1..1aaa5e268 100644
--- a/backend/src/app/features/components_v2.clj
+++ b/backend/src/app/features/components_v2.clj
@@ -1442,7 +1442,7 @@
                           data)))
         (fmg/migrate-file))))
 
-(defn- get-team
+(defn get-team
   [system team-id]
   (-> (db/get system :team {:id team-id}
               {::db/remove-deleted false
@@ -1496,17 +1496,19 @@
       AND f.deleted_at IS NULL
       FOR UPDATE")
 
-(defn- get-and-lock-files
+(defn get-and-lock-files
   [conn team-id]
   (->> (db/cursor conn [sql:get-and-lock-team-files team-id])
        (map :id)))
 
-(defn- update-team-features!
-  [conn team-id features]
-  (let [features (db/create-array conn "text" features)]
+(defn update-team!
+  [conn team]
+  (let [params (-> team
+                   (update :features db/encode-pgarray conn "text")
+                   (dissoc :id))]
     (db/update! conn :team
-                {:features features}
-                {:id team-id})))
+                params
+                {:id (:id team)})))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; PUBLIC API
@@ -1590,7 +1592,7 @@
                          :skip-on-graphic-error? skip-on-graphic-error?))
         migrate-team
         (fn [{:keys [::db/conn] :as system} team-id]
-          (let [{:keys [id features name]} (get-team system team-id)]
+          (let [{:keys [id features] :as team} (get-team system team-id)]
             (if (contains? features "components/v2")
               (l/inf :hint "team already migrated")
               (let [features (-> features
@@ -1601,13 +1603,14 @@
 
                 (events/tap :progress
                             {:op :migrate-team
-                             :name name
+                             :name (:name team)
                              :id id})
 
                 (run! (partial migrate-file system)
                       (get-and-lock-files conn id))
 
-                (update-team-features! conn id features)))))]
+                (->> (assoc team :features features)
+                     (update-team! conn))))))]
 
     (binding [*team-stats* (atom {})]
       (try
diff --git a/backend/src/app/srepl/components_v2.clj b/backend/src/app/srepl/components_v2.clj
index 9bada320a..3b145d1cd 100644
--- a/backend/src/app/srepl/components_v2.clj
+++ b/backend/src/app/srepl/components_v2.clj
@@ -12,6 +12,7 @@
    [app.db :as db]
    [app.features.components-v2 :as feat]
    [app.main :as main]
+   [app.rpc.commands.files-snapshot :as rpc]
    [app.svgo :as svgo]
    [app.util.cache :as cache]
    [app.util.events :as events]
@@ -636,3 +637,44 @@
              :file-name (:name file))
       (assoc file :deleted-at (dt/now)))
     file))
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; RESTORE SNAPSHOT
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(def ^:private sql:snapshots-with-file
+  "SELECT f.id AS file_id,
+          fc.id AS id
+     FROM file AS f
+     JOIN file_change AS fc ON (fc.file_id = f.id)
+    WHERE fc.label = ? AND f.id = ANY(?)")
+
+(defn restore-team!
+  [team-id label & {:keys [rollback?] :or {rollback? true}}]
+  (let [team-id (if (string? team-id)
+                  (parse-uuid team-id)
+                  team-id)
+
+        get-file-snapshots
+        (fn [conn ids]
+          (let [label  (str "migration/" label)]
+            (db/exec! conn [sql:snapshots-with-file label
+                            (db/create-array conn "uuid" ids)])))
+
+        restore-snapshot
+        (fn [{:keys [::db/conn] :as system}]
+          (let [ids  (into #{} (feat/get-and-lock-files conn team-id))
+                snap (get-file-snapshots conn ids)
+                ids' (into #{} (map :file-id) snap)
+                team (-> (feat/get-team conn team-id)
+                         (update :features disj "components/v2"))]
+
+            (when (not= ids ids')
+              (throw (RuntimeException. "no uniform snapshot available")))
+
+            (feat/update-team! conn team)
+            (run! (partial rpc/restore-file-snapshot! system) snap)))]
+
+
+    (-> (assoc main/system ::db/rollback rollback?)
+        (db/tx-run! restore-snapshot))))