Git of dcc - anni/blob - lib/pleroma/repo.ex

   1 # Pleroma: A lightweight social networking server
   2 # Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
   3 # SPDX-License-Identifier: AGPL-3.0-only
   4
   5 defmodule Pleroma.Repo do
   6   use Ecto.Repo,
   7     otp_app: :pleroma,
   8     adapter: Ecto.Adapters.Postgres,
   9     migration_timestamps: [type: :naive_datetime_usec]
  10
  11   import Ecto.Query
  12   require Logger
  13
  14   @doc """
  15   Dynamically loads the repository url from the
  16   DATABASE_URL environment variable.
  17   """
  18   def init(_, opts) do
  19     {:ok, Keyword.put(opts, :url, System.get_env("DATABASE_URL"))}
  20   end
  21
  22   @doc "find resource based on prepared query"
  23   @spec find_resource(Ecto.Query.t()) :: {:ok, struct()} | {:error, :not_found}
  24   def find_resource(%Ecto.Query{} = query) do
  25     case __MODULE__.one(query) do
  26       nil -> {:error, :not_found}
  27       resource -> {:ok, resource}
  28     end
  29   end
  30
  31   def find_resource(_query), do: {:error, :not_found}
  32
  33   @doc """
  34   Gets association from cache or loads if need
  35
  36   ## Examples
  37
  38     iex> Repo.get_assoc(token, :user)
  39     %User{}
  40
  41   """
  42   @spec get_assoc(struct(), atom()) :: {:ok, struct()} | {:error, :not_found}
  43   def get_assoc(resource, association) do
  44     case __MODULE__.preload(resource, association) do
  45       %{^association => assoc} when not is_nil(assoc) -> {:ok, assoc}
  46       _ -> {:error, :not_found}
  47     end
  48   end
  49
  50   @doc """
  51   Returns a lazy enumerable that emits all entries from the data store matching the given query.
  52
  53   `returns_as` use to group records. use the `batches` option to fetch records in bulk.
  54
  55   ## Examples
  56
  57   # fetch records one-by-one
  58   iex> Pleroma.Repo.chunk_stream(Pleroma.Activity.Queries.by_actor(ap_id), 500)
  59
  60   # fetch records in bulk
  61   iex> Pleroma.Repo.chunk_stream(Pleroma.Activity.Queries.by_actor(ap_id), 500, :batches)
  62   """
  63   @spec chunk_stream(Ecto.Query.t(), integer(), atom()) :: Enumerable.t()
  64   def chunk_stream(query, chunk_size, returns_as \\ :one, query_options \\ []) do
  65     # We don't actually need start and end functions of resource streaming,
  66     # but it seems to be the only way to not fetch records one-by-one and
  67     # have individual records be the elements of the stream, instead of
  68     # lists of records
  69     Stream.resource(
  70       fn -> 0 end,
  71       fn
  72         last_id ->
  73           query
  74           |> order_by(asc: :id)
  75           |> where([r], r.id > ^last_id)
  76           |> limit(^chunk_size)
  77           |> all(query_options)
  78           |> case do
  79             [] ->
  80               {:halt, last_id}
  81
  82             records ->
  83               last_id = List.last(records).id
  84
  85               if returns_as == :one do
  86                 {records, last_id}
  87               else
  88                 {[records], last_id}
  89               end
  90           end
  91       end,
  92       fn _ -> :ok end
  93     )
  94   end
  95 end