git.squeep.com Git - akkoma/blob - lib/pleroma/repo.ex

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