supabase_postgrest

alias Ecto.Adapters.SQL, as: EctoSQL

if Code.ensure_loaded?(EctoSQL) do
  defmodule Supabase.PostgREST.RLS do
    @moduledoc """
    Helpers for working with PostgreSQL Row Level Security (RLS) policies in Ecto.

    This module provides utilities to set session context for RLS policies when using
    the Ecto schemas generated by `mix supabase.gen.schema`.

    ## Usage

    Wrap operations in a transaction and set session variables that are automatically
    cleaned up when the transaction completes:

        defmodule MyApp.Repo do
          use Ecto.Repo, otp_app: :my_app

          import Supabase.PostgREST.RLS

          def with_user_context(user_id, fun) do
            transaction(fn ->
              set_rls_context(__MODULE__, "request.jwt.claims.sub", user_id)
              fun.()
            end)
          end
        end

        # Usage
        MyApp.Repo.with_user_context(user_id, fn ->
          MyApp.Repo.all(MyApp.Accounts.User)
        end)

    ## Common Supabase RLS Patterns

    Supabase RLS policies commonly use these session variables:

    - `request.jwt.claims.sub` - User ID from JWT
    - `request.jwt.claims.role` - User role from JWT
    - `request.jwt.claims.email` - User email from JWT

    Example policy in your database:

        CREATE POLICY "Users can view their own data"
          ON users FOR SELECT
          USING (id = current_setting('request.jwt.claims.sub')::uuid);

    ## Testing with RLS

    When testing, you may want to bypass RLS policies. You can disable RLS for specific tables:

        Ecto.Adapters.SQL.query!(MyApp.Repo, "ALTER TABLE users DISABLE ROW LEVEL SECURITY")
    """

    @doc """
    Sets a PostgreSQL session configuration variable for the current transaction.

    This function executes `set_config/3` with the `is_local` parameter set to `true`,
    ensuring the configuration is transaction-scoped and automatically cleared when
    the transaction completes.

    ## Parameters

      * `repo` - The Ecto repository module
      * `key` - The configuration parameter name (e.g., "request.jwt.claims.sub")
      * `value` - The value to set (will be converted to string)

    ## Examples

        # Setting user ID for RLS policies
        Repo.transaction(fn ->
          set_rls_context(repo, "request.jwt.claims.sub", user.id)
          Repo.all(User)
        end)

        # Setting user role
        Repo.transaction(fn ->
          set_rls_context(repo, "request.jwt.claims.role", "authenticated")
          Repo.all(User)
        end)

        # Setting user email
        Repo.transaction(fn ->
          set_rls_context(repo, "request.jwt.claims.email", user.email)
          Repo.all(User)
        end)

        # With custom variable name for multi-tenancy
        Repo.transaction(fn ->
          set_rls_context(repo, "app.current_tenant", tenant_id)
          Repo.all(from u in User, where: ...)
        end)

    ## Security Considerations

    - Always validate and sanitize the `value` parameter before passing it to this function
    - Never pass user input directly without validation
    - Consider using prepared statements for complex values

    """
    def set_rls_context(repo, key, value) when is_binary(key) do
      value_str = to_string(value)

      EctoSQL.query!(
        repo,
        "SELECT set_config($1, $2, true)",
        [key, value_str]
      )
    end

    @doc """
    Sets multiple RLS context variables at once in a single database query.

    This is more efficient than calling `set_rls_context/3` multiple times when you need
    to set several session variables.

    ## Parameters

      * `repo` - The Ecto repository module
      * `claims` - A map of configuration keys to values

    ## Examples

        # Set multiple JWT claims at once
        Repo.transaction(fn ->
          set_rls_claims(repo, %{
            "request.jwt.claims.sub" => user.id,
            "request.jwt.claims.role" => "authenticated",
            "request.jwt.claims.email" => user.email
          })
          Repo.all(User)
        end)

    """
    def set_rls_claims(repo, claims) when is_map(claims) and map_size(claims) > 0 do
      {keys, values} = Enum.unzip(Map.to_list(claims))

      set_config_calls =
        keys
        |> Enum.with_index(1)
        |> Enum.map_join(", ", fn {_key, idx} ->
          key_param = idx * 2 - 1
          val_param = idx * 2
          "set_config($#{key_param}, $#{val_param}, true)"
        end)

      query = "SELECT " <> set_config_calls

      params =
        Enum.zip(keys, Enum.map(values, &to_string/1))
        |> Enum.flat_map(&Tuple.to_list/1)

      EctoSQL.query!(repo, query, params)
    end

    @doc """
    Gets the current value of an RLS context variable.

    Returns `nil` if the variable is not set.

    ## Examples

        current_user_id = get_rls_context(MyApp.Repo, "request.jwt.claims.sub")

    """
    def get_rls_context(repo, key) when is_binary(key) do
      case EctoSQL.query(repo, "SELECT current_setting($1, true)", [key]) do
        {:ok, %{rows: [[value]]}} when value != "" -> value
        _ -> nil
      end
    end
  end
end