Deploy notebooks from CLI (#3034)

Author: aleDsz

lib/livebook.ex

@@ -157,6 +157,10 @@ defmodule Livebook do
       config :livebook, teams_url: url, warn_on_live_teams_server: false
     end
 
+    if System.get_env("LIVEBOOK_TEAMS_AUTH") do
+      config :livebook, :persist_storage, false
+    end
+
     if Livebook.Config.boolean!("LIVEBOOK_SHUTDOWN_ENABLED", false) do
       config :livebook, :shutdown_callback, {System, :stop, []}
     end

lib/livebook/hubs.ex

@@ -311,6 +311,7 @@ defmodule Livebook.Hubs do
   @spec get_app_specs() :: list(Livebook.Apps.AppSpec.t())
   def get_app_specs() do
     for hub <- get_hubs(),
+        Provider.connection_spec(hub),
         app_spec <- Provider.get_app_specs(hub),
         do: app_spec
   end

lib/livebook/hubs/team.ex

@@ -94,22 +94,17 @@ defmodule Livebook.Hubs.Team do
       changeset
     end
   end
-
-  @doc """
-  Returns the public key prefix
-  """
-  @spec public_key_prefix() :: String.t()
-  def public_key_prefix(), do: "lb_opk_"
 end
 
 defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   alias Livebook.Hubs.{Team, TeamClient}
-  alias Livebook.Teams.Requests
+  alias Livebook.Teams
   alias Livebook.FileSystem
   alias Livebook.Secrets.Secret
 
-  @teams_key_prefix Livebook.Teams.Org.teams_key_prefix()
-  @public_key_prefix Team.public_key_prefix()
+  @teams_key_prefix Teams.Constants.teams_key_prefix()
+  @public_key_prefix Teams.Constants.public_key_prefix()
+  @deploy_key_prefix Teams.Constants.deploy_key_prefix()
 
   def load(team, fields) do
     {offline?, fields} = Map.pop(fields, :offline?, false)
@@ -137,6 +132,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
 
   def type(_team), do: "team"
 
+  def connection_spec(%{session_token: @deploy_key_prefix <> _}), do: nil
   def connection_spec(team), do: {TeamClient, team}
 
   def disconnect(team), do: TeamClient.stop(team.id)
@@ -160,7 +156,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
     @teams_key_prefix <> teams_key = team.teams_key
     token = Livebook.Stamping.chapoly_encrypt(metadata, notebook_source, teams_key)
 
-    case Requests.org_sign(team, token) do
+    case Teams.Requests.org_sign(team, token) do
       {:ok, %{"signature" => token_signature}} ->
         stamp = %{"version" => 1, "token" => token, "token_signature" => token_signature}
         {:ok, stamp}
@@ -201,23 +197,23 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   def get_secrets(team), do: TeamClient.get_secrets(team.id)
 
   def create_secret(%Team{} = team, %Secret{} = secret) do
-    case Requests.create_secret(team, secret) do
+    case Teams.Requests.create_secret(team, secret) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_secret_errors(errors)}
       any -> any
     end
   end
 
   def update_secret(%Team{} = team, %Secret{} = secret) do
-    case Requests.update_secret(team, secret) do
+    case Teams.Requests.update_secret(team, secret) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_secret_errors(errors)}
       any -> any
     end
   end
 
   def delete_secret(%Team{} = team, %Secret{} = secret) do
-    case Requests.delete_secret(team, secret) do
+    case Teams.Requests.delete_secret(team, secret) do
       {:ok, _} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_secret_errors(errors)}
       any -> any
@@ -227,23 +223,23 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   def get_file_systems(team), do: TeamClient.get_file_systems(team.id)
 
   def create_file_system(%Team{} = team, file_system) do
-    case Requests.create_file_system(team, file_system) do
+    case Teams.Requests.create_file_system(team, file_system) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_file_system_errors(file_system, errors)}
       any -> any
     end
   end
 
   def update_file_system(%Team{} = team, file_system) do
-    case Requests.update_file_system(team, file_system) do
+    case Teams.Requests.update_file_system(team, file_system) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_file_system_errors(file_system, errors)}
       any -> any
     end
   end
 
   def delete_file_system(%Team{} = team, file_system) do
-    case Requests.delete_file_system(team, file_system) do
+    case Teams.Requests.delete_file_system(team, file_system) do
       {:ok, _} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_file_system_errors(file_system, errors)}
       any -> any
@@ -264,12 +260,12 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   defp parse_secret_errors(errors_map) do
-    Requests.to_error_list(Secret, errors_map)
+    Teams.Requests.to_error_list(Secret, errors_map)
   end
 
   defp parse_file_system_errors(%struct{} = file_system, errors_map) do
     %{error_field: field} = FileSystem.external_metadata(file_system)
     errors_map = Map.new(errors_map, fn {_key, values} -> {field, values} end)
-    Requests.to_error_list(struct, errors_map)
+    Teams.Requests.to_error_list(struct, errors_map)
   end
 end

lib/livebook/storage.ex

