Allow users to assign folders to their apps (#3088)

Author: aleDsz

lib/livebook/hubs/personal.ex

@@ -281,4 +281,6 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Personal do
   def deployment_groups(_personal), do: nil
 
   def get_app_specs(_personal), do: []
+
+  def get_app_folders(_personal), do: []
 end

lib/livebook/hubs/provider.ex

@@ -155,4 +155,10 @@ defprotocol Livebook.Hubs.Provider do
   """
   @spec get_app_specs(t()) :: list(Livebook.Apps.AppSpec.t())
   def get_app_specs(hub)
+
+  @doc """
+  Gets the app folders from the given hub.
+  """
+  @spec get_app_folders(t()) :: list(%{id: String.t(), name: String.t()})
+  def get_app_folders(hub)
 end

lib/livebook/hubs/team.ex

@@ -259,6 +259,12 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
     end
   end
 
+  def get_app_folders(team) do
+    team.id
+    |> TeamClient.get_app_folders()
+    |> Enum.sort_by(& &1.name)
+  end
+
   defp parse_secret_errors(errors_map) do
     Teams.Requests.to_error_list(Secret, errors_map)
   end

lib/livebook/hubs/team_client.ex

@@ -24,6 +24,7 @@ defmodule Livebook.Hubs.TeamClient do
     deployment_groups: [],
     app_deployments: [],
     agents: [],
+    app_folders: [],
     app_deployment_statuses: nil
   ]
 
@@ -172,6 +173,14 @@ defmodule Livebook.Hubs.TeamClient do
     GenServer.call(registry_name(id), {:user_can_deploy?, user_id, deployment_group_id})
   end
 
+  @doc """
+  Returns a list of cached app folders.
+  """
+  @spec get_app_folders(String.t()) :: list(Teams.AppFolder.t())
+  def get_app_folders(id) do
+    GenServer.call(registry_name(id), :get_app_folders)
+  end
+
   @doc """
   Returns if the Team client is connected.
   """
@@ -370,6 +379,10 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
+  def handle_call(:get_app_folders, _caller, state) do
+    {:reply, state.app_folders, state}
+  end
+
   @impl true
   def handle_info(:connected, state) do
     Hubs.Broadcasts.hub_connected(state.hub.id)
@@ -611,7 +624,8 @@ defmodule Livebook.Hubs.TeamClient do
       file: nil,
       deployed_by: app_deployment.deployed_by,
       deployed_at: DateTime.from_gregorian_seconds(app_deployment.deployed_at),
-      authorization_groups: authorization_groups
+      authorization_groups: authorization_groups,
+      app_folder_id: nullify(app_deployment.app_folder_id)
     }
   end
 
@@ -630,7 +644,8 @@ defmodule Livebook.Hubs.TeamClient do
     for authorization_group <- authorization_groups do
       %Teams.AuthorizationGroup{
         provider_id: authorization_group.provider_id,
-        group_name: authorization_group.group_name
+        group_name: authorization_group.group_name,
+        app_folder_id: nullify(authorization_group.app_folder_id)
       }
     end
   end
@@ -664,6 +679,24 @@ defmodule Livebook.Hubs.TeamClient do
     }
   end
 
+  defp put_app_folder(state, app_folder) do
+    state = remove_app_folder(state, app_folder)
+
+    %{state | app_folders: [app_folder | state.app_folders]}
+  end
+
+  defp remove_app_folder(state, app_folder) do
+    %{state | app_folders: Enum.reject(state.app_folders, &(&1.id == app_folder.id))}
+  end
+
+  defp build_app_folder(state, %LivebookProto.AppFolder{} = app_folder) do
+    %Teams.AppFolder{
+      id: app_folder.id,
+      name: app_folder.name,
+      hub_id: state.hub.id
+    }
+  end
+
   defp handle_event(:secret_created, %Secrets.Secret{} = secret, state) do
     Hubs.Broadcasts.secret_created(secret)
 
@@ -787,6 +820,7 @@ defmodule Livebook.Hubs.TeamClient do
     |> dispatch_deployment_groups(user_connected)
     |> dispatch_app_deployments(user_connected)
     |> dispatch_agents(user_connected)
+    |> dispatch_app_folders(user_connected)
     |> dispatch_connection()
   end
 
@@ -798,6 +832,7 @@ defmodule Livebook.Hubs.TeamClient do
     |> dispatch_deployment_groups(agent_connected)
     |> dispatch_app_deployments(agent_connected)
     |> dispatch_agents(agent_connected)
+    |> dispatch_app_folders(agent_connected)
     |> dispatch_connection()
   end
 
@@ -873,6 +908,43 @@ defmodule Livebook.Hubs.TeamClient do
     update_hub(state, org_updated)
   end
 
+  defp handle_event(:app_folder_created, %Teams.AppFolder{} = app_folder, state) do
+    Teams.Broadcasts.app_folder_created(app_folder)
+    put_app_folder(state, app_folder)
+  end
+
+  defp handle_event(:app_folder_created, app_folder_created, state) do
+    handle_event(
+      :app_folder_created,
+      build_app_folder(state, app_folder_created.app_folder),
+      state
+    )
+  end
+
+  defp handle_event(:app_folder_updated, %Teams.AppFolder{} = app_folder, state) do
+    Teams.Broadcasts.app_folder_updated(app_folder)
+    put_app_folder(state, app_folder)
+  end
+
+  defp handle_event(:app_folder_updated, app_folder_updated, state) do
+    handle_event(
+      :app_folder_updated,
+      build_app_folder(state, app_folder_updated.app_folder),
+      state
+    )
+  end
+
+  defp handle_event(:app_folder_deleted, %Teams.AppFolder{} = app_folder, state) do
+    Teams.Broadcasts.app_folder_deleted(app_folder)
+    remove_app_folder(state, app_folder)
+  end
+
+  defp handle_event(:app_folder_deleted, %{id: id}, state) do
+    with {:ok, app_folder} <- fetch_app_folder(id, state) do
+      handle_event(:app_folder_deleted, app_folder, state)
+    end
+  end
+
   defp dispatch_secrets(state, %{secrets: secrets}) do
     decrypted_secrets = Enum.map(secrets, &build_secret(state, &1))
 
@@ -936,6 +1008,19 @@ defmodule Livebook.Hubs.TeamClient do
     dispatch_events(state, agent_joined: joined, agent_left: left)
   end
 
+  defp dispatch_app_folders(state, %{app_folders: app_folders}) do
+    app_folders = Enum.map(app_folders, &build_app_folder(state, &1))
+
+    {created, deleted, updated} =
+      diff(state.app_folders, app_folders, &(&1.id == &2.id))
+
+    dispatch_events(state,
+      app_folder_deleted: deleted,
+      app_folder_created: created,
+      app_folder_updated: updated
+    )
+  end
+
   defp dispatch_connection(%{hub: %{id: id}} = state) do
     Teams.Broadcasts.client_connected(id)
     state
@@ -1064,6 +1149,8 @@ defmodule Livebook.Hubs.TeamClient do
   defp fetch_app_deployment_from_slug(slug, state),
     do: fetch_entry(state.app_deployments, &(&1.slug == slug), state)
 
+  defp fetch_app_folder(id, state), do: fetch_entry(state.app_folders, &(&1.id == id), state)
+
   defp fetch_entry(entries, fun, state) do
     if entry = Enum.find(entries, fun) do
       {:ok, entry}

lib/livebook/live_markdown/export.ex

@@ -113,7 +113,8 @@ defmodule Livebook.LiveMarkdown.Export do
       :auto_shutdown_ms,
       :access_type,
       :show_source,
-      :output_type
+      :output_type,
+      :app_folder_id
     ]
 
     put_unless_default(

lib/livebook/live_markdown/import.ex

@@ -495,6 +495,9 @@ defmodule Livebook.LiveMarkdown.Import do
       {"show_source", show_source}, attrs ->
         Map.put(attrs, :show_source, show_source)
 
+      {"app_folder_id", app_folder_id}, attrs ->
+        Map.put(attrs, :app_folder_id, app_folder_id)
+
       {"output_type", output_type}, attrs when output_type in ["all", "rich"] ->
         Map.put(attrs, :output_type, String.to_atom(output_type))
 
@@ -664,7 +667,25 @@ defmodule Livebook.LiveMarkdown.Import do
     # validate it against the public key).
     teams_enabled = is_struct(hub, Livebook.Hubs.Team) and (hub.offline == nil or stamp_verified?)
 
-    {%{notebook | teams_enabled: teams_enabled}, stamp_verified?, messages}
+    {app_settings, messages} =
+      if app_folder_id = notebook.app_settings.app_folder_id do
+        app_folders = Hubs.Provider.get_app_folders(hub)
+
+        if Enum.any?(app_folders, &(&1.id == app_folder_id)) do
+          {notebook.app_settings, messages}
+        else
+          {Map.replace!(notebook.app_settings, :app_folder_id, nil),
+           messages ++
+             [
+               "notebook is assigned to a non-existent app folder, defaulting to ungrouped app folder"
+             ]}
+        end
+      else
+        {notebook.app_settings, messages}
+      end
+
+    {%{notebook | app_settings: app_settings, teams_enabled: teams_enabled}, stamp_verified?,
+     messages}
   end
 
   defp safe_binary_split(binary, offset)

lib/livebook/notebook/app_settings.ex

@@ -14,7 +14,8 @@ defmodule Livebook.Notebook.AppSettings do
           access_type: access_type(),
           password: String.t() | nil,
           show_source: boolean(),
-          output_type: output_type()
+          output_type: output_type(),
+          app_folder_id: String.t() | nil
         }
 
   @type access_type :: :public | :protected
@@ -33,6 +34,7 @@ defmodule Livebook.Notebook.AppSettings do
     field :password, :string
     field :show_source, :boolean
     field :output_type, Ecto.Enum, values: [:all, :rich]
+    field :app_folder_id, :string
   end
 
   @doc """
@@ -49,7 +51,8 @@ defmodule Livebook.Notebook.AppSettings do
       access_type: :protected,
       password: generate_password(),
       show_source: false,
-      output_type: :all
+      output_type: :all,
+      app_folder_id: nil
     }
   end
 
