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