diff --git a/lib/livebook/hubs.ex b/lib/livebook/hubs.ex index 52a636ab1..616e2ce61 100644 --- a/lib/livebook/hubs.ex +++ b/lib/livebook/hubs.ex @@ -3,6 +3,7 @@ defmodule Livebook.Hubs do alias Livebook.Storage alias Livebook.Hubs.{Broadcasts, Enterprise, Fly, Local, Metadata, Provider} + alias Livebook.Secrets.Secret @namespace :hubs @@ -22,6 +23,16 @@ defmodule Livebook.Hubs do end end + @doc """ + Gets a list of hubs from storage with given capabilities. + """ + @spec get_hubs(Provider.capabilities()) :: list(Provider.t()) + def get_hubs(capabilities) do + for hub <- get_hubs(), + capability?(hub, capabilities), + do: hub + end + @doc """ Gets a list of metadatas from storage. """ @@ -79,10 +90,7 @@ defmodule Livebook.Hubs do @doc false def delete_hub(id) do with {:ok, hub} <- get_hub(id) do - if connected_hub = get_connected_hub(hub) do - GenServer.stop(connected_hub.pid, :shutdown) - end - + :ok = Provider.disconnect(hub) :ok = Storage.delete(@namespace, id) :ok = Broadcasts.hubs_metadata_changed() end @@ -167,7 +175,9 @@ defmodule Livebook.Hubs do """ @spec connect_hubs() :: :ok def connect_hubs do - for hub <- get_hubs(), do: connect_hub(hub) + for hub <- get_hubs(), + capability?(hub, [:connect]), + do: connect_hub(hub) :ok end @@ -181,23 +191,36 @@ defmodule Livebook.Hubs do end @doc """ - Gets a list of connected hubs. - - ## Example - - iex> get_connected_hubs() - [%{pid: #PID<0.178.0>, hub: %Enterprise{}}, ...] + Gets a list of hub secrets. + It gets from all hubs with secret management. """ - @spec get_connected_hubs() :: connected_hubs() - def get_connected_hubs do - for hub <- get_hubs(), connected = get_connected_hub(hub), do: connected + @spec get_secrets() :: list(Secret.t()) + def get_secrets do + for hub <- get_hubs([:secrets]), + secret <- Provider.get_secrets(hub), + do: secret end - defp get_connected_hub(hub) do - case Registry.lookup(Livebook.HubsRegistry, hub.id) do - [{pid, _}] -> %{pid: pid, hub: hub} - [] -> nil + @doc """ + Creates a secret for given hub. + """ + @spec create_secret(Secret.t()) :: :ok | {:error, list({String.t(), list(String.t())})} + def create_secret(%Secret{origin: {:hub, id}} = secret) do + case get_hub(id) do + {:ok, hub} -> + if capability?(hub, [:secrets]) do + Provider.create_secret(hub, secret) + else + {:error, %{errors: [{"hub_id", {"is invalid", []}}]}} + end + + :error -> + {:error, %{errors: [{"hub_id", {"doest not exists", []}}]}} end end + + defp capability?(hub, capabilities) do + capabilities -- Provider.capabilities(hub) == [] + end end diff --git a/lib/livebook/hubs/enterprise.ex b/lib/livebook/hubs/enterprise.ex index c85e02e45..db9dc57c4 100644 --- a/lib/livebook/hubs/enterprise.ex +++ b/lib/livebook/hubs/enterprise.ex @@ -105,7 +105,9 @@ defmodule Livebook.Hubs.Enterprise do end defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Enterprise do - def load(%Livebook.Hubs.Enterprise{} = enterprise, fields) do + alias Livebook.Hubs.EnterpriseClient + + def load(enterprise, fields) do %{ enterprise | id: fields.id, @@ -117,7 +119,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Enterprise do } end - def normalize(%Livebook.Hubs.Enterprise{} = enterprise) do + def normalize(enterprise) do %Livebook.Hubs.Metadata{ id: enterprise.id, name: enterprise.hub_name, @@ -128,10 +130,35 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Enterprise do def type(_enterprise), do: "enterprise" - def connect(%Livebook.Hubs.Enterprise{} = enterprise), - do: {Livebook.Hubs.EnterpriseClient, enterprise} + def connect(enterprise), do: {EnterpriseClient, enterprise} - def connected?(%Livebook.Hubs.Enterprise{id: id}) do - Livebook.Hubs.EnterpriseClient.connected?(id) + def connected?(enterprise) do + EnterpriseClient.connected?(enterprise.id) + end + + def disconnect(enterprise) do + EnterpriseClient.stop(enterprise.id) + end + + def capabilities(_enterprise), do: [:connect, :secrets] + + def get_secrets(enterprise) do + EnterpriseClient.get_secrets(enterprise.id) + end + + def create_secret(enterprise, %Livebook.Secrets.Secret{name: name, value: value}) do + create_secret_request = LivebookProto.CreateSecretRequest.new!(name: name, value: value) + + case EnterpriseClient.send_request(enterprise.id, create_secret_request) do + {:create_secret, _} -> + :ok + + {:changeset_error, errors} -> + errors = + for {field, values} <- errors, + do: {to_string(field), values} + + {:error, %{errors: errors}} + end end end diff --git a/lib/livebook/hubs/enterprise_client.ex b/lib/livebook/hubs/enterprise_client.ex index f7bfa6b9d..3f3363946 100644 --- a/lib/livebook/hubs/enterprise_client.ex +++ b/lib/livebook/hubs/enterprise_client.ex @@ -8,9 +8,12 @@ defmodule Livebook.Hubs.EnterpriseClient do alias Livebook.WebSocket.ClientConnection @registry Livebook.HubsRegistry + @supervisor Livebook.HubsSupervisor defstruct [:server, :hub, connected?: false, secrets: []] + @type registry_name :: {:via, Registry, {Livebook.HubsRegistry, String.t()}} + @doc """ Connects the Enterprise client with WebSocket server. """ @@ -22,16 +25,23 @@ defmodule Livebook.Hubs.EnterpriseClient do @doc """ Stops the WebSocket server. """ - @spec stop(pid()) :: :ok - def stop(pid) do - pid |> GenServer.call(:get_server) |> GenServer.stop() - GenServer.stop(pid) + @spec stop(String.t()) :: :ok + def stop(id) do + if pid = GenServer.whereis(registry_name(id)) do + DynamicSupervisor.terminate_child(@supervisor, pid) + end + + :ok end @doc """ Sends a request to the WebSocket server. """ - @spec send_request(pid(), WebSocket.proto()) :: {atom(), term()} + @spec send_request(String.t() | registry_name() | pid(), WebSocket.proto()) :: {atom(), term()} + def send_request(id, %_struct{} = data) when is_binary(id) do + send_request(registry_name(id), data) + end + def send_request(pid, %_struct{} = data) do ClientConnection.send_request(GenServer.call(pid, :get_server), data) end @@ -39,9 +49,9 @@ defmodule Livebook.Hubs.EnterpriseClient do @doc """ Returns a list of cached secrets. """ - @spec list_cached_secrets(pid()) :: list(Secret.t()) - def list_cached_secrets(pid) do - GenServer.call(pid, :list_cached_secrets) + @spec get_secrets(String.t()) :: list(Secret.t()) + def get_secrets(id) do + GenServer.call(registry_name(id), :get_secrets) end @doc """ @@ -51,8 +61,7 @@ defmodule Livebook.Hubs.EnterpriseClient do def connected?(id) do GenServer.call(registry_name(id), :connected?) catch - :exit, _ -> - false + :exit, _ -> false end ## GenServer callbacks @@ -70,7 +79,7 @@ defmodule Livebook.Hubs.EnterpriseClient do {:reply, state.server, state} end - def handle_call(:list_cached_secrets, _caller, state) do + def handle_call(:get_secrets, _caller, state) do {:reply, state.secrets, state} end diff --git a/lib/livebook/hubs/fly.ex b/lib/livebook/hubs/fly.ex index cb56e4cbd..1479eedfd 100644 --- a/lib/livebook/hubs/fly.ex +++ b/lib/livebook/hubs/fly.ex @@ -110,7 +110,7 @@ defmodule Livebook.Hubs.Fly do end defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Fly do - def load(%Livebook.Hubs.Fly{} = fly, fields) do + def load(fly, fields) do %{ fly | id: fields.id, @@ -124,7 +124,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Fly do } end - def normalize(%Livebook.Hubs.Fly{} = fly) do + def normalize(fly) do %Livebook.Hubs.Metadata{ id: fly.id, name: fly.hub_name, @@ -138,4 +138,12 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Fly do def connect(_fly), do: nil def connected?(_fly), do: false + + def disconnect(_fly), do: :ok + + def capabilities(_fly), do: [] + + def get_secrets(_fly), do: [] + + def create_secret(_fly, _secret), do: :ok end diff --git a/lib/livebook/hubs/local.ex b/lib/livebook/hubs/local.ex index 5f96a93cd..78debf51d 100644 --- a/lib/livebook/hubs/local.ex +++ b/lib/livebook/hubs/local.ex @@ -5,11 +5,11 @@ defmodule Livebook.Hubs.Local do end defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Local do - def load(%Livebook.Hubs.Local{} = local, fields) do + def load(local, fields) do %{local | id: fields.id, hub_name: fields.hub_name, hub_emoji: fields.hub_emoji} end - def normalize(%Livebook.Hubs.Local{} = local) do + def normalize(local) do %Livebook.Hubs.Metadata{ id: local.id, name: local.hub_name, @@ -23,4 +23,12 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Local do def connect(_local), do: nil def connected?(_local), do: false + + def disconnect(_local), do: :ok + + def capabilities(_local), do: [] + + def get_secrets(_local), do: [] + + def create_secret(_local, _secret), do: :ok end diff --git a/lib/livebook/hubs/provider.ex b/lib/livebook/hubs/provider.ex index 833fa1a76..e124aca80 100644 --- a/lib/livebook/hubs/provider.ex +++ b/lib/livebook/hubs/provider.ex @@ -1,33 +1,63 @@ defprotocol Livebook.Hubs.Provider do @moduledoc false + alias Livebook.Secrets.Secret + + @type capability :: :connect | :secrets + @type capabilities :: list(capability()) + @type changeset_errors :: %{required(:errors) => list({String.t(), {Stirng.t(), list()}})} + @doc """ - Normalize given struct to `Livebook.Hubs.Metadata` struct. + Normalize given hub to `Livebook.Hubs.Metadata` struct. """ @spec normalize(struct()) :: Livebook.Hubs.Metadata.t() def normalize(struct) @doc """ - Loads fields into given struct. + Loads fields into given hub. """ @spec load(struct(), map() | keyword()) :: struct() def load(struct, fields) @doc """ - Gets the type from struct. + Gets the type from hub. """ @spec type(struct()) :: String.t() def type(struct) @doc """ - Gets the child spec of the given struct. + Gets the child spec of the given hub. """ @spec connect(struct()) :: Supervisor.child_spec() | module() | {module(), any()} | nil def connect(struct) @doc """ - Gets the connection status of the given struct. + Gets the connection status of the given hub. """ @spec connected?(struct()) :: boolean() def connected?(struct) + + @doc """ + Disconnects the given hub. + """ + @spec disconnect(struct()) :: :ok + def disconnect(struct) + + @doc """ + Gets the capabilities of the given hub. + """ + @spec capabilities(struct()) :: capabilities() + def capabilities(struct) + + @doc """ + Gets the secrets of the given hub. + """ + @spec get_secrets(struct()) :: list(Secret.t()) + def get_secrets(struct) + + @doc """ + Creates a secret of the given hub. + """ + @spec create_secret(struct(), Secret.t()) :: :ok | {:error, changeset_errors()} + def create_secret(struct, secret) end diff --git a/lib/livebook/secrets.ex b/lib/livebook/secrets.ex index abab1141b..b6c756988 100644 --- a/lib/livebook/secrets.ex +++ b/lib/livebook/secrets.ex @@ -31,6 +31,16 @@ defmodule Livebook.Secrets do to_struct(fields) end + @doc """ + Gets a secret from storage. + """ + @spec get_secret(String.t()) :: {:ok, Secret.t()} | :error + def get_secret(id) do + with {:ok, fields} <- Storage.fetch(:secrets, id) do + {:ok, to_struct(fields)} + end + end + @doc """ Checks if the secret already exists. """ @@ -66,8 +76,7 @@ defmodule Livebook.Secrets do """ @spec unset_secret(String.t()) :: :ok def unset_secret(id) do - if secret_exists?(id) do - secret = fetch_secret!(id) + with {:ok, secret} <- get_secret(id) do Storage.delete(:secrets, id) broadcast_secrets_change({:unset_secret, secret}) end diff --git a/lib/livebook_web/live/hub/new/enterprise_component.ex b/lib/livebook_web/live/hub/new/enterprise_component.ex index 36fee4243..82e558619 100644 --- a/lib/livebook_web/live/hub/new/enterprise_component.ex +++ b/lib/livebook_web/live/hub/new/enterprise_component.ex @@ -122,7 +122,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do receive do {:hub_connection_failed, reason} -> - EnterpriseClient.stop(pid) + EnterpriseClient.stop(base.id) {:noreply, socket @@ -141,7 +141,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do {:noreply, assign(socket, pid: pid, changeset: changeset, base: base)} {:error, reason} -> - EnterpriseClient.stop(pid) + EnterpriseClient.stop(base.id) {:noreply, socket diff --git a/lib/livebook_web/live/session_live.ex b/lib/livebook_web/live/session_live.ex index 39c4cfebf..8a6bfbbf9 100644 --- a/lib/livebook_web/live/session_live.ex +++ b/lib/livebook_web/live/session_live.ex @@ -9,7 +9,6 @@ defmodule LivebookWeb.SessionLive do alias Livebook.Notebook.{Cell, ContentLoader} alias Livebook.JSInterop alias Livebook.Hubs - alias Livebook.Hubs.EnterpriseClient on_mount LivebookWeb.SidebarHook @@ -2277,11 +2276,6 @@ defmodule LivebookWeb.SessionLive do end defp get_saved_secrets do - hub_secrets = - for connected_hub <- Hubs.get_connected_hubs(), - secret <- EnterpriseClient.list_cached_secrets(connected_hub.pid), - do: secret - - Enum.sort(hub_secrets ++ Secrets.get_secrets()) + Enum.sort(Hubs.get_secrets() ++ Secrets.get_secrets()) end end diff --git a/lib/livebook_web/live/session_live/secrets_component.ex b/lib/livebook_web/live/session_live/secrets_component.ex index 28ef737b3..228e7e4ec 100644 --- a/lib/livebook_web/live/session_live/secrets_component.ex +++ b/lib/livebook_web/live/session_live/secrets_component.ex @@ -1,7 +1,6 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do use LivebookWeb, :live_component - alias Livebook.Hubs.EnterpriseClient alias Livebook.Secrets.Secret @impl true @@ -9,7 +8,7 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do socket = socket |> assign(assigns) - |> assign(connected_hubs: Livebook.Hubs.get_connected_hubs()) + |> assign(hubs: Livebook.Hubs.get_hubs([:secrets])) prefill_form = prefill_secret_name(socket) @@ -127,17 +126,12 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do <%= if Livebook.Config.feature_flag_enabled?(:hub) do %> <%= label class: "flex items-center gap-2 text-gray-600" do %> <%= radio_button(f, :store, "hub", - disabled: @connected_hubs == [], + disabled: @hubs == [], checked: @data["store"] == "hub" ) %> in the Hub <% end %> <%= if @data["store"] == "hub" do %> - <%= select( - f, - :connected_hub, - connected_hubs_options(@connected_hubs, @data["connected_hub"]), - class: "input" - ) %> + <%= select(f, :hub_id, hubs_options(@hubs, @data["hub_id"]), class: "input") %> <% end %> <% end %> @@ -323,7 +317,7 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do defp build_origin(%{"store" => "session"}), do: :session defp build_origin(%{"store" => "app"}), do: :app - defp build_origin(%{"store" => "hub", "connected_hub" => id}), do: {:hub, id} + defp build_origin(%{"store" => "hub", "hub_id" => id}), do: {:hub, id} defp build_attrs(%{"name" => name, "value" => value} = attrs) do %{name: name, value: value, origin: build_origin(attrs)} @@ -338,21 +332,8 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do Livebook.Session.set_secret(socket.assigns.session.pid, secret) end - defp set_secret(socket, %Secret{origin: {:hub, id}} = secret) when is_binary(id) do - if hub = Enum.find(socket.assigns.connected_hubs, &(&1.hub.id == id)) do - create_secret_request = - LivebookProto.CreateSecretRequest.new!( - name: secret.name, - value: secret.value - ) - - case EnterpriseClient.send_request(hub.pid, create_secret_request) do - {:create_secret, _} -> :ok - {:error, reason} -> {:error, put_flash(socket, :error, reason)} - end - else - {:error, %{errors: [{"connected_hub", {"can't be blank", []}}]}} - end + defp set_secret(_socket, %Secret{origin: {:hub, id}} = secret) when is_binary(id) do + Livebook.Hubs.create_secret(secret) end defp grant_access(secrets, secret_name, origin, socket) do @@ -385,11 +366,10 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do Enum.any?(socket.assigns.saved_secrets, &(&1.name == secret_name)) end - # TODO: Livebook.Hubs.fetch_hubs_with_secrets_storage() - defp connected_hubs_options(connected_hubs, selected_hub) do + defp hubs_options(hubs, hub_id) do [[key: "Select one Hub", value: "", selected: true, disabled: true]] ++ - for %{hub: %{id: id, hub_name: name}} <- connected_hubs do - [key: name, value: id, selected: id == selected_hub] + for hub <- hubs do + [key: hub.hub_name, value: hub.id, selected: hub.id == hub_id] end end end diff --git a/test/livebook_web/live/session_live/secrets_component_test.exs b/test/livebook_web/live/session_live/secrets_component_test.exs index 09e9e68d5..8877c65a0 100644 --- a/test/livebook_web/live/session_live/secrets_component_test.exs +++ b/test/livebook_web/live/session_live/secrets_component_test.exs @@ -8,8 +8,7 @@ defmodule LivebookWeb.SessionLive.SecretsComponentTest do describe "enterprise" do setup %{url: url, token: token} do - node = EnterpriseServer.get_node() - id = :erpc.call(node, Enterprise.Integration, :fetch_env!, ["ENTERPRISE_ID"]) + id = Livebook.Utils.random_short_id() hub_id = "enterprise-#{id}" Livebook.Hubs.subscribe([:connection, :secrets]) @@ -65,7 +64,7 @@ defmodule LivebookWeb.SessionLive.SecretsComponentTest do name: secret.name, value: secret.value, store: "hub", - connected_hub: enterprise.id + hub_id: enterprise.id } }