Deploy apps to Livebook Agent (#2511)

lib/livebook/hubs/team_client.ex

@@ -279,6 +279,21 @@ defmodule Livebook.Hubs.TeamClient do
     }
   end
 
+  defp put_app_deployment(deployment_group, app_deployment) do
+    deployment_group = remove_app_deployment(deployment_group, app_deployment)
+    app_deployments = [app_deployment | deployment_group.app_deployments]
+
+    %{deployment_group | app_deployments: Enum.sort_by(app_deployments, & &1.slug)}
+  end
+
+  defp remove_app_deployment(deployment_group, app_deployment) do
+    %{
+      deployment_group
+      | app_deployments:
+          Enum.reject(deployment_group.app_deployments, &(&1.slug == app_deployment.slug))
+    }
+  end
+
   defp build_agent_key(agent_key) do
     %Teams.AgentKey{
       id: agent_key.id,
@@ -290,6 +305,7 @@ defmodule Livebook.Hubs.TeamClient do
   defp build_deployment_group(state, deployment_group) do
     secrets = Enum.map(deployment_group.secrets, &build_secret(state, &1))
     agent_keys = Enum.map(deployment_group.agent_keys, &build_agent_key/1)
+    app_deployments = Enum.map(deployment_group.deployed_apps, &build_app_deployment/1)
 
     %Teams.DeploymentGroup{
       id: deployment_group.id,
@@ -298,12 +314,29 @@ defmodule Livebook.Hubs.TeamClient do
       hub_id: state.hub.id,
       secrets: secrets,
       agent_keys: agent_keys,
+      app_deployments: app_deployments,
       clustering: nullify(deployment_group.clustering),
       zta_provider: atomize(deployment_group.zta_provider),
       zta_key: nullify(deployment_group.zta_key)
     }
   end
 
+  defp build_app_deployment(app_deployment) do
+    path = URI.parse(app_deployment.archive_url).path
+
+    %Teams.AppDeployment{
+      id: app_deployment.id,
+      filename: Path.basename(path),
+      slug: app_deployment.slug,
+      sha: app_deployment.sha,
+      title: app_deployment.title,
+      deployment_group_id: app_deployment.deployment_group_id,
+      file: {:url, app_deployment.archive_url},
+      deployed_by: app_deployment.deployed_by,
+      deployed_at: NaiveDateTime.from_iso8601!(app_deployment.deployed_at)
+    }
+  end
+
   defp handle_event(:secret_created, %Secrets.Secret{} = secret, state) do
     Hubs.Broadcasts.secret_created(secret)
 
@@ -382,6 +415,7 @@ defmodule Livebook.Hubs.TeamClient do
   end
 
   defp handle_event(:deployment_group_updated, %Teams.DeploymentGroup{} = deployment_group, state) do
+    deployment_group = deploy_apps(state.deployment_group_id, deployment_group, state.derived_key)
     Teams.Broadcasts.deployment_group_updated(deployment_group)
 
     put_deployment_group(state, deployment_group)
@@ -397,6 +431,7 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp handle_event(:deployment_group_deleted, deployment_group_deleted, state) do
     with {:ok, deployment_group} <- fetch_deployment_group(deployment_group_deleted.id, state) do
+      undeploy_apps(state.deployment_group_id, deployment_group)
       Teams.Broadcasts.deployment_group_deleted(deployment_group)
 
       remove_deployment_group(state, deployment_group)
@@ -436,6 +471,24 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
+  defp handle_event(:app_deployment_created, app_deployment_created, state) do
+    app_deployment = build_app_deployment(app_deployment_created)
+    deployment_group_id = app_deployment.deployment_group_id
+
+    with {:ok, deployment_group} <- fetch_deployment_group(deployment_group_id, state) do
+      app_deployment =
+        if deployment_group_id == state.deployment_group_id do
+          {:ok, app_deployment} = download_and_deploy(app_deployment, state.derived_key)
+          app_deployment
+        else
+          app_deployment
+        end
+
+      deployment_group = put_app_deployment(deployment_group, app_deployment)
+      handle_event(:deployment_group_updated, deployment_group, state)
+    end
+  end
+
   defp dispatch_secrets(state, %{secrets: secrets}) do
     decrypted_secrets = Enum.map(secrets, &build_secret(state, &1))
 
@@ -534,4 +587,59 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp nullify(""), do: nil
   defp nullify(value), do: value
+
+  defp deploy_apps(id, %{id: id} = deployment_group, derived_key) do
+    app_deployments =
+      for app_deployment <- deployment_group.app_deployments do
+        {:ok, app_deployment} = download_and_deploy(app_deployment, derived_key)
+        app_deployment
+      end
+
+    %{deployment_group | app_deployments: app_deployments}
+  end
+
+  defp deploy_apps(_, deployment_group, _), do: deployment_group
+
+  defp undeploy_apps(id, %{id: id} = deployment_group) do
+    for %{slug: slug} <- deployment_group.app_deployments do
+      :ok = undeploy_app(slug)
+    end
+  end
+
+  defp undeploy_apps(_, _), do: :noop
+
+  defp download_and_deploy(%{file: nil} = app_deployment, _) do
+    app_deployment
+  end
+
+  defp download_and_deploy(%{file: {:url, archive_url}} = app_deployment, derived_key) do
+    destination_path = app_deployment_path(app_deployment.slug)
+
+    with {:ok, %{status: 200} = response} <- Req.get(archive_url),
+         :ok <- undeploy_app(app_deployment.slug),
+         {:ok, decrypted_content} <- Teams.decrypt(response.body, derived_key),
+         :ok <- unzip_app(decrypted_content, destination_path),
+         :ok <- Livebook.Apps.deploy_apps_in_dir(destination_path) do
+      {:ok, %{app_deployment | file: nil}}
+    end
+  end
+
+  defp undeploy_app(slug) do
+    with {:ok, app} <- Livebook.Apps.fetch_app(slug) do
+      Livebook.App.close(app.pid)
+    end
+
+    :ok
+  end
+
+  defp app_deployment_path(slug) do
+    Path.join([Livebook.Config.tmp_path(), "apps", slug <> Livebook.Utils.random_short_id()])
+  end
+
+  defp unzip_app(content, destination_path) do
+    case :zip.extract(content, cwd: to_charlist(destination_path)) do
+      {:ok, _} -> :ok
+      {:error, error} -> FileSystem.Utils.posix_error(error)
+    end
+  end
 end

lib/livebook/teams.ex

@@ -4,7 +4,7 @@ defmodule Livebook.Teams do
   alias Livebook.Hubs
   alias Livebook.Hubs.Team
   alias Livebook.Hubs.TeamClient
-  alias Livebook.Teams.{AgentKey, DeploymentGroup, Org, Requests}
+  alias Livebook.Teams.{AgentKey, AppDeployment, DeploymentGroup, Org, Requests}
 
   import Ecto.Changeset,
     only: [add_error: 3, apply_action: 2, apply_action!: 2, get_field: 2]
@@ -249,4 +249,27 @@ defmodule Livebook.Teams do
       []
     end
   end
+
+  @doc """
+  Creates a new app deployment.
+  """
+  @spec deploy_app(Team.t(), AppDeployment.t()) ::
+          :ok
+          | {:error, Ecto.Changeset.t()}
+          | {:transport_error, String.t()}
+  def deploy_app(%Team{} = team, %AppDeployment{} = app_deployment) do
+    case Requests.deploy_app(team, app_deployment) do
+      {:ok, %{"id" => _id}} ->
+        :ok
+
+      {:error, %{"errors" => %{"detail" => error}}} ->
+        {:error, Requests.add_errors(app_deployment, %{"file" => [error]})}
+
+      {:error, %{"errors" => errors}} ->
+        {:error, Requests.add_errors(app_deployment, errors)}
+
+      any ->
+        any
+    end
+  end
 end

lib/livebook/teams/app_deployment.ex

@@ -0,0 +1,88 @@
+defmodule Livebook.Teams.AppDeployment do
+  use Ecto.Schema
+  alias Livebook.FileSystem
+
+  @file_extension ".zip"
+  @max_size 20 * 1024 * 1024
+
+  @type t :: %__MODULE__{
+          id: String.t() | nil,
+          filename: String.t() | nil,
+          slug: String.t() | nil,
+          sha: String.t() | nil,
+          title: String.t() | nil,
+          deployment_group_id: String.t() | nil,
+          file: {:url, String.t()} | {:content, binary()} | nil,
+          deployed_by: String.t() | nil,
+          deployed_at: NaiveDateTime.t() | nil
+        }
+
+  @primary_key {:id, :string, autogenerate: false}
+  embedded_schema do
+    field :filename, :string
+    field :slug, :string
+    field :sha, :string
+    field :title, :string
+    field :deployment_group_id, :string
+    field :file, :string
+    field :deployed_by, :string
+
+    timestamps(updated_at: nil, inserted_at: :deployed_at)
+  end
+
+  @doc """
+  Creates a new app deployment from notebook.
+  """
+  @spec new(Livebook.Notebook.t(), Livebook.FileSystem.File.t()) ::
+          {:ok, t()} | {:warning, list(String.t())} | {:error, FileSystem.error()}
+  def new(notebook, files_dir) do
+    with {:ok, source} <- fetch_notebook_source(notebook),
+         {:ok, files} <- build_and_check_file_entries(notebook, source, files_dir),
+         {:ok, {_, zip_content}} <- :zip.create(~c"app_deployment.zip", files, [:memory]),
+         :ok <- validate_size(zip_content) do
+      md5_hash = :crypto.hash(:md5, zip_content)
+      shasum = Base.encode16(md5_hash, case: :lower)
+
+      {:ok,
+       %__MODULE__{
+         filename: shasum <> @file_extension,
+         slug: notebook.app_settings.slug,
+         sha: shasum,
+         title: notebook.name,
+         deployment_group_id: notebook.deployment_group_id,
+         file: {:content, zip_content}
+       }}
+    end
+  end
+
+  defp fetch_notebook_source(notebook) do
+    case Livebook.LiveMarkdown.notebook_to_livemd(notebook) do
+      {source, []} -> {:ok, source}
+      {_, warnings} -> {:warning, warnings}
+    end
+  end
+
+  defp build_and_check_file_entries(notebook, source, files_dir) do
+    notebook.file_entries
+    |> Enum.filter(&(&1.type == :attachment))
+    |> Enum.reduce_while({:ok, [{~c"notebook.livemd", source}]}, fn %{name: name}, {:ok, acc} ->
+      file = Livebook.FileSystem.File.resolve(files_dir, name)
+
+      with {:ok, true} <- Livebook.FileSystem.File.exists?(file),
+           {:ok, content} <- Livebook.FileSystem.File.read(file) do
+        {:cont, {:ok, [{to_charlist("files/" <> name), content} | acc]}}
+      else
+        {:ok, false} -> {:halt, {:error, "files/#{name}: doesn't exist"}}
+        {:error, reason} -> {:halt, {:error, "files/#{name}: #{reason}"}}
+      end
+    end)
+  end
+
+  defp validate_size(data) do
+    if byte_size(data) <= @max_size do
+      :ok
+    else
+      {:error, "the notebook and its attachments have exceeded the maximum size of 20MB"}
+    end
+  end
+end

lib/livebook/teams/deployment_group.ex

@@ -3,7 +3,7 @@ defmodule Livebook.Teams.DeploymentGroup do
   import Ecto.Changeset
 
   alias Livebook.Secrets.Secret
-  alias Livebook.Teams.AgentKey
+  alias Livebook.Teams.{AgentKey, AppDeployment}
 
   # If this list is updated, it must also be mirrored on Livebook Teams Server.
   @zta_providers ~w(cloudflare google_iap tailscale teleport)a
@@ -17,7 +17,8 @@ defmodule Livebook.Teams.DeploymentGroup do
           zta_provider: :cloudflare | :google_iap | :tailscale | :teleport,
           zta_key: String.t(),
           secrets: [Secret.t()],
-          agent_keys: [AgentKey.t()]
+          agent_keys: [AgentKey.t()],
+          app_deployments: [AppDeployment.t()]
         }
 
   @primary_key {:id, :string, autogenerate: false}
