Load persisted environment variables from settings into session's runtime (#1409)

lib/livebook/runtime.ex

@@ -473,8 +473,14 @@ defprotocol Livebook.Runtime do
   def search_packages(runtime, send_to, search)
 
   @doc """
-  Adds Livebook secrets as environment variables
+  Sets the given environment variables.
   """
   @spec put_system_envs(t(), list({String.t(), String.t()})) :: :ok
-  def put_system_envs(runtime, secrets)
+  def put_system_envs(runtime, env_vars)
+
+  @doc """
+  Unsets the given environment variables.
+  """
+  @spec delete_system_envs(t(), list(String.t())) :: :ok
+  def delete_system_envs(runtime, names)
 end

lib/livebook/runtime/attached.ex

@@ -127,4 +127,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
   def put_system_envs(runtime, secrets) do
     RuntimeServer.put_system_envs(runtime.server_pid, secrets)
   end
+
+  def delete_system_envs(runtime, names) do
+    RuntimeServer.delete_system_envs(runtime.server_pid, names)
+  end
 end

lib/livebook/runtime/elixir_standalone.ex

@@ -224,4 +224,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do
   def put_system_envs(runtime, secrets) do
     RuntimeServer.put_system_envs(runtime.server_pid, secrets)
   end
+
+  def delete_system_envs(runtime, names) do
+    RuntimeServer.delete_system_envs(runtime.server_pid, names)
+  end
 end

lib/livebook/runtime/embedded.ex

@@ -126,6 +126,10 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do
     RuntimeServer.put_system_envs(runtime.server_pid, secrets)
   end
 
+  def delete_system_envs(runtime, names) do
+    RuntimeServer.delete_system_envs(runtime.server_pid, names)
+  end
+
   defp config() do
     Application.get_env(:livebook, Livebook.Runtime.Embedded, [])
   end

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -166,10 +166,22 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     GenServer.cast(pid, {:stop_smart_cell, ref})
   end
 
+  @doc """
+  Sets the given environment variables.
+  """
+  @spec put_system_envs(pid(), map()) :: :ok
   def put_system_envs(pid, secrets) do
     GenServer.cast(pid, {:put_system_envs, secrets})
   end
 
+  @doc """
+  Unsets the given environment variables.
+  """
+  @spec delete_system_envs(pid(), list(String.t())) :: :ok
+  def delete_system_envs(pid, names) do
+    GenServer.cast(pid, {:delete_system_envs, names})
+  end
+
   @doc """
   Stops the manager.
 
@@ -466,6 +478,13 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
 
   def handle_cast({:put_system_envs, secrets}, state) do
     System.put_env(secrets)
+
+    {:noreply, state}
+  end
+
+  def handle_cast({:delete_system_envs, names}, state) do
+    Enum.each(names, &System.delete_env/1)
+
     {:noreply, state}
   end
 

lib/livebook/session.ex

@@ -567,6 +567,7 @@ defmodule Livebook.Session do
 
   @impl true
   def init(opts) do
+    Livebook.Settings.subscribe()
     id = Keyword.fetch!(opts, :id)
 
     {:ok, worker_pid} = Livebook.Session.Worker.start_link(id)
@@ -583,6 +584,7 @@ defmodule Livebook.Session do
              else: :ok
            ) do
       state = schedule_autosave(state)
+
       {:ok, state}
     else
       {:error, error} ->
@@ -1121,6 +1123,22 @@ defmodule Livebook.Session do
     {:noreply, state}
   end
 
+  def handle_info({:env_var_set, env_var}, state) do
+    if Runtime.connected?(state.data.runtime) do
+      Runtime.put_system_envs(state.data.runtime, %{env_var.key => env_var.value})
+    end
+
+    {:noreply, state}
+  end
+
+  def handle_info({:env_var_unset, env_var}, state) do
+    if Runtime.connected?(state.data.runtime) do
+      Runtime.delete_system_envs(state.data.runtime, [env_var.key])
+    end
+
+    {:noreply, state}
+  end
+
   def handle_info(_message, state), do: {:noreply, state}
 
   @impl true
@@ -1330,6 +1348,8 @@ defmodule Livebook.Session do
   defp after_operation(state, _prev_state, {:set_runtime, _client_id, runtime}) do
     if Runtime.connected?(runtime) do
       set_runtime_secrets(state, state.data.secrets)
+      set_runtime_env_vars(state)
+
       state
     else
       state
@@ -1555,6 +1575,11 @@ defmodule Livebook.Session do
     Runtime.put_system_envs(state.data.runtime, secrets)
   end
 
+  defp set_runtime_env_vars(state) do
+    env_vars = Enum.map(Livebook.Settings.fetch_env_vars(), &{&1.key, &1.value})
+    Runtime.put_system_envs(state.data.runtime, env_vars)
+  end
+
   defp notify_update(state) do
     session = self_from_state(state)
     Livebook.Sessions.update_session(session)

lib/livebook/settings.ex

@@ -163,9 +163,9 @@ defmodule Livebook.Settings do
   end
 
   @doc """
