Add support for accessing notebook files(#319)

Co-authored-by: Jonatan Kłosko <[email protected]>

lib/kino/bridge.ex

@@ -76,6 +76,24 @@ defmodule Kino.Bridge do
     with {:ok, reply} <- io_request({:livebook_get_file_path, file_ref}), do: reply
   end
 
+  @doc """
+  Requests the file path for notebook file with the given name.
+  """
+  @spec get_file_entry_path(String.t()) ::
+          {:ok, term()} | {:error, request_error() | :forbidden | String.t()}
+  def get_file_entry_path(name) do
+    with {:ok, reply} <- io_request({:livebook_get_file_entry_path, name}), do: reply
+  end
+
+  @doc """
+  Requests the file spec for notebook file with the given name.
+  """
+  @spec get_file_entry_spec(String.t()) ::
+          {:ok, term()} | {:error, request_error() | :forbidden | String.t()}
+  def get_file_entry_spec(name) do
+    with {:ok, reply} <- io_request({:livebook_get_file_entry_spec, name}), do: reply
+  end
+
   @doc """
   Associates `object` with `pid`.
 

lib/kino/fs.ex

@@ -0,0 +1,100 @@
+defmodule Kino.FS do
+  @moduledoc """
+  Provides access to notebook files.
+  """
+
+  defmodule ForbiddenError do
+    @moduledoc """
+    Exception raised when access to a notebook file is forbidden.
+    """
+
+    defexception [:name]
+
+    @impl true
+    def message(exception) do
+      "forbidden access to file #{inspect(exception.name)}"
+    end
+  end
+
+  @doc """
+  Accesses notebook file with the given name and returns a local path
+  to read its contents from.
+
+  This invocation may take a while, in case the file is downloaded
+  from a URL and is not in the cache.
+
+  > #### File operations {: .info}
+  >
+  > You should treat the file as read-only. To avoid unnecessary
+  > copies the path may potentially be pointing to the original file,
+  > in which case any write operations would be persisted. This
+  > behaviour is not always the case, so you should not rely on it
+  > either.
+  """
+  @spec file_path(String.t()) :: String.t()
+  def file_path(name) when is_binary(name) do
+    case Kino.Bridge.get_file_entry_path(name) do
+      {:ok, path} ->
+        path
+
+      {:error, :forbidden} ->
+        raise ForbiddenError, name: name
+
+      {:error, message} when is_binary(message) ->
+        raise message
+
+      {:error, reason} when is_atom(reason) ->
+        raise "failed to access file path, reason: #{inspect(reason)}"
+    end
+  end
+
+  @doc """
+  Accesses notebook file with the given name and returns a specification
+  of the file location.
+
+  This does not copy any files and moves the responsibility of reading
+  the file to the caller. If you need to read a file directly, use
+  `file_path/1`.
+  """
+  @spec file_spec(String.t()) :: FSS.entry()
+  def file_spec(name) do
+    case Kino.Bridge.get_file_entry_spec(name) do
+      {:ok, spec} ->
+        file_spec_to_fss(spec)
+
+      {:error, :forbidden} ->
+        raise ForbiddenError, name: name
+
+      {:error, message} when is_binary(message) ->
+        raise message
+
+      {:error, reason} when is_atom(reason) ->
+        raise "failed to access file spec, reason: #{inspect(reason)}"
+    end
+  end
+
+  defp file_spec_to_fss(%{type: :local} = file_spec) do
+    FSS.Local.from_path(file_spec.path)
+  end
+
+  defp file_spec_to_fss(%{type: :url} = file_spec) do
+    case FSS.HTTP.parse(file_spec.url) do
+      {:ok, entry} -> entry
+      {:error, error} -> raise error
+    end
+  end
+
+  defp file_spec_to_fss(%{type: :s3} = file_spec) do
+    case FSS.S3.parse("s3:///" <> file_spec.key,
+           config: [
+             region: file_spec.region,
+             endpoint: file_spec.bucket_url,
+             access_key_id: file_spec.access_key_id,
+             secret_access_key: file_spec.secret_access_key
+           ]
+         ) do
+      {:ok, entry} -> entry
+      {:error, error} -> raise error
+    end
+  end
+end

mix.exs

@@ -32,6 +32,7 @@ defmodule Kino.MixProject do
   defp deps do
     [
       {:table, "~> 0.1.2"},
+      {:fss, github: "elixir-explorer/fss", branch: "main"},
       {:nx, "~> 0.1", optional: true},
       {:ex_doc, "~> 0.28", only: :dev, runtime: false}
     ]

mix.lock

@@ -2,6 +2,7 @@
   "complex": {:hex, :complex, "0.4.2", "923e5db0be13dbb3ea00cf8459d9f75f3afdd9ff5a82742ded21064330d28273", [:mix], [], "hexpm", "069a085ef820ce675a2619fd125b963ff4514af2102c7f7d7965128e5ec0a429"},
   "earmark_parser": {:hex, :earmark_parser, "1.4.33", "3c3fd9673bb5dcc9edc28dd90f50c87ce506d1f71b70e3de69aa8154bc695d44", [:mix], [], "hexpm", "2d526833729b59b9fdb85785078697c72ac5e5066350663e5be6a1182da61b8f"},
   "ex_doc": {:hex, :ex_doc, "0.29.4", "6257ecbb20c7396b1fe5accd55b7b0d23f44b6aa18017b415cb4c2b91d997729", [:mix], [{:earmark_parser, "~> 1.4.31", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "2c6699a737ae46cb61e4ed012af931b57b699643b24dabe2400a8168414bc4f5"},
+  "fss": {:git, "https://github.com/elixir-explorer/fss.git", "19ed6ce8359e9790e818b732ba8d552aaf36e029", [branch: "main"]},
   "makeup": {:hex, :makeup, "1.1.0", "6b67c8bc2882a6b6a445859952a602afc1a41c2e08379ca057c0f525366fc3ca", [:mix], [{:nimble_parsec, "~> 1.2.2 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "0a45ed501f4a8897f580eabf99a2e5234ea3e75a4373c8a52824f6e873be57a6"},
   "makeup_elixir": {:hex, :makeup_elixir, "0.16.1", "cc9e3ca312f1cfeccc572b37a09980287e243648108384b97ff2b76e505c3555", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "e127a341ad1b209bd80f7bd1620a15693a9908ed780c3b763bccf7d200c767c6"},
   "makeup_erlang": {:hex, :makeup_erlang, "0.1.1", "3fcb7f09eb9d98dc4d208f49cc955a34218fc41ff6b84df7c75b3e6e533cc65f", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "174d0809e98a4ef0b3309256cbf97101c6ec01c4ab0b23e926a9e17df2077cbb"},

test/kino/fs_test.exs

@@ -0,0 +1,96 @@
+defmodule Kino.FSTest do
+  use ExUnit.Case, async: true
+
+  describe "file_spec/1" do
+    test "returns a file spec" do
+      name = "file.txt"
+      path = "/home/bob/file.txt"
+      spec = %{type: :local, path: path}
+
+      configure_gl_with_reply({:livebook_get_file_entry_spec, name}, {:ok, spec})
+
+      assert %FSS.Local.Entry{path: ^path} = Kino.FS.file_spec(name)
+    end
+
+    test "returns an HTTP FSS entry" do
+      name = "remote-file.txt"
+      url = "https://example.com/remote-file.txt"
+      spec = %{type: :url, url: url}
+
+      configure_gl_with_reply({:livebook_get_file_entry_spec, name}, {:ok, spec})
+
+      assert %FSS.HTTP.Entry{url: ^url, config: %FSS.HTTP.Config{headers: []}} =
+               Kino.FS.file_spec(name)
+    end
+
+    test "returns a S3 FSS entry" do
+      name = "file-from-s3.txt"
+      bucket_url = "https://s3.us-west-1.amazonaws.com/my-bucket"
+
+      spec = %{
+        type: :s3,
+        bucket_url: bucket_url,
+        region: "us-west-1",
+        access_key_id: "access-key-1",
+        secret_access_key: "secret-access-key-1",
+        key: "file-from-s3.txt"
+      }
+
+      configure_gl_with_reply({:livebook_get_file_entry_spec, name}, {:ok, spec})
+
+      assert %FSS.S3.Entry{} = s3 = Kino.FS.file_spec(name)
+
+      assert s3.key == spec.key
+
+      assert s3.config.region == spec.region
+      assert s3.config.endpoint == bucket_url
+      assert s3.config.access_key_id == spec.access_key_id
+      assert s3.config.secret_access_key == spec.secret_access_key
+      assert s3.config.bucket == nil
+    end
+
+    test "raises an error in case s3 file_spec has something nil" do
+      name = "file-from-s3.txt"
+
+      spec = %{
+        type: :s3,
+        bucket_url: nil,
+        region: "us-west-1",
+        access_key_id: "access-key-1",
+        secret_access_key: "secret-access-key-1",
+        key: name
+      }
+
+      configure_gl_with_reply({:livebook_get_file_entry_spec, name}, {:ok, spec})
+
+      assert_raise ArgumentError, "endpoint is required when bucket is nil", fn ->
+        Kino.FS.file_spec(name)
+      end
+
+      bucket_url = "https://s3.us-west-1.amazonaws.com/my-bucket"
+
+      spec =
+        spec
+        |> Map.replace!(:bucket_url, bucket_url)
+        |> Map.replace!(:access_key_id, nil)
+
+      configure_gl_with_reply({:livebook_get_file_entry_spec, name}, {:ok, spec})
+
+      assert_raise ArgumentError,
+                   "missing :access_key_id for FSS.S3 (set the key or the AWS_ACCESS_KEY_ID env var)",
+                   fn ->
+                     Kino.FS.file_spec(name)
+                   end
+    end
+  end
+
+  defp configure_gl_with_reply(request, reply) do
+    gl =
+      spawn(fn ->
+        assert_receive {:io_request, from, reply_as, ^request}
+        send(from, {:io_reply, reply_as, reply})
+      end)
+
+    Process.group_leader(self(), gl)
+  end
+end