@@ -82,7 +85,8 @@ defmodule Livebook.Notebook.AppSettings do
       :auto_shutdown_ms,
       :access_type,
       :show_source,
-      :output_type
+      :output_type,
+      :app_folder_id
     ])
     |> validate_required([
       :slug,

lib/livebook/session.ex

@@ -904,6 +904,7 @@ defmodule Livebook.Session do
   def init({caller_pid, opts}) do
     Livebook.Settings.subscribe()
     Livebook.Hubs.Broadcasts.subscribe([:crud, :secrets, :file_systems])
+    Livebook.Teams.Broadcasts.subscribe(:app_folders)
 
     id = Keyword.fetch!(opts, :id)
 
@@ -2028,6 +2029,13 @@ defmodule Livebook.Session do
     {:noreply, handle_operation(state, operation)}
   end
 
+  def handle_info({event, app_folder}, state)
+      when event in [:app_folder_created, :app_folder_updated, :app_folder_deleted] and
+             app_folder.hub_id == state.data.notebook.hub_id do
+    operation = {:sync_hub_app_folders, @client_id}
+    {:noreply, handle_operation(state, operation)}
+  end
+
   def handle_info({:hub_deleted, id}, %{data: %{notebook: %{hub_id: id}}} = state) do
     # Since the hub got deleted, we close all sessions using that hub.
     # This way we clean up all secrets and other in-memory state that

lib/livebook/session/data.ex

@@ -37,6 +37,7 @@ defmodule Livebook.Session.Data do
     :secrets,
     :hub_secrets,
     :hub_file_systems,
+    :hub_app_folders,
     :mode,
     :deployed_app_slug,
     :app_data
@@ -247,6 +248,7 @@ defmodule Livebook.Session.Data do
           | {:set_notebook_hub, client_id(), String.t()}
           | {:sync_hub_secrets, client_id()}
           | {:sync_hub_file_systems, client_id()}
+          | {:sync_hub_app_folders, client_id()}
           | {:add_file_entries, client_id(), list(Notebook.file_entry())}
           | {:rename_file_entry, client_id(), name :: String.t(), new_name :: String.t()}
           | {:delete_file_entry, client_id(), String.t()}
@@ -305,6 +307,7 @@ defmodule Livebook.Session.Data do
     hub = Livebook.Hubs.fetch_hub!(notebook.hub_id)
     hub_secrets = Livebook.Hubs.get_secrets(hub)
     hub_file_systems = Livebook.Hubs.get_file_systems(hub)
+    hub_app_folders = Livebook.Hubs.Provider.get_app_folders(hub)
 
     startup_secrets =
       for secret <- Livebook.Secrets.get_startup_secrets(),
@@ -338,6 +341,7 @@ defmodule Livebook.Session.Data do
       secrets: secrets,
       hub_secrets: hub_secrets,
       hub_file_systems: hub_file_systems,
+      hub_app_folders: hub_app_folders,
       mode: opts[:mode],
       deployed_app_slug: nil,
       app_data: app_data
@@ -1074,6 +1078,14 @@ defmodule Livebook.Session.Data do
     |> wrap_ok()
   end
 
+  def apply_operation(data, {:sync_hub_app_folders, _client_id}) do
+    data
+    |> with_actions()
+    |> sync_hub_app_folders()
+    |> set_dirty()
+    |> wrap_ok()
+  end
+
   def apply_operation(data, {:add_file_entries, _client_id, file_entries}) do
     data
     |> with_actions()
@@ -1965,7 +1977,8 @@ defmodule Livebook.Session.Data do
           teams_enabled: is_struct(hub, Livebook.Hubs.Team)
       },
       hub_secrets: Livebook.Hubs.get_secrets(hub),
-      hub_file_systems: Livebook.Hubs.get_file_systems(hub)
+      hub_file_systems: Livebook.Hubs.get_file_systems(hub),
+      hub_app_folders: Livebook.Hubs.Provider.get_app_folders(hub)
     )
   end
 
@@ -1985,6 +1998,12 @@ defmodule Livebook.Session.Data do
     set!(data_actions, hub_file_systems: file_systems)
   end
 
+  defp sync_hub_app_folders({data, _} = data_actions) do
+    hub = Livebook.Hubs.fetch_hub!(data.notebook.hub_id)
+    app_folders = Livebook.Hubs.Provider.get_app_folders(hub)
+    set!(data_actions, hub_app_folders: app_folders)
+  end
+
   defp update_notebook_hub_secret_names({data, _} = data_actions) do
     hub_secret_names =
       for {_name, secret} <- data.secrets, secret.hub_id == data.notebook.hub_id, do: secret.name

lib/livebook/teams.ex

@@ -305,4 +305,12 @@ defmodule Livebook.Teams do
   def user_can_deploy?(%Team{} = team, %Teams.DeploymentGroup{} = deployment_group) do
     TeamClient.user_can_deploy?(team.id, team.user_id, deployment_group.id)
   end
+
+  @doc """
+  Gets a list of app folders for a given Hub.
+  """
+  @spec get_app_folders(Team.t()) :: list(Teams.AppFolder.t())
+  def get_app_folders(team) do
+    Hubs.Provider.get_app_folders(team)
+  end
 end