@@ -186,10 +186,11 @@ defmodule Livebook.Storage do
     # in case it is persisting to disk. terminate/2 is still a no-op.
     Process.flag(:trap_exit, true)
 
-    table = load_or_create_table()
+    persist_storage? = Application.get_env(:livebook, :persist_storage, true)
+    table = load_or_create_table(persist_storage?)
     :persistent_term.put(__MODULE__, table)
 
-    {:ok, %{table: table}}
+    {:ok, %{table: table, persist?: persist_storage?}}
   end
 
   @impl true
@@ -220,6 +221,10 @@ defmodule Livebook.Storage do
   end
 
   @impl true
+  def handle_continue(:save_to_file, %{persist?: false} = state) do
+    {:noreply, state}
+  end
+
   def handle_continue(:save_to_file, %{table: table} = state) do
     file_path = String.to_charlist(config_file_path())
     :ok = :ets.tab2file(table, file_path)
@@ -228,7 +233,11 @@ defmodule Livebook.Storage do
 
   defp table_name(), do: :persistent_term.get(__MODULE__)
 
-  defp load_or_create_table() do
+  defp load_or_create_table(false) do
+    :ets.new(__MODULE__, [:protected, :duplicate_bag])
+  end
+
+  defp load_or_create_table(true) do
     tab =
       if path = config_file_path_for_restore() do
         path

lib/livebook/teams.ex

@@ -11,7 +11,7 @@ defmodule Livebook.Teams do
   import Ecto.Changeset,
     only: [add_error: 3, apply_action: 2, apply_action!: 2, get_field: 2]
 
-  @prefix Org.teams_key_prefix()
+  @teams_key_prefix Teams.Constants.teams_key_prefix()
 
   @doc """
   Creates an Org.
@@ -148,7 +148,7 @@ defmodule Livebook.Teams do
   Derives the secret and sign secret from given `teams_key`.
   """
   @spec derive_key(String.t()) :: bitstring()
-  def derive_key(@prefix <> teams_key) do
+  def derive_key(@teams_key_prefix <> teams_key) do
     binary_key = Base.url_decode64!(teams_key, padding: false)
     Plug.Crypto.KeyGenerator.generate(binary_key, "notebook secret", cache: Plug.Crypto.Keys)
   end
@@ -235,6 +235,45 @@ defmodule Livebook.Teams do
     TeamClient.get_environment_variables(team.id)
   end
 
+  @doc """
+  Fetches the CLI session using a deploy key.
+  """
+  @spec fetch_cli_session(map()) ::
+          {:ok, Team.t()} | {:error, String.t()} | {:transport_error, String.t()}
+  def fetch_cli_session(%{session_token: _, teams_key: _} = config) do
+    with {:ok, %{"name" => name} = attrs} <- Requests.fetch_cli_session(config) do
+      id = "team-#{name}"
+
+      hub =
+        Hubs.save_hub(%Team{
+          id: id,
+          hub_name: name,
+          hub_emoji: "🚀",
+          user_id: nil,
+          org_id: attrs["org_id"],
+          org_key_id: attrs["org_key_id"],
+          session_token: config.session_token,
+          teams_key: config.teams_key,
+          org_public_key: attrs["public_key"]
+        })
+
+      {:ok, hub}
+    end
+  end
+
+  @doc """
+  Deploys the given app deployment to given deployment group using a deploy key.
+  """
+  @spec deploy_app_from_cli(Team.t(), Teams.AppDeployment.t(), String.t()) ::
+          {:ok, String.t()} | {:error, map()} | {:transport_error, String.t()}
+  def deploy_app_from_cli(%Team{} = team, %Teams.AppDeployment{} = app_deployment, name) do
+    case Requests.deploy_app_from_cli(team, app_deployment, name) do
+      {:ok, %{"url" => url}} -> {:ok, url}
+      {:error, %{"errors" => errors}} -> {:error, errors}
+      any -> any
+    end
+  end
+
   defp map_teams_field_to_livebook_field(map, teams_field, livebook_field) do
     if value = map[teams_field] do
       Map.put_new(map, livebook_field, value)

lib/livebook/teams/constants.ex

@@ -0,0 +1,25 @@
+defmodule Livebook.Teams.Constants do
+  @doc """
+  Returns the public key prefix
+  """
+  @spec public_key_prefix() :: String.t()
+  def public_key_prefix(), do: "lb_opk_"
+
+  @doc """
+  Returns the Agent Key prefix
+  """
+  @spec agent_key_prefix() :: String.t()
+  def agent_key_prefix, do: "lb_ak_"
+
+  @doc """
+  Returns the Deploy Key prefix
+  """
+  @spec deploy_key_prefix() :: String.t()
+  def deploy_key_prefix, do: "lb_dk_"
+
+  @doc """
+  Returns the Teams Key prefix
+  """
+  @spec teams_key_prefix() :: String.t()
+  def teams_key_prefix(), do: "lb_tk_"
+end

lib/livebook/teams/org.ex