@@ -31,6 +32,7 @@ defmodule Livebook.Teams.DeploymentGroup do
 
     has_many :secrets, Secret
     has_many :agent_keys, AgentKey
+    has_many :app_deployments, AppDeployment
   end
 
   def changeset(deployment_group, attrs \\ %{}) do

lib/livebook/teams/requests.ex

@@ -4,7 +4,9 @@ defmodule Livebook.Teams.Requests do
   alias Livebook.Hubs.Team
   alias Livebook.Secrets.Secret
   alias Livebook.Teams
-  alias Livebook.Teams.{AgentKey, DeploymentGroup, Org}
+  alias Livebook.Teams.{AgentKey, AppDeployment, DeploymentGroup, Org}
+
+  @error_message "Something went wrong, try again later or please file a bug if it persists"
 
   @doc """
   Send a request to Livebook Team API to create a new org.
@@ -216,6 +218,26 @@ defmodule Livebook.Teams.Requests do
     delete("/api/v1/org/deployment-groups/agent-keys", params, team)
   end
 
+  @doc """
+  Send a request to Livebook Team API to deploy an app.
+  """
+  @spec deploy_app(Team.t(), AppDeployment.t()) ::
+          {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
+  def deploy_app(team, %{file: {:content, content}} = app_deployment) do
+    secret_key = Teams.derive_key(team.teams_key)
+
+    params = %{
+      filename: app_deployment.filename <> ".encrypted",
+      title: app_deployment.title,
+      slug: app_deployment.slug,
+      deployment_group_id: app_deployment.deployment_group_id,
+      sha: app_deployment.sha
+    }
+
+    encrypted_content = Teams.encrypt(content, secret_key)
+    upload("/api/v1/org/apps", encrypted_content, params, team)
+  end
+
   @doc """
   Add requests errors to a `changeset` for the given `fields`.
   """
@@ -235,6 +257,9 @@ defmodule Livebook.Teams.Requests do
     value |> Ecto.Changeset.change() |> add_errors(struct.__schema__(:fields), errors_map)
   end
 
+  @doc false
+  def error_message(), do: @error_message
+
   defp post(path, json, team \\ nil) do
     build_req()
     |> add_team_auth(team)
@@ -261,6 +286,14 @@ defmodule Livebook.Teams.Requests do
     |> request(method: :get, url: path, params: params)
   end
 
+  defp upload(path, content, params, team) do
+    build_req()
+    |> add_team_auth(team)
+    |> Req.Request.put_header("content-length", "#{byte_size(content)}")
+    |> request(method: :post, url: path, params: params, body: content)
+    |> dispatch_messages(team)
+  end
+
   defp build_req() do
     Req.new(
       base_url: Livebook.Config.teams_url(),
@@ -301,8 +334,7 @@ defmodule Livebook.Teams.Requests do
          "You are not authorized to perform this action, make sure you have the access or you are not in a Livebook Agent instance"}
 
       _otherwise ->
-        {:transport_error,
-         "Something went wrong, try again later or please file a bug if it persists"}
+        {:transport_error, @error_message}
     end
   end