Extract secrets storage to personal hub (#2132)

lib/livebook/hubs/personal.ex

@@ -5,7 +5,10 @@ defmodule Livebook.Hubs.Personal do
   import Ecto.Changeset
 
   alias Livebook.Hubs
+  alias Livebook.Storage
+  alias Livebook.Secrets.Secret
 
+  @secrets_namespace :hub_secrets
   @secret_key_size 64
 
   @type t :: %__MODULE__{
@@ -76,6 +79,51 @@ defmodule Livebook.Hubs.Personal do
     |> put_change(:id, id())
   end
 
+  @doc """
+  Get the secrets list from storage.
+  """
+  @spec get_secrets :: [Secret.t()]
+  def get_secrets do
+    Enum.map(Storage.all(@secrets_namespace), &to_secret/1)
+  end
+
+  @doc """
+  Gets a secret from storage.
+
+  Raises `RuntimeError` if the secret doesn't exist.
+  """
+  @spec fetch_secret!(String.t()) :: Secret.t()
+  def fetch_secret!(id) do
+    Storage.fetch!(@secrets_namespace, id) |> to_secret()
+  end
+
+  @doc """
+  Stores the given secret as is, without validation.
+  """
+  @spec set_secret(Secret.t()) :: Secret.t()
+  def set_secret(secret) do
+    attributes = Map.from_struct(secret)
+    :ok = Storage.insert(@secrets_namespace, secret.name, Map.to_list(attributes))
+    secret
+  end
+
+  @doc """
+  Unset secret from given id.
+  """
+  @spec set_secret(String.t()) :: :ok
+  def unset_secret(id) do
+    Storage.delete(@secrets_namespace, id)
+    :ok
+  end
+
+  defp to_secret(%{name: name, value: value}) do
+    %Secret{
+      name: name,
+      value: value,
+      hub_id: Livebook.Hubs.Personal.id()
+    }
+  end
+
   @doc """
   Generates a random secret key used for stamping the notebook.
   """
@@ -87,7 +135,7 @@ end
 
 defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Personal do
   alias Livebook.Hubs.Broadcasts
-  alias Livebook.Secrets
+  alias Livebook.Hubs.Personal
 
   def load(personal, fields) do
     %{
@@ -117,22 +165,22 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Personal do
 
   def capabilities(_personal), do: ~w(list_secrets create_secret)a
 
-  def get_secrets(personal) do
+  def get_secrets(_personal) do
-    Secrets.get_secrets(personal)
+    Personal.get_secrets()
   end
 
   def create_secret(_personal, secret) do
-    Secrets.set_secret(secret)
+    Personal.set_secret(secret)
     :ok = Broadcasts.secret_created(secret)
   end
 
   def update_secret(_personal, secret) do
-    Secrets.set_secret(secret)
+    Personal.set_secret(secret)
     :ok = Broadcasts.secret_updated(secret)
   end
 
-  def delete_secret(personal, secret) do
+  def delete_secret(_personal, secret) do
-    :ok = Secrets.unset_secret(personal, secret.name)
+    :ok = Personal.unset_secret(secret.name)
     :ok = Broadcasts.secret_deleted(secret)
   end
 

lib/livebook/migration.ex

@@ -37,7 +37,7 @@ defmodule Livebook.Migration do
         hub_id: Livebook.Hubs.Personal.id()
       }
 
-      Livebook.Secrets.set_secret(secret)
+      Livebook.Hubs.Personal.set_secret(secret)
       Livebook.Storage.delete(:secrets, name)
     end
   end

lib/livebook/secrets.ex

@@ -1,48 +1,9 @@
 defmodule Livebook.Secrets do
-  # This module is used to store secrets on Livebook.Storage for specific hubs.
+  # Shared secret functionality across all hubs.
-  # Currently it is only used by personal hub.
   @moduledoc false
 
-  alias Livebook.Hubs.Provider
-  alias Livebook.Storage
   alias Livebook.Secrets.Secret
 
-  @namespace :hub_secrets
-
-  @doc """
-  Get the secrets list from storage.
-  """
-  @spec get_secrets(Provider.t()) :: list(Secret.t())
-  def get_secrets(hub) do
-    for fields <- Storage.all(@namespace),
-        from_hub?(fields, hub),
-        do: to_struct(fields)
-  end
-
-  @doc """
-  Gets a secret from storage.
-  Raises `RuntimeError` if the secret doesn't exist.
-  """
-  @spec fetch_secret!(Provider.t(), String.t()) :: Secret.t()
-  def fetch_secret!(hub, id) do
-    fields = Storage.fetch!(@namespace, id)
-    true = from_hub?(fields, hub)
-
-    to_struct(fields)
-  end
-
-  @doc """
-  Gets a secret from storage.
-  """
-  @spec get_secret(Provider.t(), String.t()) :: {:ok, Secret.t()} | :error
-  def get_secret(hub, id) do
-    with {:ok, fields} <- Storage.fetch(@namespace, id) do
-      if from_hub?(fields, hub),
-        do: {:ok, to_struct(fields)},
-        else: :error
-    end
-  end
-
   @doc """
   Returns an `%Ecto.Changeset{}` for tracking secret changes.
   """
@@ -61,42 +22,6 @@ defmodule Livebook.Secrets do
     |> Ecto.Changeset.apply_action(:update)
   end
 
-  @doc """
-  Stores the given secret as is, without validation.
-  """
-  @spec set_secret(Secret.t()) :: Secret.t()
-  def set_secret(secret) do
-    attributes = Map.from_struct(secret)
-    :ok = Storage.insert(@namespace, secret.name, Map.to_list(attributes))
-
-    secret
-  end
-
-  @doc """
-  Unset secret from given id.
-  """
-  @spec unset_secret(Provider.t(), String.t()) :: :ok
-  def unset_secret(hub, id) do
-    with {:ok, _secret} <- get_secret(hub, id) do
-      Storage.delete(@namespace, id)
-    end
-
-    :ok
-  end
-
-  defp to_struct(%{name: name, value: value} = fields) do
-    %Secret{
-      name: name,
-      value: value,
-      hub_id: fields[:hub_id] || Livebook.Hubs.Personal.id()
-    }
-  end
-
-  defp from_hub?(fields, hub) do
-    hub_id = fields[:hub_id] || Livebook.Hubs.Personal.id()
-    hub_id == hub.id
-  end
-
   @secret_startup_key :livebook_startup_secrets
 
   @doc """

test/livebook/hubs/personal_test.exs

@@ -0,0 +1,28 @@
+defmodule Livebook.Hubs.PersonalTest do
+  use Livebook.DataCase
+
+  alias Livebook.Hubs.Personal
+
+  test "get_secrets/1 returns a list of secrets from storage" do
+    secret = build(:secret, name: "FOO", value: "111")
+
+    Personal.set_secret(secret)
+    assert secret in Personal.get_secrets()
+
+    Personal.unset_secret(secret.name)
+    refute secret in Personal.get_secrets()
+  end
+
+  test "fetch an specific secret" do
+    secret = insert_secret(name: "FOO", value: "111")
+
+    assert_raise Livebook.Storage.NotFoundError,
+                 ~s(could not find entry in "hub_secrets" with ID "NOT_HERE"),
+                 fn ->
+                   Personal.fetch_secret!("NOT_HERE")
+                 end
+
+    assert Personal.fetch_secret!(secret.name) == secret
+    Personal.unset_secret(secret.name)
+  end
+end

test/livebook/secrets_test.exs

@@ -1,37 +1,9 @@
 defmodule Livebook.SecretsTest do
-  use ExUnit.Case
+  use Livebook.DataCase, async: true
-  use Livebook.DataCase
 
   alias Livebook.Secrets
   alias Livebook.Secrets.Secret
 
-  setup do
-    {:ok, hub: build(:personal)}
-  end
-
-  test "get_secrets/1 returns a list of secrets from storage", %{hub: hub} do
-    secret = build(:secret, name: "FOO", value: "111")
-
-    Secrets.set_secret(secret)
-    assert secret in Secrets.get_secrets(hub)
-
-    Secrets.unset_secret(hub, secret.name)
-    refute secret in Secrets.get_secrets(hub)
-  end
-
-  test "fetch an specific secret", %{hub: hub} do
-    secret = insert_secret(name: "FOO", value: "111")
-
-    assert_raise Livebook.Storage.NotFoundError,
-                 ~s(could not find entry in "hub_secrets" with ID "NOT_HERE"),
-                 fn ->
-                   Secrets.fetch_secret!(hub, "NOT_HERE")
-                 end
-
-    assert Secrets.fetch_secret!(hub, secret.name) == secret
-    Secrets.unset_secret(hub, secret.name)
-  end
-
   describe "update_secret/2" do
     test "returns a valid secret" do
       attrs = params_for(:secret, name: "FOO", value: "111")