Repository: livebook-dev/kino
Pull request #340 (2023-09-13 05:06:43)

Global attrs and global name and cookie for RemoteCell
I did the synchronization between secrets and plain text implicit in `Global`. That way, following our pattern of using the suffix "_secret", the SmartCell wouldn't have to worry about synchronizing it manually when using globals.
WDYT?

cristineguadelupe: I also wondered if it wouldn't be better to have a single module for global and counter 🤔 
jonatanklosko: > I also wondered if it wouldn't be better to have a single module for global and counter 🤔

Sounds good, we can reuse the same ETS table :)

Patch:
diff --git a/lib/kino/application.ex b/lib/kino/application.ex
index de6c174..d154830 100644
--- a/lib/kino/application.ex
+++ b/lib/kino/application.ex
@@ -6,7 +6,7 @@ defmodule Kino.Application do
   def start(_type, _args) do
     Application.put_env(:elixir, :dbg_callback, {Kino.Debug, :dbg, []})
 
-    Kino.Counter.initialize()
+    Kino.AttributeStore.initialize()
 
     Kino.SmartCell.register(Kino.RemoteExecutionCell)
 
diff --git a/lib/kino/attribute_store.ex b/lib/kino/attribute_store.ex
new file mode 100644
index 0000000..1c27a9a
--- /dev/null
+++ b/lib/kino/attribute_store.ex
@@ -0,0 +1,54 @@
+defmodule Kino.AttributeStore do
+  @moduledoc false
+
+  @table_name __MODULE__
+
+  @doc """
+  Initializes resources for global attrs.
+  """
+  @spec initialize() :: :ok
+  def initialize() do
+    :ets.new(@table_name, [:set, :named_table, :public])
+    :ok
+  end
+
+  @doc """
+  Increments the given counter and returns the new value.
+  """
+  @spec counter_next(term()) :: integer()
+  def counter_next(key) do
+    :ets.update_counter(@table_name, key, {2, 1}, {key, 0})
+  end
+
+  @doc """
+  Sets the counter to `value` unless it already has a higher value.
+
+  Returns the new counter value.
+  """
+  @spec counter_put_max(term(), integer()) :: integer()
+  def counter_put_max(key, value) do
+    [_, counter] =
+      :ets.update_counter(@table_name, key, [{2, -1, value, value - 1}, {2, 1}], {key, 0})
+
+    counter
+  end
+
+  @doc """
+  Puts the shared attribute value.
+  """
+  @spec put_attribute(term(), term()) :: boolean()
+  def put_attribute(key, value) do
+    :ets.insert(@table_name, {key, value})
+  end
+
+  @doc """
+  Returns the attribute value for a given key.
+  """
+  @spec get_attribute(term()) :: term()
+  def get_attribute(key, default \\ nil) do
+    case :ets.lookup(@table_name, key) do
+      [] -> default
+      [{_key, value} | _] -> value
+    end
+  end
+end
diff --git a/lib/kino/counter.ex b/lib/kino/counter.ex
deleted file mode 100644
index cbb5098..0000000
--- a/lib/kino/counter.ex
+++ /dev/null
@@ -1,35 +0,0 @@
-defmodule Kino.Counter do
-  @moduledoc false
-
-  @table_name __MODULE__
-
-  @doc """
-  Initializes resources for global counters.
-  """
-  @spec initialize() :: :ok
-  def initialize() do
-    :ets.new(@table_name, [:set, :named_table, :public])
-    :ok
-  end
-
-  @doc """
-  Increments the given counter and returns the new value.
-  """
-  @spec next(term()) :: integer()
-  def next(key) do
-    :ets.update_counter(@table_name, key, {2, 1}, {key, 0})
-  end
-
-  @doc """
-  Sets the counter to `value` unless it already has a higher value.
-
-  Returns the new counter value.
-  """
-  @spec put_max(term(), integer()) :: integer()
-  def put_max(key, value) do
-    [_, counter] =
-      :ets.update_counter(@table_name, key, [{2, -1, value, value - 1}, {2, 1}], {key, 0})
-
-    counter
-  end
-end
diff --git a/lib/kino/remote_ execution_cell.ex b/lib/kino/remote_execution_cell.ex
similarity index 70%
rename from lib/kino/remote_ execution_cell.ex
rename to lib/kino/remote_execution_cell.ex
index fcc7027..e0b769b 100644
--- a/lib/kino/remote_ execution_cell.ex
+++ b/lib/kino/remote_execution_cell.ex
@@ -5,16 +5,24 @@ defmodule Kino.RemoteExecutionCell do
   use Kino.JS.Live
   use Kino.SmartCell, name: "Remote execution"
 
+  alias Kino.AttributeStore
+
   @default_code ":ok"
+  @global_key __MODULE__
+  @global_attrs ["node", "cookie", "cookie_secret"]
 
   @impl true
   def init(attrs, ctx) do