-  Persists the given environment variable.
+  Sets the given environment variable.
 
-  With success, notifies interested processes about environment variables
+  With success, notifies interested processes about environment variable
   data change. Otherwise, it will return an error tuple with changeset.
   """
   @spec set_env_var(EnvVar.t(), map()) ::
@@ -182,20 +182,26 @@ defmodule Livebook.Settings do
     attributes = env_var |> Map.from_struct() |> Map.to_list()
 
     with :ok <- storage().insert(:env_vars, env_var.key, attributes),
-         :ok <- broadcast_env_vars_change() do
+         :ok <- broadcast_env_vars_change({:env_var_set, env_var}) do
       {:ok, env_var}
     end
   end
 
   @doc """
-  Deletes an environment variable from given id.
+  Unsets an environment variable from given id.
 
-  Also, it notifies interested processes about environment variables data change.
+  With success, notifies interested processes about environment variable
+  deletion. Otherwise, it does nothing.
   """
-  @spec delete_env_var(String.t()) :: :ok
-  def delete_env_var(id) do
-    storage().delete(:env_vars, id)
-    broadcast_env_vars_change()
+  @spec unset_env_var(String.t()) :: :ok
+  def unset_env_var(id) do
+    if env_var_exists?(id) do
+      env_var = fetch_env_var!(id)
+      storage().delete(:env_vars, id)
+      broadcast_env_vars_change({:env_var_unset, env_var})
+    end
+
+    :ok
   end
 
   @doc """
@@ -213,7 +219,8 @@ defmodule Livebook.Settings do
 
   ## Messages
 
