Deploy notebooks from CLI

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

@@ -110,6 +110,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
 
   @teams_key_prefix Livebook.Teams.Org.teams_key_prefix()
   @public_key_prefix Team.public_key_prefix()
+  @deploy_key_prefix Requests.deploy_key_prefix()
 
   def load(team, fields) do
     {offline?, fields} = Map.pop(fields, :offline?, false)
@@ -137,6 +138,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)

lib/livebook/migration.ex

@@ -11,6 +11,7 @@ defmodule Livebook.Migration do
   def run() do
     insert_personal_hub()
     remove_offline_hub()
+    remove_cli_hub()
 
     storage_version =
       case Storage.fetch_key(:system, "global", :migration_version) do
@@ -46,6 +47,19 @@ defmodule Livebook.Migration do
     end
   end
 
+  @deploy_key_prefix Livebook.Teams.Requests.deploy_key_prefix()
+
+  defp remove_cli_hub() do
+    # The CLI hub will only be present in the storage if the
+    # user doesn't have the Team hub already persisted with the
+    # user credentials. Consequently, we always remove it and
+    # insert on CLI if applicable.
+
+    for %{id: "team-" <> _ = id, session_token: @deploy_key_prefix <> _} <- Storage.all(:hubs) do
+      :ok = Storage.delete(:hubs, id)
+    end
+  end
+
   defp migration(1) do
     v1_add_personal_hub_secret_key()
     v1_delete_local_host_hub()

lib/livebook/teams.ex

@@ -235,6 +235,49 @@ 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 =
+        if Hubs.hub_exists?(id) do
+          %{Hubs.fetch_hub!(id) | user_id: nil, session_token: config.session_token}
+        else
+          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"]
+          })
+        end
+
+      {: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/requests.ex

@@ -5,6 +5,7 @@ defmodule Livebook.Teams.Requests do
   alias Livebook.Secrets.Secret
   alias Livebook.Teams
 
+  @deploy_key_prefix "lb_dk_"
   @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"
 
@@ -14,6 +15,9 @@ defmodule Livebook.Teams.Requests do
   @doc false
   def error_message(), do: @error_message
 
+  @doc false
+  def deploy_key_prefix(), do: @deploy_key_prefix
+
   @doc """
   Send a request to Livebook Team API to create a new org.
   """
@@ -227,6 +231,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,6 +303,7 @@ 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(:cli, path =~ "cli")
     |> Req.Request.put_private(:deploy, true)
     |> Req.post(url: path, params: params, body: content)
     |> handle_response()
@@ -291,6 +324,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}"
@@ -305,6 +343,14 @@ defmodule Livebook.Teams.Requests do
 
   defp transform_response({request, response}) do
     case {request, response} do
+      {request, %{status: 404}} when request.private.cli and request.private.deploy ->
+        {request,
+         %{
+           response
+           | status: 422,
+             body: %{"errors" => %{"deployment_group" => ["does not exist"]}}
+         }}
+
       {request, %{status: 400, body: %{"errors" => %{"detail" => error}}}}
       when request.private.deploy ->
         {request, %{response | status: 422, body: %{"errors" => %{"file" => [error]}}}}

lib/livebook_cli.ex

@@ -1,19 +1,23 @@
 defmodule LivebookCLI do
-  def usage() do
-    """
+  alias LivebookCLI.{Task, Utils}
+
+  @help_args ["--help", "-h"]
+  @version_args ["--version", "-v"]
+
+  def usage,
+    do: """
     Usage: livebook [command] [options]
 
     Available commands:
 
       livebook server    Starts the Livebook web application
+      livebook deploy    Deploys a notebook to Livebook Teams
 
-    The --help and --version options can be given instead of a command for usage and versioning information.
+    The --help and --version options can be given instead of a command for usage and versioning information.\
     """
-  end
 
   def main(args) do
     {:ok, _} = Application.ensure_all_started(:elixir)
-
     extract_priv!()
 
     :ok = Application.load(:livebook)
@@ -22,62 +26,32 @@ defmodule LivebookCLI do
       Application.put_env(:elixir, :ansi_enabled, true)
     end
 
-    call(args)
-  end
-
-  defp unix?(), do: match?({:unix, _}, :os.type())
-
-  defp call([arg]) when arg in ["--help", "-h"], do: display_help()
-  defp call([arg]) when arg in ["--version", "-v"], do: display_version()
-
-  defp call([task_name | args]) do
-    case find_task(task_name) do
-      nil ->
-        IO.ANSI.format([:red, "Unknown command #{task_name}\n"]) |> IO.puts()
-        IO.write(usage())
-
-      task ->
-        call_task(task, args)
+    case args do
+      [arg] when arg in @help_args -> display_help()
+      [arg] when arg in @version_args -> display_version()
+      [name | [arg]] when arg in @help_args -> Task.usage(name)
+      [name | args] -> Task.call(name, List.delete(args, name))
+      _args -> Utils.print_text(usage())
     end
   end
 
-  defp call(_args), do: IO.write(usage())
-
-  defp find_task("server"), do: LivebookCLI.Server
-  defp find_task(_), do: nil
-
-  defp call_task(task, [arg]) when arg in ["--help", "-h"] do
-    IO.write(task.usage())
-  end
-
-  defp call_task(task, args) do
-    try do
-      task.call(args)
-    rescue
-      error in OptionParser.ParseError ->
-        IO.ANSI.format([
-          :red,
-          Exception.message(error),
-          "\n\nFor more information try --help"
-        ])
-        |> IO.puts()
-
-      error ->
-        IO.ANSI.format([:red, Exception.format(:error, error, __STACKTRACE__), "\n"]) |> IO.puts()
-    end
-  end
+  defp unix?(), do: match?({:unix, _}, :os.type())
 
   defp display_help() do
-    IO.puts("Livebook is an interactive notebook system for Elixir\n")
-    IO.write(usage())
+    Utils.print_text("""
+    Livebook is an interactive notebook system for Elixir
+
+    #{usage()}\
+    """)
   end
 
   defp display_version() do
-    IO.puts(:erlang.system_info(:system_version))
-    IO.puts("Elixir " <> System.build_info()[:build])
+    Utils.print_text("""
+    #{:erlang.system_info(:system_version)}
+    Elixir #{System.build_info()[:build]}
 
-    version = Livebook.Config.app_version()
-    IO.puts("\nLivebook #{version}")
+    Livebook #{Livebook.Config.app_version()}\
+    """)
   end
 
   import Record
@@ -104,14 +78,10 @@ defmodule LivebookCLI do
         List.starts_with?(name, in_archive_priv_path)
       end
 
-      case :zip.extract(archive, cwd: String.to_charlist(archive_dir), file_filter: file_filter) do
-        {:ok, _} ->
-          :ok
+      opts = [cwd: String.to_charlist(archive_dir), file_filter: file_filter]
 
-        {:error, error} ->
-          print_error_and_exit(
-            "Livebook failed to extract archive files, reason: #{inspect(error)}"
-          )
+      with {:error, error} <- :zip.extract(archive, opts) do
+        raise "Livebook failed to extract archive files, reason: #{inspect(error)}"
       end
 
       File.touch!(extracted_path)
@@ -120,10 +90,4 @@ defmodule LivebookCLI do
     priv_dir = Path.join(archive_dir, in_archive_priv_path)
     Application.put_env(:livebook, :priv_dir, priv_dir, persistent: true)
   end
-
-  @spec print_error_and_exit(String.t()) :: no_return()
-  defp print_error_and_exit(message) do
-    IO.ANSI.format([:red, message]) |> IO.puts()
-    System.halt(1)
-  end
 end