+    {shared_cookie, shared_cookie_secret} =
+      AttributeStore.get_attribute({@global_key, :cookie}, {nil, nil})
+
     fields = %{
       "assign_to" => attrs["assign_to"] || "",
-      "node" => attrs["node"] || "",
-      "cookie" => attrs["cookie"] || "",
-      "cookie_secret" => attrs["cookie_secret"] || "",
-      "use_cookie_secret" => Map.get(attrs, "use_cookie_secret", true)
+      "node" => attrs["node"] || AttributeStore.get_attribute({@global_key, :node}) || "",
+      "cookie" => attrs["cookie"] || shared_cookie || "",
+      "cookie_secret" => attrs["cookie_secret"] || shared_cookie_secret || "",
+      "use_cookie_secret" =>
+        if(shared_cookie, do: false, else: Map.get(attrs, "use_cookie_secret", true))
     }
 
     ctx = assign(ctx, fields: fields)
@@ -31,6 +39,7 @@ defmodule Kino.RemoteExecutionCell do
   @impl true
   def handle_event("update_field", %{"field" => field, "value" => value}, ctx) do
     ctx = update(ctx, :fields, &Map.put(&1, field, value))
+    if field in @global_attrs, do: put_shared_attr(field, value)
     broadcast_event(ctx, "update_field", %{"fields" => %{field => value}})
 
     {:noreply, ctx}
@@ -95,4 +104,16 @@ defmodule Kino.RemoteExecutionCell do
   end
 
   defp build_set_cookie(%{"cookie" => cookie}), do: String.to_atom(cookie)
+
+  defp put_shared_attr("cookie", value) do
+    AttributeStore.put_attribute({@global_key, :cookie}, {value, nil})
+  end
+
+  defp put_shared_attr("cookie_secret", value) do
+    AttributeStore.put_attribute({@global_key, :cookie}, {nil, value})
+  end
+
+  defp put_shared_attr(field, value) do
+    AttributeStore.put_attribute({@global_key, String.to_atom(field)}, value)
+  end
 end
diff --git a/lib/kino/smart_cell.ex b/lib/kino/smart_cell.ex
index 75db928..cb1accb 100644
--- a/lib/kino/smart_cell.ex
+++ b/lib/kino/smart_cell.ex
@@ -338,7 +338,7 @@ defmodule Kino.SmartCell do
   def prefixed_var_name(prefix, var_name)
 
   def prefixed_var_name(prefix, nil) do
-    case Kino.Counter.next(var_counter_key(prefix)) do
+    case Kino.AttributeStore.counter_next(var_counter_key(prefix)) do
       1 -> prefix
       n -> "#{prefix}#{n}"
     end
@@ -347,7 +347,7 @@ defmodule Kino.SmartCell do
   def prefixed_var_name(prefix, var_name) do
     with {:ok, suffix} <- parse_var_prefix(var_name, prefix),
          {:ok, n} <- parse_var_suffix(suffix) do
-      Kino.Counter.put_max(var_counter_key(prefix), n)
+      Kino.AttributeStore.counter_put_max(var_counter_key(prefix), n)
     end
 
     var_name


