Store autogenerated user id seprate from identity data (#2879)

lib/livebook_web/live/hooks/user_hook.ex

@@ -4,16 +4,16 @@ defmodule LivebookWeb.UserHook do
 
   alias Livebook.Users.User
 
-  def on_mount(:default, _params, %{"identity_data" => identity_data} = session, socket) do
+  def on_mount(:default, _params, session, socket) do
-    if connected?(socket) do
-      Livebook.Users.subscribe(identity_data.id)
-    end
-
     socket =
       socket
       |> assign_new(:current_user, fn -> build_current_user(session, socket) end)
       |> attach_hook(:current_user_subscription, :handle_info, &info/2)
 
+    if connected?(socket) do
+      Livebook.Users.subscribe(socket.assigns.current_user.id)
+    end
+
     {:cont, socket}
   end
 
@@ -43,7 +43,7 @@ defmodule LivebookWeb.UserHook do
         attrs -> attrs
       end
 
-    user = User.new(attrs["id"])
+    user = User.new(session["user_id"])
 
     case Livebook.Users.update_user(user, attrs) do
       {:ok, user} -> user

lib/livebook_web/plugs/user_plug.ex

@@ -37,11 +37,11 @@ defmodule LivebookWeb.UserPlug do
 
       identity_data ->
         # Ensure we have a unique ID to identify this user/session.
-        id =
+        id = identity_data[:id] || get_session(conn, :user_id) || Livebook.Utils.random_long_id()
-          identity_data[:id] || get_session(conn, :identity_data)[:id] ||
-            Livebook.Utils.random_long_id()
 
-        put_session(conn, :identity_data, Map.put(identity_data, :id, id))
+        conn
+        |> put_session(:identity_data, identity_data)
+        |> put_session(:user_id, id)
 
       true ->
         conn

test/livebook_web/plugs/user_plug_test.exs

@@ -13,17 +13,17 @@ defmodule LivebookWeb.UserPlugTest do
       |> fetch_cookies()
       |> call()
 
-    assert get_session(conn, :identity_data)[:id] != nil
+    assert get_session(conn, :user_id) != nil
   end
 
   test "keeps user id in the session if present" do
     conn =
       conn(:get, "/")
-      |> init_test_session(%{identity_data: %{id: "valid_user_id"}})
+      |> init_test_session(%{user_id: "valid_user_id"})
       |> fetch_cookies()
       |> call()
 
-    assert get_session(conn, :identity_data)[:id] != nil
+    assert get_session(conn, :user_id) != nil
   end
 
   test "given no user_data cookie, generates and stores new data" do