@@ -2,8 +2,6 @@ defmodule Livebook.Teams.Org do
   use Ecto.Schema
   import Ecto.Changeset
 
-  @prefix "lb_tk_"
-
   @type t :: %__MODULE__{
           id: pos_integer() | nil,
           emoji: String.t() | nil,
@@ -31,7 +29,7 @@ defmodule Livebook.Teams.Org do
   @spec teams_key() :: String.t()
   def teams_key() do
     key = :crypto.strong_rand_bytes(@secret_key_size)
-    @prefix <> Base.url_encode64(key, padding: false)
+    Livebook.Teams.Constants.teams_key_prefix() <> Base.url_encode64(key, padding: false)
   end
 
   @doc """
@@ -50,10 +48,4 @@ defmodule Livebook.Teams.Org do
       message: "should only contain lowercase alphanumeric characters and dashes"
     )
   end
-
-  @doc """
-  Returns the teams key prefix
-  """
-  @spec teams_key_prefix() :: String.t()
-  def teams_key_prefix(), do: @prefix
 end

lib/livebook/teams/requests.ex

@@ -5,6 +5,7 @@ defmodule Livebook.Teams.Requests do
   alias Livebook.Secrets.Secret
   alias Livebook.Teams
 
+  @deploy_key_prefix Teams.Constants.deploy_key_prefix()
   @error_message "Something went wrong, try again later or please file a bug if it persists"
   @unauthorized_error_message "You are not authorized to perform this action, make sure you have the access and you are not in a Livebook App Server/Offline instance"
 
@@ -227,6 +228,34 @@ defmodule Livebook.Teams.Requests do
     get("/api/v1/org/identity", %{access_token: access_token}, team)
   end
 
+  @doc """
+  Send a request to Livebook Team API to return a session using a deploy key.
+  """
+  @spec fetch_cli_session(map()) :: api_result()
+  def fetch_cli_session(config) do
+    post("/api/v1/cli/auth", %{}, config)
+  end
+
+  @doc """
+  Send a request to Livebook Team API to deploy an app using a deploy key.
+  """
+  @spec deploy_app_from_cli(Team.t(), Teams.AppDeployment.t(), String.t()) :: api_result()
+  def deploy_app_from_cli(team, app_deployment, deployment_group_name) do
+    secret_key = Teams.derive_key(team.teams_key)
+
+    params = %{
+      title: app_deployment.title,
+      slug: app_deployment.slug,
+      multi_session: app_deployment.multi_session,
+      access_type: app_deployment.access_type,
+      deployment_group_name: deployment_group_name,
+      sha: app_deployment.sha
+    }
+
+    encrypted_content = Teams.encrypt(app_deployment.file, secret_key)
+    upload("/api/v1/cli/org/apps", encrypted_content, params, team)
+  end
+
   @doc """
   Normalizes errors map into errors for the given schema.
   """
@@ -271,7 +300,6 @@ defmodule Livebook.Teams.Requests do
   defp upload(path, content, params, team) do
     build_req(team)
     |> Req.Request.put_header("content-length", "#{byte_size(content)}")
-    |> Req.Request.put_private(:deploy, true)
     |> Req.post(url: path, params: params, body: content)
     |> handle_response()
     |> dispatch_messages(team)
@@ -280,7 +308,6 @@ defmodule Livebook.Teams.Requests do
   defp build_req(team) do
     Req.new(base_url: Livebook.Config.teams_url())
     |> Req.Request.put_new_header("x-lb-version", Livebook.Config.app_version())
-    |> Req.Request.append_response_steps(transform_teams_response: &transform_response/1)
     |> Livebook.Utils.req_attach_defaults()
     |> add_team_auth(team)
   end
@@ -291,6 +318,11 @@ defmodule Livebook.Teams.Requests do
     Req.Request.append_request_steps(req, unauthorized: &{&1, Req.Response.new(status: 401)})
   end
 
+  defp add_team_auth(req, %{session_token: @deploy_key_prefix <> _} = team) do
+    token = "#{team.session_token}:#{Teams.Org.key_hash(%Teams.Org{teams_key: team.teams_key})}"
+    Req.Request.merge_options(req, auth: {:bearer, token})
+  end
+
   defp add_team_auth(req, %{user_id: nil} = team) do
     agent_name = Livebook.Config.agent_name()
     token = "#{team.session_token}:#{agent_name}:#{team.org_id}:#{team.org_key_id}"
@@ -303,17 +335,6 @@ defmodule Livebook.Teams.Requests do
     Req.Request.merge_options(req, auth: {:bearer, token})
   end
 
-  defp transform_response({request, response}) do
-    case {request, response} do
-      {request, %{status: 400, body: %{"errors" => %{"detail" => error}}}}
-      when request.private.deploy ->
-        {request, %{response | status: 422, body: %{"errors" => %{"file" => [error]}}}}
-
-      _otherwise ->
-        {request, response}
-    end
-  end
-
   defp handle_response(response) do
     case response do
       {:ok, %{status: status} = response} when status in 200..299 -> {:ok, response.body}