nebula_graph_ex

defmodule NebulaGraphEx.Connection do
  @moduledoc """
  `DBConnection` behaviour implementation for NebulaGraph.

  Each connection in the pool owns one TCP (or TLS) socket and one
  authenticated NebulaGraph session. The session ID is held in the connection
  state and is automatically reused for every query checked out from the pool.

  You should not use this module directly — `NebulaGraphEx.Graph` is the
  public interface. This module is documented for library developers who want
  to understand the connection lifecycle or implement custom pool strategies.

  ## Connection lifecycle

  1. `connect/1` — opens the socket, performs the version handshake, and
     authenticates to get a `session_id`.
  2. `ping/1` — executes `YIELD 1` to keep the session alive. Called by
     `DBConnection` on idle connections at `:idle_interval`.
  3. `handle_execute/4` — encodes and sends the query, decodes the response
     into `%NebulaGraphEx.Result{}`.
  4. `disconnect/2` — signs out the session and closes the socket.

  ## Connection state

  The state map held by DBConnection:

  * `:socket` — the `:gen_tcp` or `:ssl` socket
  * `:session_id` — integer NebulaGraph session ID
  * `:opts` — resolved pool options from `NebulaGraphEx.Options`
  * `:seq` — monotonically incrementing Thrift sequence number
  """

  use DBConnection

  alias NebulaGraphEx.Error
  alias NebulaGraphEx.Options
  alias NebulaGraphEx.Query
  alias NebulaGraphEx.Result
  alias NebulaGraphEx.Thrift.GraphService
  alias NebulaGraphEx.Transport

  @impl DBConnection
  def connect(opts) do
    host = Keyword.get(opts, :hostname, "localhost")
    port = Keyword.get(opts, :port, 9669)

    case Transport.connect(host, port, opts) do
      {:ok, socket} -> authenticate(socket, opts)
      {:error, reason} -> {:error, conn_error(reason, opts)}
    end
  end

  @impl DBConnection
  def disconnect(_err, %{socket: socket, session_id: session_id} = _state) do
    if is_integer(session_id) do
      GraphService.signout(socket, session_id)
    end

    Transport.close(socket)
    :ok
  end

  @impl DBConnection
  def ping(%{socket: socket, session_id: session_id, opts: opts} = state) do
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    case GraphService.execute(socket, session_id, "YIELD 1", recv_timeout: recv_timeout) do
      {:ok, %{error_code: 0}} ->
        {:ok, state}

      {:ok, %{error_code: code, error_msg: msg}} ->
        {:disconnect, Error.from_response(code, msg), state}

      {:error, reason} ->
        {:disconnect, Error.connection_error(reason), state}
    end
  end

  @impl DBConnection
  def checkout(state), do: {:ok, state}

  @impl DBConnection
  def handle_status(_opts, state), do: {:idle, state}

  @impl DBConnection
  def handle_execute(
        %Query{statement: stmt, params: params, opts: query_opts},
        _params_arg,
        call_opts,
        state
      ) do
    opts = Keyword.merge(state.opts, Keyword.merge(query_opts, call_opts))
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)
    prefer_json = Keyword.get(opts, :prefer_json, false)
    override_space = Keyword.get(opts, :space)

    with :ok <- maybe_use_space(state, override_space, recv_timeout),
         {:ok, response} <- run_query(state, stmt, params, prefer_json, recv_timeout) do
      case response do
        %{error_code: 0} ->
          result = Result.from_response(response, stmt)
          {:ok, %Query{statement: stmt}, result, state}

        %{error_code: code, error_msg: msg} ->
          {:error, Error.from_response(code, msg, stmt), state}
      end
    else
      {:error, reason} when is_binary(reason) ->
        {:error, Error.client_error(reason), state}

      {:error, reason} ->
        {:disconnect, Error.connection_error(reason), state}
    end
  end

  @impl DBConnection
  def handle_prepare(query, _opts, state), do: {:ok, query, state}

  @impl DBConnection
  def handle_close(_query, _opts, state), do: {:ok, nil, state}

  @impl DBConnection
  def handle_declare(_query, _params, _opts, state), do: {:error, :not_supported, state}

  @impl DBConnection
  def handle_fetch(_query, _cursor, _opts, state), do: {:error, :not_supported, state}

  @impl DBConnection
  def handle_deallocate(_query, _cursor, _opts, state), do: {:ok, nil, state}

  @impl DBConnection
  def handle_begin(_opts, state),
    do: {:error, Error.client_error("NebulaGraph does not support transactions"), state}

  @impl DBConnection
  def handle_commit(_opts, state),
    do: {:error, Error.client_error("NebulaGraph does not support transactions"), state}

  @impl DBConnection
  def handle_rollback(_opts, state),
    do: {:error, Error.client_error("NebulaGraph does not support transactions"), state}

  # ─── Private ───────────────────────────────────────────────────────────────

  defp authenticate(socket, opts) do
    username = Keyword.get(opts, :username, "root")
    password = Options.resolve_password(Keyword.get(opts, :password, "nebula"))
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    with {:ok, :ok} <- GraphService.verify_client_version(socket, recv_timeout: recv_timeout),
         {:ok, auth} <-
           GraphService.authenticate(socket, username, password, recv_timeout: recv_timeout) do
      case auth do
        %{error_code: 0, session_id: session_id} ->
          state = %{
            socket: socket,
            session_id: session_id,
            opts: opts,
            seq: 6,
            time_zone_offset: auth[:time_zone_offset_seconds] || 0
          }

          maybe_select_space(state, opts)

        %{error_code: code, error_msg: msg} ->
          Transport.close(socket)
          {:error, Error.from_response(code, msg)}
      end
    else
      {:error, reason} ->
        Transport.close(socket)
        {:error, conn_error(reason, opts)}
    end
  end

  defp maybe_select_space(state, opts) do
    case Keyword.get(opts, :space) do
      nil ->
        {:ok, state}

      space when is_binary(space) ->
        recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

        case GraphService.execute(state.socket, state.session_id, "USE #{space}",
               recv_timeout: recv_timeout
             ) do
          {:ok, %{error_code: 0}} -> {:ok, state}
          {:ok, %{error_code: code, error_msg: msg}} -> {:error, Error.from_response(code, msg)}
          {:error, reason} -> {:error, conn_error(reason, opts)}
        end
    end
  end

  defp maybe_use_space(_state, nil, _timeout), do: :ok

  defp maybe_use_space(state, space, recv_timeout) do
    case GraphService.execute(state.socket, state.session_id, "USE #{space}",
           recv_timeout: recv_timeout
         ) do
      {:ok, %{error_code: 0}} -> :ok
      {:ok, %{error_code: code, error_msg: msg}} -> {:error, Error.from_response(code, msg)}
      {:error, reason} -> {:error, reason}
    end
  end

  defp run_query(state, stmt, params, _prefer_json = false, recv_timeout) do
    if map_size(params) == 0 do
      GraphService.execute(state.socket, state.session_id, stmt, recv_timeout: recv_timeout)
    else
      GraphService.execute_with_parameter(state.socket, state.session_id, stmt, params,
        recv_timeout: recv_timeout
      )
    end
  end

  defp run_query(state, stmt, params, _prefer_json = true, recv_timeout) do
    if map_size(params) == 0 do
      case GraphService.execute_json(state.socket, state.session_id, stmt,
             recv_timeout: recv_timeout
           ) do
        {:ok, json} ->
          {:ok,
           %{
             error_code: 0,
             data: json,
             latency_in_us: 0,
             space_name: nil,
             error_msg: nil,
             comment: nil
           }}

        err ->
          err
      end
    else
      run_query(state, stmt, params, false, recv_timeout)
    end
  end

  defp conn_error(reason, opts) do
    if Keyword.get(opts, :show_sensitive_data_on_connection_error, false) do
      Error.connection_error(reason)
    else
      Error.connection_error(:redacted)
    end
  end
end