Restrict teams features based on org status (#2897)

lib/livebook/hubs/team.ex

@@ -41,6 +41,7 @@ defmodule Livebook.Hubs.Team do
     field :session_token, :string, redact: true
     field :hub_name, :string
     field :hub_emoji, :string
+    field :disabled, :boolean, default: false
 
     embeds_one :offline, Offline
   end

lib/livebook/hubs/team_client.ex

@@ -641,6 +641,7 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp handle_event(:user_connected, user_connected, state) do
     state
+    |> update_hub(user_connected)
     |> dispatch_secrets(user_connected)
     |> dispatch_file_systems(user_connected)
     |> dispatch_deployment_groups(user_connected)
@@ -728,6 +729,10 @@ defmodule Livebook.Hubs.TeamClient do
     state
   end
 
+  defp handle_event(:org_updated, org_updated, state) do
+    update_hub(state, org_updated)
+  end
+
   defp dispatch_secrets(state, %{secrets: secrets}) do
     decrypted_secrets = Enum.map(secrets, &build_secret(state, &1))
 
@@ -796,14 +801,26 @@ defmodule Livebook.Hubs.TeamClient do
     state
   end
 
-  defp update_hub(state, %{public_key: org_public_key}) do
+  defp update_hub(state, %LivebookProto.UserConnected{org_disabled: disabled}) do
-    hub = %{state.hub | org_public_key: org_public_key}
+    update_hub(state, &put_in(&1.disabled, disabled))
+  end
 
-    if Livebook.Hubs.hub_exists?(hub.id) do
+  defp update_hub(state, %LivebookProto.OrgUpdated{disabled: disabled}) do
+    update_hub(state, &put_in(&1.disabled, disabled))
+  end
+
+  defp update_hub(state, %LivebookProto.AgentConnected{public_key: org_public_key}) do
+    update_hub(state, &put_in(&1.org_public_key, org_public_key))
+  end
+
+  defp update_hub(state, fun) when is_function(fun, 1) do
+    hub = fun.(state.hub)
+
+    if Hubs.hub_exists?(hub.id) do
       Hubs.save_hub(hub)
     end
 
-    %{state | hub: hub}
+    put_in(state.hub, hub)
   end
 
   defp diff(old_list, new_list, fun, deleted_fun \\ nil, updated_fun \\ nil) do

lib/livebook_web/live/hub/edit/personal_component.ex

@@ -123,6 +123,7 @@ defmodule LivebookWeb.Hub.Edit.PersonalComponent do
               id="hub-file-systems-list"
               hub_id={@hub.id}
               file_systems={@file_systems}
+              disabled={false}
             />
           </div>
 
@@ -195,6 +196,7 @@ defmodule LivebookWeb.Hub.Edit.PersonalComponent do
           hub={@hub}
           secret_name={@secret_name}
           secret_value={@secret_value}
+          disabled={false}
           return_to={~p"/hub/#{@hub.id}"}
         />
       </.modal>
@@ -210,6 +212,7 @@ defmodule LivebookWeb.Hub.Edit.PersonalComponent do
           module={LivebookWeb.Hub.FileSystemFormComponent}
           id="file-systems"
           hub={@hub}
+          disabled={false}
           file_system={@file_system}
           file_system_id={@file_system_id}
           return_to={~p"/hub/#{@hub.id}"}

lib/livebook_web/live/hub/edit/team_component.ex

@@ -63,6 +63,15 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
         {Provider.connection_status(@hub)}
       </LayoutComponents.topbar>
 
+      <LayoutComponents.topbar :if={@hub.disabled} variant="warning">
+        <h2>
+          Workspace disabled: your organization doesn't have an active subscription. Please contact your <.link
+            href={org_url(@hub, "/users")}
+            class="underline"
+          >org's admin</.link>.
+        </h2>
+      </LayoutComponents.topbar>
+
       <div class="p-4 md:px-12 md:py-7 max-w-screen-md mx-auto">
         <div id={"#{@id}-component"}>
           <div class="mb-8 flex flex-col space-y-2">
@@ -176,10 +185,15 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
                 secrets={@secrets}
                 edit_path={"hub/#{@hub.id}/secrets/edit"}
                 return_to={~p"/hub/#{@hub.id}"}
+                disabled={@hub.disabled}
               />
 
               <div>
-                <.button patch={~p"/hub/#{@hub.id}/secrets/new"} id="add-secret">
+                <.button
+                  patch={~p"/hub/#{@hub.id}/secrets/new"}
+                  id="add-secret"
+                  disabled={@hub.disabled}
+                >
                   Add secret
                 </.button>
               </div>
@@ -200,6 +214,7 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
                 hub_id={@hub.id}
                 file_systems={@file_systems}
                 target={@myself}
+                disabled={@hub.disabled}
               />
             </div>
 
@@ -233,7 +248,11 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
               </div>
 
               <div>
-                <.button patch={~p"/hub/#{@hub.id}/groups/new"} id="add-deployment-group">
+                <.button
+                  patch={~p"/hub/#{@hub.id}/groups/new"}
+                  id="add-deployment-group"
+                  disabled={@hub.disabled}
+                >
                   Add deployment group
                 </.button>
               </div>
@@ -289,6 +308,7 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
           secret_name={@secret_name}
           secret_value={@secret_value}
           return_to={~p"/hub/#{@hub.id}"}
+          disabled={@hub.disabled}
         />
       </.modal>
 
@@ -303,6 +323,7 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
           module={LivebookWeb.Hub.FileSystemFormComponent}
           id="file-systems"
           hub={@hub}
+          disabled={@hub.disabled}
           file_system={@file_system}
           file_system_id={@file_system_id}
           return_to={~p"/hub/#{@hub.id}"}

lib/livebook_web/live/hub/edit_live.ex

@@ -11,6 +11,7 @@ defmodule LivebookWeb.Hub.EditLive do
   def mount(_params, _session, socket) do
     if connected?(socket) do
       Hubs.Broadcasts.subscribe([:connection])
+
       Livebook.Teams.Broadcasts.subscribe([:deployment_groups, :app_deployments, :agents])
     end
 
@@ -111,6 +112,10 @@ defmodule LivebookWeb.Hub.EditLive do
     {:noreply, load_hub(socket, id)}
   end
 
+  def handle_info({:hub_changed, id}, %{assigns: %{hub: %{id: id}}} = socket) do
+    {:noreply, load_hub(socket, id)}
+  end
+
   def handle_info(_message, socket) do
     {:noreply, socket}
   end

lib/livebook_web/live/hub/file_system_form_component.ex

@@ -88,7 +88,7 @@ defmodule LivebookWeb.Hub.FileSystemFormComponent do
             </p>
           <% end %>
           <div class="flex space-x-2">
-            <.button type="submit" disabled={not @changeset.valid?}>
+            <.button type="submit" disabled={@disabled or not @changeset.valid?}>
               <.remix_icon icon={@button.icon} />
               <span class="font-normal">{@button.label}</span>
             </.button>

lib/livebook_web/live/hub/file_system_list_component.ex

@@ -60,7 +60,11 @@ defmodule LivebookWeb.Hub.FileSystemListComponent do
         </div>
       </div>
       <div class="flex">
-        <.button patch={~p"/hub/#{@hub_id}/file-systems/new"} id="add-file-system">
+        <.button
+          patch={~p"/hub/#{@hub_id}/file-systems/new"}
+          id="add-file-system"
+          disabled={@disabled}
+        >
           Add file storage
         </.button>
       </div>

lib/livebook_web/live/hub/secret_form_component.ex

@@ -75,7 +75,7 @@ defmodule LivebookWeb.Hub.SecretFormComponent do
             <.hidden_field field={f[:hub_id]} value={@hub.id} />
             <.hidden_field field={f[:deployment_group_id]} value={@deployment_group_id} />
             <div class="flex space-x-2">
-              <.button type="submit" disabled={not @changeset.valid?}>
+              <.button type="submit" disabled={@disabled or not @changeset.valid?}>
                 <.remix_icon icon={@button.icon} />
                 <span class="font-normal">{@button.label}</span>
               </.button>

proto/lib/livebook_proto/agent_connected.pb.ex

@@ -18,4 +18,5 @@ defmodule LivebookProto.AgentConnected do
     json_name: "appDeployments"
 
   field :agents, 9, repeated: true, type: LivebookProto.Agent
+  field :org_disabled, 10, type: :bool, json_name: "orgDisabled"
 end

proto/lib/livebook_proto/event.pb.ex

@@ -71,4 +71,6 @@ defmodule LivebookProto.Event do
     type: LivebookProto.AppDeploymentStopped,
     json_name: "appDeploymentStopped",
     oneof: 0
+
+  field :org_updated, 17, type: LivebookProto.OrgUpdated, json_name: "orgUpdated", oneof: 0
 end

proto/lib/livebook_proto/org_updated.pb.ex

@@ -0,0 +1,6 @@
+defmodule LivebookProto.OrgUpdated do
+  use Protobuf, protoc_gen_elixir_version: "0.13.0", syntax: :proto3
+
+  field :id, 1, type: :string
+  field :disabled, 2, type: :bool
+end

proto/lib/livebook_proto/user_connected.pb.ex

@@ -16,4 +16,5 @@ defmodule LivebookProto.UserConnected do
     json_name: "appDeployments"
 
   field :agents, 6, repeated: true, type: LivebookProto.Agent
+  field :org_disabled, 7, type: :bool, json_name: "orgDisabled"
 end

proto/messages.proto

@@ -110,6 +110,7 @@ message UserConnected {
   repeated DeploymentGroup deployment_groups = 4;
   repeated AppDeployment app_deployments = 5;
   repeated Agent agents = 6;
+  bool org_disabled = 7;
 }
 
 message AgentConnected {
@@ -121,6 +122,7 @@ message AgentConnected {
   repeated DeploymentGroup deployment_groups = 7;
   repeated AppDeployment app_deployments = 8;
   repeated Agent agents = 9;
+  bool org_disabled = 10;
 }
 
 message AppDeployment {
@@ -157,6 +159,11 @@ message AgentLeft {
   string id = 1;
 }
 
+message OrgUpdated {
+  string id = 1;
+  bool disabled = 2;
+}
+
 message Agent {
   string id = 1;
   string name = 2;
@@ -210,5 +217,6 @@ message Event {
     AgentJoined agent_joined = 14;
     AgentLeft agent_left = 15;
     AppDeploymentStopped app_deployment_stopped = 16;
+    OrgUpdated org_updated = 17;
   }
 }