Test patch:
diff --git a/test/kino/attribute_store_test.exs b/test/kino/attribute_store_test.exs
new file mode 100644
index 0000000..93227de
--- /dev/null
+++ b/test/kino/attribute_store_test.exs
@@ -0,0 +1,50 @@
+defmodule Kino.AttributeStoreTets do
+  use ExUnit.Case, async: true
+
+  alias Kino.AttributeStore
+
+  describe "counter_next/2" do
+    test "returns the default value if the counter doesn't exist", %{test: test} do
+      assert AttributeStore.counter_next(test) == 1
+    end
+
+    test "increments the counter on every call", %{test: test} do
+      AttributeStore.counter_next(test)
+
+      assert AttributeStore.counter_next(test) == 2
+      assert AttributeStore.counter_next(test) == 3
+      assert AttributeStore.counter_next(test) == 4
+    end
+  end
+
+  describe "counter_put_max/2" do
+    test "sets the given value if the counter doesn't exist", %{test: test} do
+      assert AttributeStore.counter_put_max(test, 10) == 10
+    end
+
+    test "keeps existing value if the given one is lower or equal", %{test: test} do
+      AttributeStore.counter_put_max(test, 10)
+
+      assert AttributeStore.counter_put_max(test, 5) == 10
+      assert AttributeStore.counter_put_max(test, 10) == 10
+    end
+
+    test "sets the given value if higher than the current one", %{test: test} do
+      AttributeStore.counter_put_max(test, 10)
+
+      assert AttributeStore.counter_put_max(test, 11) == 11
+      assert AttributeStore.counter_put_max(test, 15) == 15
+    end
+  end
+
+  describe "get_attr/1" do
+    test "returns the default value if the shared attr doesn't exist", %{test: test} do
+      assert AttributeStore.get_attribute({test, :no_attr}) == nil
+    end
+
+    test "returns the value if the shared attr exist", %{test: test} do
+      AttributeStore.put_attribute({test, :shared_attr}, "plain value")
+      assert AttributeStore.get_attribute({test, :shared_attr}) == "plain value"
+    end
+  end
+end
diff --git a/test/kino/counter_test.exs b/test/kino/counter_test.exs
deleted file mode 100644
index ad6c3cd..0000000
--- a/test/kino/counter_test.exs
+++ /dev/null
@@ -1,39 +0,0 @@
-defmodule Kino.CounterTets do
-  use ExUnit.Case, async: true
-
-  alias Kino.Counter
-
-  describe "next/2" do
-    test "returns the default value if the counter doesn't exist", %{test: test} do
-      assert Counter.next(test) == 1
-    end
-
-    test "increments the counter on every call", %{test: test} do
-      Counter.next(test)
-
-      assert Counter.next(test) == 2
-      assert Counter.next(test) == 3
-      assert Counter.next(test) == 4
-    end
-  end
-
-  describe "bump/2" do
-    test "sets the given value if the counter doesn't exist", %{test: test} do
-      assert Counter.put_max(test, 10) == 10
-    end
-
-    test "keeps existing value if the given one is lower or equal", %{test: test} do
-      Counter.put_max(test, 10)
-
-      assert Counter.put_max(test, 5) == 10
-      assert Counter.put_max(test, 10) == 10
-    end
-
-    test "sets the given value if higher than the current one", %{test: test} do
-      Counter.put_max(test, 10)
-
-      assert Counter.put_max(test, 11) == 11
-      assert Counter.put_max(test, 15) == 15
-    end
-  end
-end
diff --git a/test/kino/remote_execution_cell_test.exs b/test/kino/remote_execution_cell_test.exs
index 3ceb82d..33a732b 100644
--- a/test/kino/remote_execution_cell_test.exs
+++ b/test/kino/remote_execution_cell_test.exs
@@ -4,6 +4,7 @@ defmodule Kino.RemoteExecutionCellTest do
   import Kino.Test
 
   alias Kino.RemoteExecutionCell
+  alias Kino.AttributeStore
 
   setup :configure_livebook_bridge
 
@@ -159,4 +160,68 @@ defmodule Kino.RemoteExecutionCellTest do
       assert RemoteExecutionCell.to_source(attrs) == ""
     end
   end
+
+  defmodule Global do
+    use ExUnit.Case, async: false
+
+    setup do
+      AttributeStore.put_attribute({Kino.RemoteExecutionCell, :node}, "name@node@global")
+
+      AttributeStore.put_attribute(
+        {Kino.RemoteExecutionCell, :cookie},
+        {"node-cookie-global", nil}
+      )
+
+      :ok
+    end
+
+    test "from prefilled attrs" do
+      {_kino, source} = start_smart_cell!(RemoteExecutionCell, %{})
+
+      assert source == """
+             node = :name@node@global
+             Node.set_cookie(node, :"node-cookie-global")
+             :erpc.call(node, fn -> :ok end)\
+             """
+    end
+
+    test "from prefilled attrs with cookie as a secret" do
+      AttributeStore.put_attribute(
+        {Kino.RemoteExecutionCell, :cookie},
+        {nil, "COOKIE_SECRET_GLOBAL"}
+      )
+
+      {_kino, source} = start_smart_cell!(RemoteExecutionCell, %{})
+
+      assert source == """
+             node = :name@node@global
+             Node.set_cookie(node, String.to_atom(System.fetch_env!("LB_COOKIE_SECRET_GLOBAL")))
+             :erpc.call(node, fn -> :ok end)\
+             """
+    end
+
+    test "attrs precedes globals" do
+      {_kino, source} = start_smart_cell!(RemoteExecutionCell, %{"node" => "name@node@attrs"})
+
+      assert source == """
+             node = :name@node@attrs
+             Node.set_cookie(node, :"node-cookie-global")
+             :erpc.call(node, fn -> :ok end)\
+             """
+    end
+
+    test "globals always come from the most recent edited cell" do
+      {kino, _source} = start_smart_cell!(RemoteExecutionCell, %{})
+      push_event(kino, "update_field", %{"field" => "node", "value" => "edited@node@name"})
+      assert_receive {:runtime_smart_cell_update, _, _, _, _}
+
+      {_kino, source} = start_smart_cell!(RemoteExecutionCell, %{})
+
+      assert source == """
+             node = :edited@node@name
+             Node.set_cookie(node, :"node-cookie-global")
+             :erpc.call(node, fn -> :ok end)\
+             """
+    end
+  end
 end