Implement version-based migrations (#2453)

lib/livebook/migration.ex

@@ -1,23 +1,14 @@
 defmodule Livebook.Migration do
   use GenServer, restart: :temporary
 
-  # We version our storage so we can remove migrations in the future
+  alias Livebook.Storage
-  # by deleting the whole storage when finding too old versions.
+
-  #
+  # We version out storage, so that we know which migrations to run
-  # We also document tables and IDs used in previous versions,
+  # when someone upgrades. In the future we can also remove migrations
-  # as those must be avoided in the future.
+  # by deleting the whole storage when the storage version is too old.
-  #
-  # ## v1 (Livebook v0.11, Oct 2023)
-  #
-  # * Deleted hubs.local-host
-  # * Migrated secrets to hub_secrets
-  # * Migrated filesystems to file_systems
-  # * Migrated settings.global.default_file_system_id to settings.global.default_dir
-  # * Added :file_system_type to starred/recent files
-  #
   @migration_version 1
 
-  alias Livebook.Storage
+  def migration_version(), do: @migration_version
 
   def start_link(_opts) do
     GenServer.start_link(__MODULE__, :ok)
@@ -27,23 +18,19 @@ defmodule Livebook.Migration do
     insert_personal_hub()
     remove_offline_hub()
 
-    # v1
+    storage_version =
-    add_personal_hub_secret_key()
+      case Storage.fetch_key(:system, "global", :migration_version) do
-    delete_local_host_hub()
+        {:ok, version} -> version
-    move_app_secrets_to_personal_hub()
+        :error -> 0
-    add_file_system_type_to_notebook_manager_files()
+      end
 
-    # TODO: remove on Livebook v0.12
+    for version <- (storage_version + 1)..@migration_version//1 do
-    update_file_systems_to_deterministic_ids()
+      migration(version)
-    ensure_new_file_system_attributes()
+    end
-    move_default_file_system_id_to_default_dir()
 
     Storage.insert(:system, "global", migration_version: @migration_version)
-    :ignore
-  end
 
-  defp delete_local_host_hub() do
+    :ignore
-    Storage.delete(:hubs, "local-host")
   end
 
   defp insert_personal_hub() do
@@ -62,17 +49,25 @@ defmodule Livebook.Migration do
     # be present if the environment variables are set. Consequently, we
     # always remove it and insert on startup if applicable.
 
-    for %{id: "team-" <> _ = id} = attrs <- Storage.all(:hubs), offline_hub?(attrs) do
+    for %{id: "team-" <> _ = id, offline: true} <- Storage.all(:hubs) do
       :ok = Storage.delete(:hubs, id)
     end
   end
 
-  # TODO: use just :offline on Livebook v0.12
+  defp migration(1) do
-  defp offline_hub?(%{offline: true}), do: true
+    v1_add_personal_hub_secret_key()
-  defp offline_hub?(%{user_id: 0, org_id: 0, org_key_id: 0}), do: true
+    v1_delete_local_host_hub()
-  defp offline_hub?(_attrs), do: false
+    v1_move_app_secrets_to_personal_hub()
+    v1_add_file_system_type_to_notebook_manager_files()
+    v1_remove_old_filesystems()
+    v1_remove_default_file_system_id()
+  end
 
-  defp move_app_secrets_to_personal_hub() do
+  defp v1_delete_local_host_hub() do
+    Storage.delete(:hubs, "local-host")
+  end
+
+  defp v1_move_app_secrets_to_personal_hub() do
     for %{name: name, value: value} <- Storage.all(:secrets) do
       secret = %Livebook.Secrets.Secret{
         name: name,
@@ -85,87 +80,15 @@ defmodule Livebook.Migration do
     end
   end
 
-  defp add_personal_hub_secret_key() do
+  defp v1_add_personal_hub_secret_key() do
     with :error <- Storage.fetch_key(:hubs, Livebook.Hubs.Personal.id(), :secret_key) do
       secret_key = Livebook.Hubs.Personal.generate_secret_key()
       Storage.insert(:hubs, Livebook.Hubs.Personal.id(), secret_key: secret_key)
     end
   end
 
-  # We changed S3 file system ids, such that they are deterministic
-  # for the same bucket, rather than random. We take this opportunity
-  # to rename the scope from :filesystem to :file_systems, which
-  # conveniently allows for easy check if there's anything to migrate.
-  # This migration can be removed in the future (at the cost of discarding
-  # very old file systems (which can be re-added).
-  # TODO: remove on Livebook v0.12
-  defp update_file_systems_to_deterministic_ids() do
-    case Storage.all(:filesystem) do
-      [] ->
-        :ok
-
-      configs ->
-        id_mapping =
-          for config <- configs, into: %{} do
-            old_id = config.id
-
-            # Ensure new file system fields
-            new_fields = %{
-              hub_id: Livebook.Hubs.Personal.id(),
-              external_id: nil,
-              region: Livebook.FileSystem.S3.region_from_url(config.bucket_url)
-            }
-
-            config = Map.merge(new_fields, config)
-
-            # At this point S3 is the only file system we store
-            file_system = Livebook.FileSystems.load("s3", config)
-            Livebook.Hubs.Personal.save_file_system(file_system)
-            Storage.delete(:filesystem, old_id)
-            {old_id, file_system.id}
-          end
-
-        # Remap default file system id
-        with {:ok, default_file_system_id} <-
-               Storage.fetch_key(:settings, "global", :default_file_system_id),
-             {:ok, new_id} <- Map.fetch(id_mapping, default_file_system_id) do
-          Storage.insert(:settings, "global", default_file_system_id: new_id)
-        end
-    end
-  end
-
-  # Note that this is already handled by update_file_systems_to_deterministic_ids/0,
-  # we add this migration so it also applies to people using Livebook main.
-  # TODO: remove on Livebook v0.12
-  defp ensure_new_file_system_attributes() do
-    for attrs <- Storage.all(:file_systems) do
-      new_attrs = %{
-        hub_id: Livebook.Hubs.Personal.id(),
-        external_id: nil,
-        region: Livebook.FileSystem.S3.region_from_url(attrs.bucket_url)
-      }
-
-      attrs = Map.merge(new_attrs, attrs)
-
-      Storage.insert(:file_systems, attrs.id, Map.to_list(attrs))
-    end
-  end
-
-  # TODO: remove on Livebook v0.12
-  defp move_default_file_system_id_to_default_dir() do
-    with {:ok, default_file_system_id} <-
-           Storage.fetch_key(:settings, "global", :default_file_system_id) do
-      Livebook.Hubs.get_file_systems()
-      |> Enum.find(&(&1.id == default_file_system_id))
-      |> Livebook.FileSystem.File.new()
-      |> Livebook.Settings.set_default_dir()
-
-      Storage.delete_key(:settings, "global", :default_file_system_id)
-    end
-  end
-
   # In the past, not all files had a file_system_type, so we need to detect one from the id.
-  defp add_file_system_type_to_notebook_manager_files() do
+  defp v1_add_file_system_type_to_notebook_manager_files() do
     with {:ok, %{recent_notebooks: _, starred_notebooks: _} = attrs} <-
            Storage.fetch(:notebook_manager, "global") do
       attrs =
@@ -188,4 +111,20 @@ defmodule Livebook.Migration do
 
     Map.put(file, :file_system_type, file_system_type)
   end
+
+  defp v1_remove_old_filesystems() do
+    # For a while we had a migration converting this to the new table.
+    # Now we just delete the old entries for simplicity. If someone
+    # upgrateds from an old version, they just need to re-add their
+    # file systems. This has negligible impact.
+    for id <- Storage.all(:filesystem), do: Storage.delete(:filesystem, id)
+  end
+
+  defp v1_remove_default_file_system_id() do
+    # For a while we had a migration converting this to default dir.
+    # Now we just delete the old setting for simplicity. If someone
+    # upgrateds from an old version, they just need to set the default
+    # fiel system again. This has negligible impact.
+    Storage.delete_key(:settings, "global", :default_file_system_id)
+  end
 end

lib/livebook/storage.ex

@@ -140,12 +140,44 @@ defmodule Livebook.Storage do
     GenServer.call(__MODULE__, {:delete_key, namespace, entity_id, key})
   end
 
-  @doc """
+  defp config_file_path() do
-  Returns file path where the data is persisted.
+    migration_version = Livebook.Migration.migration_version()
-  """
+    Path.join([Livebook.Config.data_path(), "livebook_config.v#{migration_version}.ets"])
-  @spec config_file_path() :: Path.t()
+  end
-  def config_file_path() do
+
-    Path.join([Livebook.Config.data_path(), "livebook_config.ets"])
+  defp config_file_path_for_restore() do
+    # We look for a file matching the current expected version, or an
+    # older one that will get migrated.
+    #
+    # There may be files with a more recent version in case the user
+    # downgrated Livebook, and those files we ignore.
+
+    migration_version = Livebook.Migration.migration_version()
+    dir = Livebook.Config.data_path()
+
+    names =
+      case File.ls(dir) do
+        {:ok, names} -> names
+        {:error, _} -> []
+      end
+
+    candidates =
+      for name <- names, version = file_version(name), version <= migration_version do
+        %{name: name, version: version}
+      end
+
+    if candidates != [] do
+      %{name: name} = Enum.max_by(candidates, & &1.version)
+      Path.join([dir, name])
+    end
+  end
+
+  defp file_version(name) do
+    case String.split(name, ".") do
+      ["livebook_config", "ets"] -> 0
+      ["livebook_config", "v" <> version, "ets"] -> String.to_integer(version)
+      _ -> nil
+    end
   end
 
   @impl true
@@ -197,24 +229,23 @@ defmodule Livebook.Storage do
   defp table_name(), do: :persistent_term.get(__MODULE__)
 
   defp load_or_create_table() do
-    path = config_file_path()
+    tab =
+      if path = config_file_path_for_restore() do
+        path
+        |> String.to_charlist()
+        |> :ets.file2tab()
+        |> case do
+          {:ok, tab} ->
+            Logger.info("Reading storage from #{path}")
+            tab
 
-    path
+          {:error, reason} ->
-    |> String.to_charlist()
+            Logger.warning("Could not open up #{path}: #{inspect(reason)}")
-    |> :ets.file2tab()
+            nil
-    |> case do
-      {:ok, tab} ->
-        Logger.info("Reading storage from #{path}")
-        tab
-
-      {:error, reason} ->
-        case reason do
-          {:read_error, {:file_error, _, :enoent}} -> :ok
-          _ -> Logger.warning("Could not open up #{config_file_path()}: #{inspect(reason)}")
         end
+      end
 
-        :ets.new(__MODULE__, [:protected, :duplicate_bag])
+    tab || :ets.new(__MODULE__, [:protected, :duplicate_bag])
-    end
   end
 
   defp delete_keys(table, namespace, entity_id, keys) do