Add download for notebook files (#2112)

lib/livebook/session.ex

@@ -684,6 +684,18 @@ defmodule Livebook.Session do
     end
   end
 
+  @doc """
+  Looks up file entry with the given name and returns a local path
+  for accessing the file.
+
+  When a file is available remotely, it is first downloaded into a
+  cached location.
+  """
+  @spec fetch_file_entry_path(pid(), String.t()) :: {:ok, String.t()} | {:error, String.t()}
+  def fetch_file_entry_path(pid, name) do
+    GenServer.call(pid, {:fetch_file_entry_path, name}, :infinity)
+  end
+
   @doc """
   Closes one or more sessions.
 
@@ -990,6 +1002,14 @@ defmodule Livebook.Session do
     {:reply, :ok, state}
   end
 
+  def handle_call({:fetch_file_entry_path, name}, from, state) do
+    file_entry_path(state, name, fn reply ->
+      GenServer.reply(from, reply)
+    end)
+
+    {:noreply, state}
+  end
+
   @impl true
   def handle_cast({:set_notebook_attributes, client_pid, attrs}, state) do
     client_id = client_id(state, client_pid)

lib/livebook_web/controllers/session_controller.ex

@@ -39,6 +39,16 @@ defmodule LivebookWeb.SessionController do
     end
   end
 
+  def download_file(conn, %{"id" => id, "name" => name}) do
+    with {:ok, session} <- Sessions.fetch_session(id),
+         {:ok, path} <- Session.fetch_file_entry_path(session.pid, name) do
+      send_download(conn, {:file, path}, filename: name)
+    else
+      _ ->
+        send_resp(conn, 404, "Not found")
+    end
+  end
+
   def download_source(conn, %{"id" => id, "format" => format}) do
     case Sessions.fetch_session(id) do
       {:ok, session} ->

lib/livebook_web/live/home_live/session_list_component.ex

@@ -151,7 +151,7 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
           <.menu_item>
             <a
               role="menuitem"
-              href={~p"/sessions/#{session.id}/export/download/livemd?include_outputs=false"}
+              href={~p"/sessions/#{session.id}/download/export/livemd?include_outputs=false"}
               download
             >
               <.remix_icon icon="download-2-line" />

lib/livebook_web/live/session_live/export_live_markdown_component.ex

@@ -51,7 +51,7 @@ defmodule LivebookWeb.SessionLive.ExportLiveMarkdownComponent do
                 class="icon-button"
                 aria-label="download source"
                 href={
-                  ~p"/sessions/#{@session.id}/export/download/livemd?include_outputs=#{@include_outputs}"
+                  ~p"/sessions/#{@session.id}/download/export/livemd?include_outputs=#{@include_outputs}"
                 }
               >
                 <.remix_icon icon="download-2-line" class="text-lg" />

lib/livebook_web/live/session_live/files_list_component.ex

@@ -97,6 +97,12 @@ defmodule LivebookWeb.SessionLive.FilesListComponent do
                 <span>Clear cache</span>
               </button>
             </.menu_item>
+            <.menu_item>
+              <a role="menuitem" href={~p"/sessions/#{@session.id}/files/download/#{file_entry.name}"}>
+                <.remix_icon icon="download-2-line" />
+                <span>Download</span>
+              </a>
+            </.menu_item>
             <.menu_item variant={:danger}>
               <button
                 role="menuitem"

lib/livebook_web/router.ex

@@ -88,7 +88,7 @@ defmodule LivebookWeb.Router do
       live "/sessions/:id/settings/app", SessionLive, :app_settings
       live "/sessions/:id/add-file/:tab", SessionLive, :add_file_entry
       live "/sessions/:id/bin", SessionLive, :bin
-      get "/sessions/:id/export/download/:format", SessionController, :download_source
+      get "/sessions/:id/download/export/:format", SessionController, :download_source
       live "/sessions/:id/export/:tab", SessionLive, :export
       live "/sessions/:id/cell-settings/:cell_id", SessionLive, :cell_settings
       live "/sessions/:id/insert-image", SessionLive, :insert_image
@@ -96,6 +96,7 @@ defmodule LivebookWeb.Router do
       live "/sessions/:id/package-search", SessionLive, :package_search
       get "/sessions/:id/files/:name", SessionController, :show_file
       get "/sessions/:id/images/:name", SessionController, :show_image
+      get "/sessions/:id/download/files/:name", SessionController, :download_file
       live "/sessions/:id/settings/custom-view", SessionLive, :custom_view_settings
       live "/sessions/:id/*path_parts", SessionLive, :catch_all
     end

test/livebook_web/controllers/session_controller_test.exs

@@ -101,10 +101,71 @@ defmodule LivebookWeb.SessionControllerTest do
     end
   end
 
+  describe "download_file" do
+    test "returns not found when the given session does not exist", %{conn: conn} do
+      id = Livebook.Utils.random_node_aware_id()
+      conn = get(conn, ~p"/sessions/#{id}/download/files/data.csv")
+
+      assert conn.status == 404
+      assert conn.resp_body == "Not found"
+    end
+
+    test "returns not found when file entry with this name does not exist", %{conn: conn} do
+      {:ok, session} = Sessions.create_session()
+
+      conn = get(conn, ~p"/sessions/#{session.id}/download/files/data.csv")
+
+      assert conn.status == 404
+      assert conn.resp_body == "Not found"
+
+      Session.close(session.pid)
+    end
+
+    test "returns the file contents when it does exist", %{conn: conn} do
+      {:ok, session} = Sessions.create_session()
+
+      :ok =
+        FileSystem.File.resolve(session.files_dir, "data.csv") |> FileSystem.File.write("hello")
+
+      Session.add_file_entries(session.pid, [%{type: :attachment, name: "data.csv"}])
+
+      conn = get(conn, ~p"/sessions/#{session.id}/download/files/data.csv")
+
+      assert conn.status == 200
+      assert get_resp_header(conn, "content-disposition") == [~s/attachment; filename="data.csv"/]
+      assert get_resp_header(conn, "content-type") == ["text/csv"]
+      assert conn.resp_body == "hello"
+
+      Session.close(session.pid)
+    end
+
+    test "downloads remote file to cache and returns the file contents", %{conn: conn} do
+      bypass = Bypass.open()
+      url = "http://localhost:#{bypass.port}/data.csv"
+
+      Bypass.expect_once(bypass, "GET", "/data.csv", fn conn ->
+        Plug.Conn.resp(conn, 200, "hello")
+      end)
+
+      {:ok, session} = Sessions.create_session()
+
+      Session.add_file_entries(session.pid, [%{type: :url, name: "data.csv", url: url}])
+
+      conn = get(conn, ~p"/sessions/#{session.id}/download/files/data.csv")
+
+      assert conn.status == 200
+      assert get_resp_header(conn, "content-disposition") == [~s/attachment; filename="data.csv"/]
+      assert get_resp_header(conn, "content-type") == ["text/csv"]
+      assert conn.resp_body == "hello"
+
+      Session.close(session.pid)
+    end
+  end
+
   describe "download_source" do
     test "returns not found when the given session does not exist", %{conn: conn} do
       id = Livebook.Utils.random_node_aware_id()
-      conn = get(conn, ~p"/sessions/#{id}/export/download/livemd")
+      conn = get(conn, ~p"/sessions/#{id}/download/export/livemd")
 
       assert conn.status == 404
       assert conn.resp_body == "Not found"
@@ -113,7 +174,7 @@ defmodule LivebookWeb.SessionControllerTest do
     test "returns bad request when given an invalid format", %{conn: conn} do
       {:ok, session} = Sessions.create_session()
 
-      conn = get(conn, ~p"/sessions/#{session.id}/export/download/invalid")
+      conn = get(conn, ~p"/sessions/#{session.id}/download/export/invalid")
 
       assert conn.status == 400
       assert conn.resp_body == "Invalid format, supported formats: livemd, exs"
@@ -124,9 +185,13 @@ defmodule LivebookWeb.SessionControllerTest do
     test "handles live markdown notebook source", %{conn: conn} do
       {:ok, session} = Sessions.create_session()
 
-      conn = get(conn, ~p"/sessions/#{session.id}/export/download/livemd")
+      conn = get(conn, ~p"/sessions/#{session.id}/download/export/livemd")
 
       assert conn.status == 200
+
+      assert get_resp_header(conn, "content-disposition") ==
+               [~s/attachment; filename="untitled_notebook.livemd"/]
+
       assert get_resp_header(conn, "content-type") == ["text/plain"]
 
       assert conn.resp_body == """
@@ -165,9 +230,13 @@ defmodule LivebookWeb.SessionControllerTest do
 
       {:ok, session} = Sessions.create_session(notebook: notebook)
 
-      conn = get(conn, ~p"/sessions/#{session.id}/export/download/livemd?include_outputs=true")
+      conn = get(conn, ~p"/sessions/#{session.id}/download/export/livemd?include_outputs=true")
 
       assert conn.status == 200
+
+      assert get_resp_header(conn, "content-disposition") ==
+               [~s/attachment; filename="my_notebook.livemd"/]
+
       assert get_resp_header(conn, "content-type") == ["text/plain"]
 
       assert conn.resp_body == """
@@ -192,9 +261,13 @@ defmodule LivebookWeb.SessionControllerTest do
     test "handles elixir notebook source", %{conn: conn} do
       {:ok, session} = Sessions.create_session()
 
-      conn = get(conn, ~p"/sessions/#{session.id}/export/download/exs")
+      conn = get(conn, ~p"/sessions/#{session.id}/download/export/exs")
 
       assert conn.status == 200
+
+      assert get_resp_header(conn, "content-disposition") ==
+               [~s/attachment; filename="untitled_notebook.exs"/]
+
       assert get_resp_header(conn, "content-type") == ["text/plain"]
 
       assert conn.resp_body == """

test/livebook_web/live/home_live_test.exs

@@ -73,7 +73,7 @@ defmodule LivebookWeb.HomeLiveTest do
         |> element(~s{[data-test-session-id="#{session.id}"] a}, "Download source")
         |> render_click
 
-      assert to == ~p"/sessions/#{session.id}/export/download/livemd?include_outputs=false"
+      assert to == ~p"/sessions/#{session.id}/download/export/livemd?include_outputs=false"
 
       Session.close(session.pid)
     end