Implement table widget for Ecto queries (#34)

* Implement table widget for Ecto queries

* Use to_query instead of empty from

* Improve schema extraction

* Update lib/kino/utils/table.ex

Co-authored-by: José Valim <[email protected]>

* Add tests for the Ecto widget

Co-authored-by: José Valim <[email protected]>

lib/kino.ex

@@ -83,6 +83,13 @@ defmodule Kino do
       | 2  | Erlang | https://www.erlang.org  |
       \"\"\")
 
+  ### Kino.Ecto
+
+  `Kino.Ecto` implements a data table output for arbitrary
+  `Ecto` queries:
+
+      Kino.Ecto.new(Weather, Repo)
+
   ### All others
 
   All other data structures are rendered as text using Elixir's

lib/kino/data_table.ex

@@ -27,6 +27,8 @@ defmodule Kino.DataTable do
 
   use GenServer, restart: :temporary
 
+  alias Kino.Utils.Table
+
   defstruct [:pid]
 
   @type t :: %__MODULE__{pid: pid()}
@@ -146,15 +148,12 @@ defmodule Kino.DataTable do
   def handle_info({:connect, pid}, state) do
     columns =
       if state.keys do
-        Enum.map(state.keys, &key_to_column/1)
+        Table.keys_to_columns(state.keys)
       else
         []
       end
 
-    features =
-      [pagination: true, sorting: state.sorting_enabled]
-      |> Enum.filter(&elem(&1, 1))
-      |> Keyword.keys()
+    features = Kino.Utils.truthy_keys(pagination: true, sorting: state.sorting_enabled)
 
     send(pid, {:connect_reply, %{name: "Data", columns: columns, features: features}})
 
@@ -168,7 +167,7 @@ defmodule Kino.DataTable do
       if state.keys do
         {:initial, state.keys}
       else
-        columns = columns_structure(records)
+        columns = Table.columns_for_records(records)
 
         columns =
           if state.show_underscored,
@@ -179,7 +178,7 @@ defmodule Kino.DataTable do
         {columns, keys}
       end
 
-    rows = Enum.map(records, &record_to_row(&1, keys))
+    rows = Enum.map(records, &Table.record_to_row(&1, keys))
 
     send(pid, {:rows, %{rows: rows, total_rows: state.total_rows, columns: columns}})
 
@@ -190,95 +189,17 @@ defmodule Kino.DataTable do
     {:stop, :shutdown, state}
   end
 
-  defp columns_structure(records) do
-    case Enum.at(records, 0) do
-      nil ->
-        []
-
-      first_record ->
-        first_record_columns = columns_structure_for_record(first_record)
-
-        all_columns =
-          records
-          |> Enum.reduce(MapSet.new(), fn record, columns ->
-            record
-            |> columns_structure_for_record()
-            |> MapSet.new()
-            |> MapSet.union(columns)
-          end)
-          |> MapSet.to_list()
-          |> Enum.sort_by(& &1.key)
-
-        # If all records have the same structure, keep the order,
-        # otherwise sort the accumulated columns
-        if length(first_record_columns) == length(all_columns) do
-          first_record_columns
-        else
-          all_columns
-        end
-    end
-  end
-
-  defp columns_structure_for_record(record) when is_tuple(record) do
-    record
-    |> Tuple.to_list()
-    |> Enum.with_index()
-    |> Enum.map(fn {_, idx} -> key_to_column(idx) end)
-  end
-
-  defp columns_structure_for_record(record) when is_map(record) do
-    record
-    |> Map.keys()
-    |> Enum.sort()
-    |> Enum.map(&key_to_column/1)
-  end
-
-  defp columns_structure_for_record(record) when is_list(record) do
-    record
-    |> Keyword.keys()
-    |> Enum.map(&key_to_column/1)
-  end
-
-  defp key_to_column(key), do: %{key: key, label: inspect(key)}
-
   defp get_records(data, rows_spec) do
     sorted_data =
       if order_by = rows_spec[:order_by] do
-        Enum.sort_by(data, fn record -> get_field(record, order_by) end, rows_spec.order)
+        Enum.sort_by(data, &Table.get_field(&1, order_by), rows_spec.order)
       else
         data
       end
 
     Enum.slice(sorted_data, rows_spec.offset, rows_spec.limit)
   end
 
-  defp get_field(record, key) when is_tuple(record) do
-    if key < tuple_size(record) do
-      elem(record, key)
-    else
-      nil
-    end
-  end
-
-  defp get_field(record, key) when is_list(record) do
-    record[key]
-  end
-
-  defp get_field(record, key) when is_map(record) do
-    Map.get(record, key)
-  end
-
-  defp record_to_row(record, keys) do
-    fields =
-      Map.new(keys, fn key ->
-        value = get_field(record, key)
-        {key, inspect(value)}
-      end)
-
-    # Note: id is opaque to the client, and we don't need it for now
-    %{id: nil, fields: fields}
-  end
-
   defp underscored?(key) when is_atom(key) do
     key |> Atom.to_string() |> String.starts_with?("_")
   end

lib/kino/ecto.ex

@@ -0,0 +1,163 @@
+defmodule Kino.Ecto do
+  @moduledoc """
+  A widget for interactively viewing `Ecto` query results.
+
+  The data must be an enumerable of records, where each
+  record is either map, struct, keyword list or tuple.
+
+  ## Examples
+
+  The widget primarly allows for viewing a database table
+  given a schema:
+
+      Kino.Ecto.new(Weather, Repo)
+
+  However, the first argument can be any queryable, so
+  you can pipe arbitrary queries directly to the widget:
+
+      from(w in Weather, where: w.city == "New York")
+      |> Kino.Ecto.new(Repo)
+  """
+
+  use GenServer, restart: :temporary
+
+  alias Kino.Utils.Table
+
+  defstruct [:pid]
+
+  import Ecto.Query, only: [from: 2]
+
+  @type t :: %__MODULE__{pid: pid()}
+
+  @typedoc false
+  @type state :: %{
+          parent_monitor_ref: reference(),
+          repo: Ecto.Repo.t(),
+          queryable: Ecto.Queryable.t()
+        }
+
+  @doc """
+  Starts a widget process with the given queryable as
+  the data source.
+  """
+  @spec new(Ecto.Queryable.t(), Ecto.Repo.t()) :: t()
+  def new(queryable, repo) when is_atom(repo) do
+    unless queryable?(queryable) do
+      raise ArgumentError,
+            "expected a term implementing the Ecto.Queryable protocol, got: #{inspect(queryable)}"
+    end
+
+    parent = self()
+    opts = [repo: repo, queryable: queryable, parent: parent]
+
+    {:ok, pid} = DynamicSupervisor.start_child(Kino.WidgetSupervisor, {__MODULE__, opts})
+
+    %__MODULE__{pid: pid}
+  end
+
+  defp queryable?(term) do
+    Ecto.Queryable.impl_for(term) != nil
+  end
+
+  @doc false
+  def start_link(opts) do
+    GenServer.start_link(__MODULE__, opts)
+  end
+
+  @impl true
+  def init(opts) do
+    repo = Keyword.fetch!(opts, :repo)
+    queryable = Keyword.fetch!(opts, :queryable)
+    parent = Keyword.fetch!(opts, :parent)
+
+    parent_monitor_ref = Process.monitor(parent)
+
+    {:ok, %{parent_monitor_ref: parent_monitor_ref, repo: repo, queryable: queryable}}
+  end
+
+  @impl true
+  def handle_info({:connect, pid}, state) do
+    name = state.queryable |> query_source() |> to_string()
+    columns = state.queryable |> keys_from_queryable() |> Table.keys_to_columns()
+
+    features =
+      Kino.Utils.truthy_keys(
+        refetch: true,
+        pagination: true,
+        # If the user specifies custom select, the record keys
+        # are not valid "order by" fields, so we disable sorting
+        sorting: default_select_query?(state.queryable)
+      )
+
+    send(
+      pid,
+      {:connect_reply, %{name: name, columns: columns, features: features}}
+    )
+
+    {:noreply, state}
+  end
+
+  def handle_info({:get_rows, pid, rows_spec}, state) do
+    {total_rows, records} = get_records(state.repo, state.queryable, rows_spec)
+
+    {columns, keys} =
+      case keys_from_queryable(state.queryable) do
+        [] ->
+          columns = Table.columns_for_records(records)
+          keys = Enum.map(columns, & &1.key)
+          {columns, keys}
+
+        keys ->
+          {:initial, keys}
+      end
+
+    rows = Enum.map(records, &Table.record_to_row(&1, keys))
+
+    send(pid, {:rows, %{rows: rows, total_rows: total_rows, columns: columns}})
+
+    {:noreply, state}
+  end
+
+  def handle_info({:DOWN, ref, :process, _object, _reason}, %{parent_monitor_ref: ref} = state) do
+    {:stop, :shutdown, state}
+  end
+
+  defp get_records(repo, queryable, rows_spec) do
+    count = repo.aggregate(queryable, :count)
+
+    query = from(q in queryable, limit: ^rows_spec.limit, offset: ^rows_spec.offset)
+
+    query =
+      if rows_spec[:order_by] do
+        query = Ecto.Query.exclude(query, :order_by)
+        order_by = [{rows_spec.order, rows_spec.order_by}]
+        from(q in query, order_by: ^order_by)
+      else
+        query
+      end
+
+    records = repo.all(query)
+
+    {count, records}
+  end
+
+  defp query_source(queryable) do
+    %{from: %{source: {source, _schema}}} = Ecto.Queryable.to_query(queryable)
+    source
+  end
+
+  defp default_select_query?(queryable) do
+    query = Ecto.Queryable.to_query(queryable)
+    query.select == nil
+  end
+
+  defp keys_from_queryable(queryable) do
+    schema = Table.ecto_schema(queryable)
+
+    if schema != nil and default_select_query?(queryable) do
+      schema.__schema__(:fields)
+    else
+      []
+    end
+  end
+end

lib/kino/ets.ex

@@ -12,6 +12,8 @@ defmodule Kino.ETS do
 
   use GenServer, restart: :temporary
 
+  alias Kino.Utils.Table
+
   defstruct [:pid]
 
   @type t :: %__MODULE__{pid: pid()}
@@ -77,7 +79,7 @@ defmodule Kino.ETS do
 
     columns =
       case :ets.match_object(state.tid, :_, 1) do
-        {[record], _} -> columns_structure_for_records([record])
+        {[record], _} -> Table.columns_for_records([record])
         :"$end_of_table" -> []
       end
 
@@ -97,7 +99,7 @@ defmodule Kino.ETS do
     columns =
       case records do
         [] -> :initial
-        records -> columns_structure_for_records(records)
+        records -> Table.columns_for_records(records)
       end
 
     send(pid, {:rows, %{rows: rows, total_rows: total_rows, columns: columns}})
@@ -109,17 +111,6 @@ defmodule Kino.ETS do
     {:stop, :shutdown, state}
   end
 
-  defp columns_structure_for_records(records) do
-    max_columns =
-      records
-      |> Enum.map(&tuple_size/1)
-      |> Enum.max()
-
-    for idx <- 0..(max_columns - 1) do
-      %{key: idx, label: to_string(idx)}
-    end
-  end
-
   defp get_records(tid, rows_spec) do
     query = :ets.table(tid)
     cursor = :qlc.cursor(query)

lib/kino/render.ex

@@ -50,6 +50,12 @@ defimpl Kino.Render, for: Kino.Markdown do
   end
 end
 
+defimpl Kino.Render, for: Kino.Ecto do
+  def to_livebook(widget) do
+    Kino.Output.table_dynamic(widget.pid)
+  end
+end
+
 # Elixir built-ins
 
 defimpl Kino.Render, for: Reference do

lib/kino/utils.ex

@@ -0,0 +1,21 @@
+defmodule Kino.Utils do
+  @moduledoc false
+
+  @doc """
+  Returns keyword list keys that hold a truthy value.
+
+  ## Examples
+
+      iex> Kino.Utils.truthy_keys(cat: true, dog: false)
+      [:cat]
+
+      iex> Kino.Utils.truthy_keys(tea: :ok, coffee: nil)
+      [:tea]
+  """
+  @spec truthy_keys(keyword()) :: list(atom())
+  def truthy_keys(keywords) when is_list(keywords) do
+    keywords
+    |> Enum.filter(&elem(&1, 1))
+    |> Keyword.keys()
+  end
+end