Adds ZTA and clustering to deployment groups (#2413)

Co-authored-by: José Valim <jose.valim@gmail.com>

lib/livebook/hubs/team_client.ex

@@ -244,9 +244,19 @@ defmodule Livebook.Hubs.TeamClient do
     }
   end
 
-  defp build_deployment_group(state, %{id: id, name: name, mode: mode, secrets: secrets}) do
-    secrets = Enum.map(secrets, &build_secret(state, &1))
-    %DeploymentGroup{id: id, name: name, mode: mode, hub_id: state.hub.id, secrets: secrets}
+  defp build_deployment_group(state, deployment_group) do
+    secrets = Enum.map(deployment_group.secrets, &build_secret(state, &1))
+
+    %DeploymentGroup{
+      id: deployment_group.id,
+      name: deployment_group.name,
+      mode: deployment_group.mode,
+      hub_id: state.hub.id,
+      secrets: secrets,
+      clustering: deployment_group.clustering,
+      zta_provider: String.to_atom(deployment_group.zta_provider),
+      zta_key: deployment_group.zta_key
+    }
   end
 
   defp handle_event(:secret_created, %Secrets.Secret{} = secret, state) do

lib/livebook/teams/deployment_group.ex

@@ -3,11 +3,17 @@ defmodule Livebook.Teams.DeploymentGroup do
   import Ecto.Changeset
   alias Livebook.Secrets.Secret
 
+  # If this list is updated, it must also be mirrored on Livebook Teams Server.
+  @zta_providers ~w(cloudflare google_iap tailscale teleport)a
+
   @type t :: %__MODULE__{
           id: String.t() | nil,
           name: String.t() | nil,
           mode: :online | :offline,
           hub_id: String.t() | nil,
+          clustering: :fly_io | nil,
+          zta_provider: :cloudflare | :google_iap | :tailscale | :teleport,
+          zta_key: String.t(),
           secrets: [Secret.t()]
         }
 
@@ -16,12 +22,15 @@ defmodule Livebook.Teams.DeploymentGroup do
     field :name, :string
     field :mode, Ecto.Enum, values: [:online, :offline]
     field :hub_id, :string
+    field :clustering, Ecto.Enum, values: [:fly_io]
+    field :zta_provider, Ecto.Enum, values: @zta_providers
+    field :zta_key, :string
     has_many :secrets, Secret
   end
 
   def changeset(deployment_group, attrs \\ %{}) do
     deployment_group
-    |> cast(attrs, [:id, :name, :mode, :hub_id])
+    |> cast(attrs, [:id, :name, :mode, :hub_id, :clustering, :zta_provider, :zta_key])
     |> validate_required([:name, :mode])
   end
 end

lib/livebook/teams/requests.ex

@@ -169,7 +169,14 @@ defmodule Livebook.Teams.Requests do
   @spec create_deployment_group(Team.t(), DeploymentGroup.t()) ::
           {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
   def create_deployment_group(team, deployment_group) do
-    params = %{name: deployment_group.name, mode: deployment_group.mode}
+    params = %{
+      name: deployment_group.name,
+      mode: deployment_group.mode,
+      clustering: deployment_group.clustering,
+      zta_provider: deployment_group.zta_provider,
+      zta_key: deployment_group.zta_key
+    }
+
     post("/api/v1/org/deployment-groups", params, team)
   end
 
@@ -179,7 +186,15 @@ defmodule Livebook.Teams.Requests do
   @spec update_deployment_group(Team.t(), DeploymentGroup.t()) ::
           {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
   def update_deployment_group(team, deployment_group) do
-    params = %{id: deployment_group.id, name: deployment_group.name, mode: deployment_group.mode}
+    params = %{
+      id: deployment_group.id,
+      name: deployment_group.name,
+      mode: deployment_group.mode,
+      clustering: deployment_group.clustering,
+      zta_provider: deployment_group.zta_provider,
+      zta_key: deployment_group.zta_key
+    }
+
     put("/api/v1/org/deployment-groups", params, team)
   end
 

lib/livebook_web/components/form_components.ex

@@ -411,6 +411,7 @@ defmodule LivebookWeb.FormComponents do
   attr :class, :string, default: ""
   attr :field, Phoenix.HTML.FormField, doc: "a form field struct retrieved from the form"
   attr :help, :string, default: nil
+  attr :disabled, :boolean, default: false
 
   attr :options, :list, default: []
   attr :prompt, :string, default: nil

lib/livebook_web/live/app_helpers.ex

@@ -102,68 +102,85 @@ defmodule LivebookWeb.AppHelpers do
         ]}
       />
       <.radio_field label="Base image" field={@form[:docker_tag]} options={docker_tag_options()} />
