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

   1 # Pleroma: A lightweight social networking server
   2 # Copyright © 2017-2020 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   def chunk_stream(query, chunk_size) do
  53     # We don't actually need start and end funcitons of resource streaming,
  54     # but it seems to be the only way to not fetch records one-by-one and
  55     # have individual records be the elements of the stream, instead of
  56     # lists of records
  57     Stream.resource(
  58       fn -> 0 end,
  59       fn
  60         last_id ->
  61           query
  62           |> order_by(asc: :id)
  63           |> where([r], r.id > ^last_id)
  64           |> limit(^chunk_size)
  65           |> all()
  66           |> case do
  67             [] ->
  68               {:halt, last_id}
  69
  70             records ->
  71               last_id = List.last(records).id
  72               {records, last_id}
  73           end
  74       end,
  75       fn _ -> :ok end
  76     )
  77   end
  78 end