postgrex

defmodule Postgrex.SimpleConnection do
  @moduledoc ~S"""
  A generic connection suitable for simple queries and pubsub functionality.

  On its own, a SimpleConnection server only maintains a connection. To execute
  queries, process results, or relay notices you must implement a callback module
  with the SimpleConnection behaviour.

  ## Example

  The SimpleConnection behaviour abstracts common client/server interactions,
  along with optional mechanisms for running queries or relaying notifications.

  Let's start with a minimal callback module that executes a query and relays
  the result back to the caller.

      defmodule MyConnection do
        @behaviour Postgrex.SimpleConnection

        @impl true
        def init(_args) do
          {:ok, %{from: nil}}
        end

        @impl true
        def handle_call({:query, query}, from, state) do
          {:query, query, %{state | from: from}}
        end

        @impl true
        def handle_result(results, state) when is_list(results) do
          SimpleConnection.reply(state.from, results)

          {:noreply, state}
        end

        @impl true
        def handle_result(%Postgrex.Error{} = error, state) do
          SimpleConnection.reply(state.from, error)

          {:noreply, state}
        end
      end

      # Start the connection
      {:ok, pid} = SimpleConnection.start_link(MyConnection, [], database: "demo")

      # Execute a literal query
      SimpleConnection.call(pid, {:query, "SELECT 1"})
      # => %Postgrex.Result{rows: [["1"]]}

  We start a connection by passing the callback module, callback options, and
  server options to `SimpleConnection.start_link/3`. The `init/1` function
  receives any callback options and returns the callback state.

  Queries are sent through `SimpleConnection.call/2`, executed on the server,
  and the result is handed off to `handle_result/2`. At that point the callback
  can process the result before replying back to the caller with
  `SimpleConnection.reply/2`.

  ## Building a PubSub Connection

  With the `notify/3` callback you can also build a pubsub server on top of
  `LISTEN/NOTIFY`. Here's a naive pubsub implementation:

      defmodule MyPubSub do
        @behaviour Postgrex.SimpleConnection

        defstruct [:from, listeners: %{}]

        @impl true
        def init(args) do
          {:ok, struct!(__MODULE__, args)}
        end

        @impl true
        def notify(channel, payload, state) do
          for pid <- state.listeners[channel] do
            send(pid, {:notice, channel, payload})
          end
        end

        @impl true
        def handle_call({:listen, channel}, {pid, _} = from, state) do
          listeners = Map.update(state.listeners, channel, [pid], &[pid | &1])

          {:query, ~s(LISTEN "#{channel}"), %{state | from: from, listeners: listeners}}
        end

        def handle_call({:query, query}, from, state) do
          {:query, query, %{state | from: from}}
        end

        @impl true
        def handle_result(_results, state) do
          SimpleConnection.reply(state.from, :ok)

          {:noreply, %{state | from: nil}}
        end
      end

      # Start the connection
      {:ok, pid} = SimpleConnection.start_link(MyPubSub, [], database: "demo")

      # Start listening to the "demo" channel
      SimpleConnection.call(pid, {:listen, "demo"})
      # => %Postgrex.Result{command: :listen}

      # Notify all listeners
      SimpleConnection.call(pid, {:query, ~s(NOTIFY "demo", 'hello')})
      # => %Postgrex.Result{command: :notify}

      # Check the inbox to see the notice message
      flush()
      # => {:notice, "demo", "hello"}

  See `Postgrex.Notifications` for a more complex implementation that can
  unlisten, handle process exits, and persist across reconnection.

  ## Name registration

  A `Postgrex.SimpleConnection` is bound to the same name registration rules as a
  `GenServer`. Read more about them in the `GenServer` docs.
  """

  @behaviour :gen_statem

  require Logger

  alias Postgrex.Protocol

  @doc false
  defstruct idle_interval: 5000,
            protocol: nil,
            auto_reconnect: false,
            reconnect_backoff: 500,
            state: nil

  ## PUBLIC API ##

  @type query :: iodata
  @type state :: term

  @typedoc since: "0.17.0"
  @type from :: {pid, term}

  @doc """
  Callback for process initialization.

  This is called once and before the Postgrex connection is established.
  """
  @callback init(term) :: {:ok, state}

  @doc """
  Callback for processing or relaying pubsub notifications.
  """
  @callback notify(binary, binary, state) :: :ok

  @doc """
  Invoked after connecting or reconnecting.

  This may be called multiple times if `:auto_reconnect` is true.
  """
  @callback handle_connect(state) :: {:noreply, state} | {:query, query, state}

  @doc """
  Invoked after disconnection.

  This is invoked regardless of the `:auto_reconnect` option.
  """
  @callback handle_disconnect(state) :: {:noreply, state}

  @doc """
  Callback for `SimpleConnection.call/3`.

  Replies must be sent with `SimpleConnection.reply/2`.
  """
  @callback handle_call(term, from, state) ::
              {:noreply, state} | {:query, query, state}

  @doc """
  Callback for `Kernel.send/2`.
  """
  @callback handle_info(term, state) :: {:noreply, state} | {:query, query, state}

  @doc """
  Callback for processing or relaying queries executed via `{:query, query, state}`.

  Either a list of successful query results or an error will be passed to this callback.
  A list is passed because the simple query protocol allows multiple commands to be
  issued in a single query.
  """
  @callback handle_result([Postgrex.Result.t()] | Postgrex.Error.t(), state) ::
              {:noreply, state}

  @optional_callbacks handle_call: 3,
                      handle_connect: 1,
                      handle_disconnect: 1,
                      handle_info: 2,
                      handle_result: 2

  @doc """
  Replies to the given client.

  Wrapper for `:gen_statem.reply/2`.
  """
  def reply({caller_pid, from} = _from, reply) when is_pid(caller_pid) do
    :gen_statem.reply(from, reply)
  end

  @doc """
  Calls the given server.

  Wrapper for `:gen_statem.call/3`.
  """
  def call(server, message, timeout \\ 5000) do
    with {__MODULE__, reason} <- :gen_statem.call(server, {message, self()}, timeout) do
      exit({reason, {__MODULE__, :call, [server, message, timeout]}})
    end
  end

  @doc false
  def child_spec(opts) do
    %{id: __MODULE__, start: {__MODULE__, :start_link, opts}}
  end

  @doc """
  Start the connection process and connect to Postgres.

  The options that this function accepts are the same as those accepted by
  `Postgrex.start_link/1`, as well as the extra options `:sync_connect`,
  `:auto_reconnect`, `:reconnect_backoff`, and `:configure`.

  ## Options

    * `:auto_reconnect` - automatically attempt to reconnect to the database
      in event of a disconnection. See the
      [note about async connect and auto-reconnects](#module-async-connect-and-auto-reconnects)
      above. Defaults to `false`, which means the process terminates.

    * `:configure` - A function to run before every connect attempt to dynamically
      configure the options as a `{module, function, args}`, where the current
      options will prepended to `args`. Defaults to `nil`.

    * `:idle_interval` - while also accepted on `Postgrex.start_link/1`, it has
      a default of `5000ms` in `Postgrex.SimpleConnection` (instead of 1000ms).

    * `:reconnect_backoff` - time (in ms) between reconnection attempts when
      `auto_reconnect` is enabled. Defaults to `500`.

    * `:sync_connect` - controls if the connection should be established on boot
      or asynchronously right after boot. Defaults to `true`.
  """
  @spec start_link(module, term, Keyword.t()) :: {:ok, pid} | {:error, Postgrex.Error.t() | term}
  def start_link(module, args, opts) do
    {name, opts} = Keyword.pop(opts, :name)

    opts = Keyword.put_new(opts, :sync_connect, true)
    connection_opts = Postgrex.Utils.default_opts(opts)
    start_args = {module, args, connection_opts}

    case name do
      nil ->
        :gen_statem.start_link(__MODULE__, start_args, [])

      atom when is_atom(atom) ->
        :gen_statem.start_link({:local, atom}, __MODULE__, start_args, [])

      {:global, _term} = tuple ->
        :gen_statem.start_link(tuple, __MODULE__, start_args, [])

      {:via, via_module, _term} = tuple when is_atom(via_module) ->
        :gen_statem.start_link(tuple, __MODULE__, start_args, [])

      other ->
        raise ArgumentError, """
        expected :name option to be one of the following:

          * nil
          * atom
          * {:global, term}
          * {:via, module, term}

        Got: #{inspect(other)}
        """
    end
  end

  ## CALLBACKS ##

  @state :no_state

  @doc false
  @impl :gen_statem
  def callback_mode, do: :handle_event_function

  @doc false
  @impl :gen_statem
  def init({mod, args, opts}) do
    case mod.init(args) do
      {:ok, mod_state} ->
        idle_timeout = opts[:idle_timeout]

        if idle_timeout do
          Logger.warning(
            ":idle_timeout in Postgrex.SimpleConnection is deprecated, " <>
              "please use :idle_interval instead"
          )
        end

        {idle_interval, opts} = Keyword.pop(opts, :idle_interval, idle_timeout || 5000)
        {auto_reconnect, opts} = Keyword.pop(opts, :auto_reconnect, false)
        {reconnect_backoff, opts} = Keyword.pop(opts, :reconnect_backoff, 500)

        state = %__MODULE__{
          idle_interval: idle_interval,
          auto_reconnect: auto_reconnect,
          reconnect_backoff: reconnect_backoff,
          state: {mod, mod_state}
        }

        put_opts(mod, opts)

        if opts[:sync_connect] do
          case handle_event(:internal, {:connect, :init}, @state, state) do
            {:keep_state, state} -> {:ok, @state, state}
            {:keep_state, state, actions} -> {:ok, @state, state, actions}
            {:stop, _reason, _state} = stop -> stop
          end
        else
          {:ok, @state, state, {:next_event, :internal, {:connect, :init}}}
        end
    end
  end

  @doc false
  @impl :gen_statem
  def handle_event(type, content, statem_state, state)

  def handle_event(:internal, {:connect, _}, @state, %{state: {mod, mod_state}} = state) do
    opts =
      case Keyword.get(opts(mod), :configure) do
        {module, fun, args} -> apply(module, fun, [opts(mod) | args])
        fun when is_function(fun, 1) -> fun.(opts(mod))
        nil -> opts(mod)
      end

    case Protocol.connect(opts) do
      {:ok, protocol} ->
        state = %{state | protocol: protocol}

        with {:keep_state, state, _} <- maybe_handle(mod, :handle_connect, [mod_state], state) do
          {:keep_state, state}
        end

      {:error, reason} ->
        if state.auto_reconnect do
          {:keep_state, state, {{:timeout, :backoff}, state.reconnect_backoff, nil}}
        else
          {:stop, reason, state}
        end
    end
  end

  def handle_event({:timeout, :backoff}, nil, @state, state) do
    {:keep_state, state, {:next_event, :internal, {:connect, :reconnect}}}
  end

  def handle_event({:call, from}, {msg, caller_pid}, @state, %{state: {mod, mod_state}} = state) do
    # We have to do a hack here to carry the actual caller PID over to the handle_call/3
    # callback, because gen_statem uses a proxy process to do calls with timeout != :infinity.
    # This results in the caller PID not being the same as the PID in the "from" tuple,
    # so things like Postgrex.Notifications cannot use that "from"'s PID to register
    # notification handlers. This approach is paired with reconstructing the proper
    # "from" tuple in the reply/2 function in this module.
    callback_from = {caller_pid, from}
    handle(mod, :handle_call, [msg, callback_from, mod_state], from, state)
  end

  def handle_event(:timeout, nil, @state, %{protocol: protocol} = state) do
    case Protocol.ping(protocol) do
      {:ok, protocol} ->
        {:keep_state, %{state | protocol: protocol}, {:timeout, state.idle_interval, nil}}

      {error, reason, protocol} ->
        reconnect_or_stop(error, reason, protocol, state)
    end
  end

  def handle_event(:info, msg, @state, %{protocol: protocol, state: {mod, mod_state}} = state) do
    opts = [notify: &mod.notify(&1, &2, mod_state)]

    case Protocol.handle_info(msg, opts, protocol) do
      {:ok, protocol} ->
        {:keep_state, %{state | protocol: protocol}, {:timeout, state.idle_interval, nil}}

      {:unknown, protocol} ->
        maybe_handle(mod, :handle_info, [msg, mod_state], %{state | protocol: protocol})

      {error, reason, protocol} ->
        reconnect_or_stop(error, reason, protocol, state)
    end
  end

  def handle_event(:info, msg, @state, %{state: {mod, mod_state}} = state) do
    maybe_handle(mod, :handle_info, [msg, mod_state], state)
  end

  @doc false
  @impl :gen_statem
  def format_status(_opts, [_pdict, @state, state]) do
    state
  end

  ## Helpers

  defp maybe_handle(mod, fun, args, state) do
    if function_exported?(mod, fun, length(args)) do
      handle(mod, fun, args, nil, state)
    else
      {:keep_state, state, {:timeout, state.idle_interval, nil}}
    end
  end

  defp handle(mod, fun, args, from, state) do
    case apply(mod, fun, args) do
      {:noreply, mod_state} ->
        {:keep_state, %{state | state: {mod, mod_state}}, {:timeout, state.idle_interval, nil}}

      {:query, query, mod_state} ->
        opts = [notify: &mod.notify(&1, &2, mod_state)]

        state = %{state | state: {mod, mod_state}}

        with {:ok, results, protocol} <- Protocol.handle_simple(query, opts, state.protocol),
             {:ok, protocol} <- Protocol.checkin(protocol) do
          state = %{state | protocol: protocol}

          handle(mod, :handle_result, [results, mod_state], from, state)
        else
          {:error, %Postgrex.Error{} = error, protocol} ->
            handle(mod, :handle_result, [error, mod_state], from, %{state | protocol: protocol})

          {:disconnect, reason, protocol} ->
            reconnect_or_stop(:disconnect, reason, protocol, state)
        end
    end
  end

  defp reconnect_or_stop(error, reason, protocol, %{state: {mod, mod_state}} = state)
       when error in [:error, :disconnect] do
    {:keep_state, state, _actions} = maybe_handle(mod, :handle_disconnect, [mod_state], state)

    if state.auto_reconnect do
      {:keep_state, state, {:next_event, :internal, {:connect, :reconnect}}}
    else
      {:stop, reason, %{state | protocol: protocol}}
    end
  end

  defp opts(mod), do: Process.get(mod)

  defp put_opts(mod, opts), do: Process.put(mod, opts)
end