-      <div class="grid grid-cols-1 md:grid-cols-2">
-        <.select_field
-          label="Clustering"
-          help={
-            ~S'''
-            When running multiple
-            instances of Livebook,
-            they need to be connected
-            into a single cluster.
-            You must either deploy
-            it as a single instance
-            or choose a platform to
-            enable clustering on.
-            '''
-          }
-          field={@form[:clustering]}
-          options={[
-            {"Single instance", ""},
-            {"Fly.io", "fly_io"}
-          ]}
-        />
-      </div>
-      <%= if Hubs.Provider.type(@hub) == "team" do %>
-        <div class="flex flex-col">
-          <div class="grid grid-cols-1 md:grid-cols-2 gap-3">
-            <.select_field
-              label="Zero Trust Authentication provider"
-              field={@form[:zta_provider]}
-              help={
-                ~S'''
-                Enable this option if you want
-                to deploy your notebooks behind
-                an authentication proxy
-                '''
-              }
-              prompt="None"
-              options={zta_options()}
-            />
-            <.text_field
-              :if={zta_metadata = zta_metadata(@form[:zta_provider].value)}
-              field={@form[:zta_key]}
-              label={zta_metadata.value}
-              help={zta_help(zta_metadata)}
-              phx-debounce
-            />
-          </div>
-          <div :if={zta_metadata = zta_metadata(@form[:zta_provider].value)} class="text-sm mt-1">
-            See the
-            <a
-              class="text-blue-800 hover:text-blue-600"
-              href={"https://hexdocs.pm/livebook/#{zta_metadata.type}"}
-            >
-              Authentication with <%= zta_metadata.name %> docs
-            </a>
-            for more information.
-          </div>
-        </div>
-      <% end %>
     </div>
     """
   end
 
+  @doc """
+  Renders form fields for Deployment Group.
+  """
+  attr :form, Phoenix.HTML.Form, required: true
+  attr :hub, :map, required: true
+
+  def deployment_group_form_content(assigns) do
+    ~H"""
+    <div class="grid grid-cols-1 md:grid-cols-2">
+      <.select_field
+        label="Clustering"
+        help={
+          ~S'''
+          When running multiple
+          instances of Livebook,
+          they need to be connected
+          into a single cluster.
+          You must either deploy
+          it as a single instance
+          or choose a platform to
+          enable clustering on.
+          '''
+        }
+        field={@form[:clustering]}
+        options={[
+          {"Single instance", ""},
+          {"Fly.io", "fly_io"}
+        ]}
+        disabled={@form[:ready_only].value}
+        class="disabled:cursor-not-allowed"
+      />
+    </div>
+    <%= if Hubs.Provider.type(@hub) == "team" do %>
+      <div class="flex flex-col">
+        <div class="grid grid-cols-1 md:grid-cols-2 gap-3">
+          <.select_field
+            label="Zero Trust Authentication provider"
+            field={@form[:zta_provider]}
+            help={
+              ~S'''
+              Enable this option if you want
+              to deploy your notebooks behind
+              an authentication proxy
+              '''
+            }
+            prompt="None"
+            options={zta_options()}
+            disabled={@form[:ready_only].value}
+            class="disabled:cursor-not-allowed"
+          />
+          <.text_field
+            :if={zta_metadata = zta_metadata(@form[:zta_provider].value)}
+            field={@form[:zta_key]}
+            label={zta_metadata.value}
+            help={zta_help(zta_metadata)}
+            phx-debounce
+            disabled={@form[:ready_only].value}
+            class="disabled:cursor-not-allowed"
+          />
+        </div>
+        <div :if={zta_metadata = zta_metadata(@form[:zta_provider].value)} class="text-sm mt-1">
+          See the
+          <a
+            class="text-blue-800 hover:text-blue-600"
+            href={"https://hexdocs.pm/livebook/#{zta_metadata.type}"}
+          >
+            Authentication with <%= zta_metadata.name %> docs
+          </a>
+          for more information.
+        </div>
+      </div>
+    <% end %>
+    """
+  end
+
   @zta_options for provider <- Livebook.Config.identity_providers(),
                    do: {provider.name, provider.type}
 

lib/livebook_web/live/hub/teams/deployment_group_form_component.ex

@@ -70,6 +70,7 @@ defmodule LivebookWeb.Hub.Teams.DeploymentGroupFormComponent do
                 {"Online", "online"}
               ]}
             />
+            <LivebookWeb.AppHelpers.deployment_group_form_content hub={@hub} form={f} />
             <div class="flex space-x-2">
               <button class="button-base button-blue" type="submit" disabled={not @changeset.valid?}>
                 <.remix_icon icon={@button.icon} class="align-middle mr-1" />

lib/livebook_web/live/session_live/app_docker_component.ex

@@ -20,7 +20,8 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
        hub_secrets: Hubs.get_secrets(assigns.hub),
        hub_file_systems: Hubs.get_file_systems(assigns.hub, hub_only: true),
        deployment_groups: deployment_groups,
-       deployment_group_form: %{"id" => assigns.deployment_group_id}
+       deployment_group_form: %{"deployment_group_id" => assigns.deployment_group_id},
+       deployment_group_id: assigns.deployment_group_id
      )
      |> assign_new(:changeset, fn -> Hubs.Dockerfile.config_changeset() end)
      |> assign_new(:save_result, fn -> nil end)
@@ -40,6 +41,7 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
         hub={@hub}
         deployment_groups={@deployment_groups}
         deployment_group_form={@deployment_group_form}
+        deployment_group_id={@deployment_group_id}
         changeset={@changeset}
         session={@session}
         dockerfile={@dockerfile}
@@ -97,7 +99,6 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
         <%= if @deployment_groups do %>
           <%= if @deployment_groups != [] do %>
             <.form
-              :let={f}
               for={@deployment_group_form}
               phx-change="select_deployment_group"
               phx-target={@myself}
@@ -109,9 +110,11 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
                   Share deployment credentials, secrets, and configuration with deployment groups.
                   '''
                 }
-                field={f[:id]}
+                field={@deployment_group_form[:deployment_group_id]}
                 options={deployment_group_options(@deployment_groups)}
                 label="Deployment Group"
+                name="deployment_group_id"
+                value={@deployment_group_id}
               />
             </.form>
           <% else %>
@@ -127,6 +130,17 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
           <%= raw(warning) %>
         </.message_box>
       </div>
+      <.form
+        :let={f}
+        for={deployment_group_form_content(assigns)}
+        as={:data}
+        phx-change="validate"
+        phx-target={@myself}
+      >
+        <div class="flex flex-col space-y-4">
+          <AppHelpers.deployment_group_form_content hub={@hub} form={f} />
+        </div>
+      </.form>
       <.form :let={f} for={@changeset} as={:data} phx-change="validate" phx-target={@myself}>
         <AppHelpers.docker_config_form_content hub={@hub} form={f} />
       </.form>
@@ -191,7 +205,7 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
     {:noreply, assign(socket, save_result: save_result)}
   end
 
-  def handle_event("select_deployment_group", %{"id" => id}, socket) do
+  def handle_event("select_deployment_group", %{"deployment_group_id" => id}, socket) do
     Livebook.Session.set_notebook_deployment_group(socket.assigns.session.pid, id)
 
     {:noreply, socket}
@@ -255,4 +269,20 @@ defmodule LivebookWeb.SessionLive.AppDockerComponent do
     for deployment_group <- [%{name: "none", id: nil}] ++ deployment_groups,
         do: {deployment_group.name, deployment_group.id}
   end
+
+  defp deployment_group_form_content(%{deployment_group_id: id} = assigns) do
+    deployment_group = if id, do: Enum.find(assigns.deployment_groups, &(&1.id == id))
+
+    if deployment_group do
+      %{
+        "deployment_group_id" => deployment_group.id,
+        "clustering" => deployment_group.clustering,
+        "zta_provider" => deployment_group.zta_provider,
+        "zta_key" => deployment_group.zta_key,
+        "ready_only" => true
+      }
+    else
+      assigns.changeset
+    end
+  end
 end

proto/lib/livebook_proto/deployment_group.pb.ex

@@ -5,4 +5,7 @@ defmodule LivebookProto.DeploymentGroup do
   field :name, 2, type: :string
   field :mode, 3, type: :string
   field :secrets, 4, repeated: true, type: LivebookProto.DeploymentGroupSecret
+  field :clustering, 5, type: :string
+  field :zta_provider, 6, type: :string
+  field :zta_key, 7, type: :string
 end

proto/lib/livebook_proto/deployment_group_created.pb.ex

@@ -5,4 +5,7 @@ defmodule LivebookProto.DeploymentGroupCreated do
   field :name, 2, type: :string
   field :mode, 3, type: :string
   field :secrets, 4, repeated: true, type: LivebookProto.DeploymentGroupSecret
+  field :clustering, 5, type: :string
+  field :zta_provider, 6, type: :string
+  field :zta_key, 7, type: :string
 end

proto/lib/livebook_proto/deployment_group_updated.pb.ex

@@ -5,4 +5,7 @@ defmodule LivebookProto.DeploymentGroupUpdated do
   field :name, 2, type: :string
   field :mode, 3, type: :string
   field :secrets, 4, repeated: true, type: LivebookProto.DeploymentGroupSecret
+  field :clustering, 5, type: :string
+  field :zta_provider, 6, type: :string
+  field :zta_key, 7, type: :string
 end