-    * `{:env_vars_changed, env_vars}`
+    * `{:env_var_set, env_var}`
+    * `{:env_var_unset, env_var}`
 
   """
   @spec subscribe() :: :ok | {:error, term()}
@@ -231,8 +238,8 @@ defmodule Livebook.Settings do
 
   # Notifies interested processes about environment variables data change.
   #
-  # Broadcasts `{:env_vars_changed, env_vars}` message under the `"settings"` topic.
-  defp broadcast_env_vars_change do
-    Phoenix.PubSub.broadcast(Livebook.PubSub, "settings", {:env_vars_changed, fetch_env_vars()})
+  # Broadcasts given message under the `"settings"` topic.
+  defp broadcast_env_vars_change(message) do
+    Phoenix.PubSub.broadcast(Livebook.PubSub, "settings", message)
   end
 end

lib/livebook_web/live/settings_live.ex

@@ -14,7 +14,7 @@ defmodule LivebookWeb.SettingsLive do
     {:ok,
      assign(socket,
        file_systems: Livebook.Settings.file_systems(),
-       env_vars: Livebook.Settings.fetch_env_vars(),
+       env_vars: Livebook.Settings.fetch_env_vars() |> Enum.sort(),
        env_var: nil,
        autosave_path_state: %{
          file: autosave_dir(),
@@ -132,7 +132,6 @@ defmodule LivebookWeb.SettingsLive do
               env_vars={@env_vars}
               return_to={Routes.settings_path(@socket, :page)}
               add_env_var_path={Routes.settings_path(@socket, :add_env_var)}
-              target={@socket.view}
             />
           </div>
         </div>
@@ -340,7 +339,7 @@ defmodule LivebookWeb.SettingsLive do
   end
 
   def handle_event("delete_env_var", %{"env_var" => key}, socket) do
-    Livebook.Settings.delete_env_var(key)
+    Livebook.Settings.unset_env_var(key)
     {:noreply, socket}
   end
 
@@ -357,7 +356,20 @@ defmodule LivebookWeb.SettingsLive do
     handle_event("set_autosave_path", %{}, socket)
   end
 
-  def handle_info({:env_vars_changed, env_vars}, socket) do
+  def handle_info({:env_var_set, env_var}, socket) do
+    idx = Enum.find_index(socket.assigns.env_vars, &(&1.key == env_var.key))
+
+    env_vars =
+      if idx,
+        do: List.replace_at(socket.assigns.env_vars, idx, env_var),
+        else: [env_var | socket.assigns.env_vars]
+
+    {:noreply, assign(socket, env_vars: Enum.sort(env_vars), env_var: nil)}
+  end
+
+  def handle_info({:env_var_unset, env_var}, socket) do
+    env_vars = Enum.reject(socket.assigns.env_vars, &(&1.key == env_var.key))
+
     {:noreply, assign(socket, env_vars: env_vars, env_var: nil)}
   end
 

test/livebook/settings_test.exs

@@ -7,7 +7,7 @@ defmodule Livebook.SettingsTest do
     env_var = insert_env_var(:env_var)
     assert env_var in Settings.fetch_env_vars()
 
-    Settings.delete_env_var(env_var.key)
+    Settings.unset_env_var(env_var.key)
     refute env_var in Settings.fetch_env_vars()
   end
 
@@ -21,7 +21,7 @@ defmodule Livebook.SettingsTest do
     env_var = insert_env_var(:env_var, key: "123456")
     assert Settings.fetch_env_var!("123456") == env_var
 
-    Settings.delete_env_var("123456")
+    Settings.unset_env_var("123456")
   end
 
   test "env_var_exists?/1" do
@@ -29,7 +29,7 @@ defmodule Livebook.SettingsTest do
     insert_env_var(:env_var, key: "FOO")
     assert Settings.env_var_exists?("FOO")
 
-    Settings.delete_env_var("FOO")
+    Settings.unset_env_var("FOO")
   end
 
   describe "set_env_var/1" do
@@ -40,7 +40,7 @@ defmodule Livebook.SettingsTest do
       assert attrs.key == env_var.key
       assert attrs.value == env_var.value
 
-      Settings.delete_env_var(env_var.key)
+      Settings.unset_env_var(env_var.key)
     end
 
     test "updates an environment variable" do
@@ -51,7 +51,7 @@ defmodule Livebook.SettingsTest do
       assert env_var.key == updated_env_var.key
       assert updated_env_var.value == attrs.value
 
-      Settings.delete_env_var(env_var.key)
+      Settings.unset_env_var(env_var.key)
     end
 
     test "returns changeset error" do

test/livebook_web/live/session_live_test.exs

@@ -3,7 +3,7 @@ defmodule LivebookWeb.SessionLiveTest do
 
   import Phoenix.LiveViewTest
 
-  alias Livebook.{Sessions, Session, Runtime, Users, FileSystem}
+  alias Livebook.{Sessions, Session, Settings, Runtime, Users, FileSystem}
   alias Livebook.Notebook.Cell
 
   setup do
@@ -942,6 +942,53 @@ defmodule LivebookWeb.SessionLiveTest do
     end
   end
 
+  test "outputs persisted env var from ets", %{conn: conn, session: session} do
+    Session.subscribe(session.id)
+    {:ok, view, _} = live(conn, "/sessions/#{session.id}")
+
+    section_id = insert_section(session.pid)
+
+    cell_id =
+      insert_text_cell(session.pid, section_id, :code, ~s{System.get_env("MY_AWESOME_ENV")})
+
+    view
+    |> element(~s{[data-el-session]})
+    |> render_hook("queue_cell_evaluation", %{"cell_id" => cell_id})
+
+    assert_receive {:operation,
+                    {:add_cell_evaluation_response, _, ^cell_id, {:text, "\e[35mnil\e[0m"}, _}}
+
+    attrs = params_for(:env_var, key: "MY_AWESOME_ENV", value: "MyEnvVarValue")
+    Settings.set_env_var(attrs)
+
+    view
+    |> element(~s{[data-el-session]})
+    |> render_hook("queue_cell_evaluation", %{"cell_id" => cell_id})
+
+    assert_receive {:operation,
+                    {:add_cell_evaluation_response, _, ^cell_id,
+                     {:text, "\e[32m\"MyEnvVarValue\"\e[0m"}, _}}
+
+    Settings.set_env_var(%{attrs | value: "OTHER_VALUE"})
+
+    view
+    |> element(~s{[data-el-session]})
+    |> render_hook("queue_cell_evaluation", %{"cell_id" => cell_id})
+
+    assert_receive {:operation,
+                    {:add_cell_evaluation_response, _, ^cell_id,
+                     {:text, "\e[32m\"OTHER_VALUE\"\e[0m"}, _}}
+
+    Settings.unset_env_var("MY_AWESOME_ENV")
+
+    view
+    |> element(~s{[data-el-session]})
+    |> render_hook("queue_cell_evaluation", %{"cell_id" => cell_id})
+
+    assert_receive {:operation,
+                    {:add_cell_evaluation_response, _, ^cell_id, {:text, "\e[35mnil\e[0m"}, _}}
+  end
+
   # Helpers
 
   defp wait_for_session_update(session_pid) do

test/support/noop_runtime.ex

@@ -40,5 +40,6 @@ defmodule Livebook.Runtime.NoopRuntime do
     def search_packages(_, _, _), do: make_ref()
 
     def put_system_envs(_, _), do: :ok
+    def delete_system_envs(_, _), do: :ok
   end
 end