Using Triggers to Map Database Relationships in Custom Claims

Custom claims are a powerful tool for implementing row-level security (RLS) policies in your applications which was popularized by Supabase. In a previous blog post (Part 1: "Using Custom Claims: Testing RLS with Supabase"), we explored the basics of custom claims and their application in Supabase. In this Part 2, we will delve deeper into the topic by introducing triggers and how they can be used to map database relationships in custom claims.

Mapping relationships between database tables is not only essential for enforcing access control but also plays a crucial role in managing data visibility. Traditionally, RLS policies would require multiple joins to determine the relationships between entities and apply appropriate access restrictions. However, with triggers and custom claims, we can achieve a more performant solution.

By leveraging triggers in PostgreSQL, we can automate the process of updating custom claims whenever there are changes in the relationships between entities. This means that instead of making multiple joins every time access control is enforced, the relationships are stored within the custom claims themselves. This approach significantly improves the efficiency and performance of your application's access control implementation.

In this blog post, we will walk you through the creation and implementation of a trigger function that updates custom claims based on changes in a relationship table. By utilizing triggers, you can simplify the management of claims and enhance the security of your Supabase-powered applications while maintaining optimal performance.

Let's explore how triggers can revolutionize the way you handle claims and take your application's security to the next level.

Mapping Relationships with Triggers

To demonstrate the concept, we'll use a relationship table called team_members. This table connects users and teams, allowing us to define which users belong to which teams. Here's the structure of the team_members table:

-- Table that we will be mapping on claims
create table
  public.team_members (
    id bigint generated by default as identity not null,
    created_at timestamp with time zone null default now(),
    user_id uuid null,
    team_id integer null,
    constraint profiles_pkey primary key (id),
    constraint team_members_team_id_fkey foreign key (team_id) references teams (id),
    constraint team_members_user_id_fkey foreign key (user_id) references auth.users (id)
  );

Now, let's dive into the trigger function responsible for updating the custom claims whenever changes occur in the team_members table.

Full Trigger Function for Updating Custom Claims

In this section, we will explore the complete trigger function responsible for updating the custom claims when changes occur in the team_members table. The trigger function plays a crucial role in automating the process of managing and maintaining accurate custom claims based on the relationships between entities. By understanding the inner workings of this function, you'll gain valuable insights into how Supabase enables seamless and efficient implementation of row-level security policies. Then, we will delve into a detailed explanation of each component of the trigger function. But first, let's take a look at the full code of the function itself:

CREATE FUNCTION team_member_claim()
RETURNS TRIGGER
LANGUAGE PLPGSQL
AS $$
DECLARE
  -- This function works to map if team_id was UUID
  old_team_id INTEGER; 
  new_team_id INTEGER;
  affected_user_id UUID;
  current_claim jsonb;
  updated_claim jsonb;
BEGIN
  -- Retrieve the relevant IDs based on the operation type
  IF TG_OP = 'DELETE' THEN
    old_team_id := OLD.team_id;
    affected_user_id := OLD.user_id;
  ELSIF TG_OP = 'UPDATE' THEN
    old_team_id := OLD.team_id;
    new_team_id := NEW.team_id;
    affected_user_id := NEW.user_id;
  ELSE
    new_team_id := NEW.team_id;
    affected_user_id := NEW.user_id;
  END IF;
  -- Retrieve the current custom claim for the affected user
  current_claim := get_claim(affected_user_id, 'team_member');
  -- Update the claim based on the operation type
  IF TG_OP = 'UPDATE' THEN
    -- Remove the old_team_id from the claim
    SELECT jsonb_agg(elem)
    INTO updated_claim
    FROM jsonb_array_elements(current_claim) elem
    WHERE elem::integer <> old_team_id;
    -- Add the new_team_id to the claim
    updated_claim := updated_claim || jsonb_build_array(new_team_id);
    -- Update the claim
    PERFORM set_claim(affected_user_id, 'team_member', updated_claim);
    RAISE NOTICE 'Set claim (update team): %', updated_claim;
  ELSE
    IF current_claim IS NULL THEN
      IF TG_OP <> 'DELETE' THEN
      -- Set the claim with a new_team_id when it was previously empty
        PERFORM set_claim(affected_user_id, 'team_member', jsonb_build_array(new_team_id));
        RAISE NOTICE 'Set claim (null): %', jsonb_build_array(new_team_id);
      END IF;
    ELSIF current_claim @> jsonb_build_array(old_team_id) THEN
      IF TG_OP = 'DELETE' THEN
        -- Remove the old_team_id from the claim
        SELECT jsonb_agg(elem)
        INTO updated_claim
        FROM jsonb_array_elements(current_claim) elem
        WHERE elem::integer <> old_team_id;

        -- Update the claim
        PERFORM set_claim(affected_user_id, 'team_member', updated_claim);
        RAISE NOTICE 'Set claim (remove team): %', updated_claim;
      END IF;
    ELSE
      IF TG_OP <> 'DELETE' THEN
        -- Add the new_team_id to the claim
        PERFORM set_claim(affected_user_id, 'team_member', current_claim || jsonb_build_array(new_team_id));
        RAISE NOTICE 'Set claim (add team): %', current_claim || jsonb_build_array(new_team_id);
      END IF;
    END IF;
  END IF;
  -- Return the appropriate value based on the operation type
  IF TG_OP = 'DELETE' THEN
    RETURN OLD;
  ELSE
    RETURN NEW;
  END IF;
END;
$$;

Explanation of the Trigger Function

Let's break down the logic behind the trigger function:

The function is defined as a trigger function, indicated by the RETURNS TRIGGER statement.
Inside the function, we declare variables to store the old and new team IDs, the affected user ID, the current custom claim, and the updated custom claim.
Depending on the operation type (DELETE, UPDATE, or INSERT), we assign values to the variables accordingly.
We retrieve the current custom claim for the affected user using the get_claim function.
Based on the operation type, we modify the claim:
- For an UPDATE operation, we remove the old_team_id from the claim and add the new_team_id to it.
- If the current claim is NULL, we set the claim with the new_team_id.
- If the claim already contains the old_team_id, we remove it.
- If none of the above conditions apply, we add the new_team_id to the claim.
We update the claim using the set_claim function and raise a notice to log the updated claim.
Finally, we return the appropriate value based on the operation type.

Deploying the trigger to leverage this function

We can create a single trigger that works on all operations and use it to manage the table relationships in our claims.

CREATE TRIGGER team_member_claim_trigger
AFTER INSERT OR UPDATE OR DELETE ON team_members
FOR EACH ROW
EXECUTE FUNCTION team_member_claim();

This huge trigger function might feel overwhelming. It is possible and fine if you split it into different functions to simplify the logic. Here's an example of splitting the logic across different trigger functions and creating the triggers for each case.

Insert trigger function:

CREATE FUNCTION team_member_claim_insert()
RETURNS TRIGGER
LANGUAGE PLPGSQL
AS $$
DECLARE
  new_team_id INTEGER := NEW.team_id;
  affected_user_id UUID := NEW.user_id;
  current_claim jsonb := get_claim(affected_user_id, 'team_member');
BEGIN
  IF current_claim IS NULL THEN
    PERFORM set_claim(affected_user_id, 'team_member', jsonb_build_array(new_team_id));
    RAISE NOTICE 'Set claim (null): %', jsonb_build_array(new_team_id);
  ELSE
    PERFORM set_claim(affected_user_id, 'team_member', current_claim || jsonb_build_array(new_team_id));
    RAISE NOTICE 'Set claim (add team): %', current_claim || jsonb_build_array(new_team_id);
  END IF;
  RETURN NEW;
END;
$$;

Delete trigger function:

CREATE FUNCTION team_member_claim_delete()
RETURNS TRIGGER
LANGUAGE PLPGSQL
AS $$
DECLARE
  old_team_id INTEGER := OLD.team_id;
  affected_user_id UUID := OLD.user_id;
  current_claim jsonb := get_claim(affected_user_id, 'team_member');
  updated_claim jsonb;
BEGIN
  IF current_claim @> jsonb_build_array(old_team_id) THEN
    SELECT jsonb_agg(elem)
    INTO updated_claim
    FROM jsonb_array_elements(current_claim) elem
    WHERE elem::integer <> old_team_id;
    PERFORM set_claim(affected_user_id, 'team_member', updated_claim);
    RAISE NOTICE 'Set claim (remove team): %', updated_claim;
  END IF;
  RETURN OLD;
END;
$$;

Update trigger function

CREATE FUNCTION team_member_claim_update()
RETURNS TRIGGER
LANGUAGE PLPGSQL
AS $$
DECLARE
  old_team_id INTEGER := OLD.team_id;
  new_team_id INTEGER := NEW.team_id;
  affected_user_id UUID := NEW.user_id;
  current_claim jsonb := get_claim(affected_user_id, 'team_member');
  updated_claim jsonb;
BEGIN
  SELECT jsonb_agg(elem)
  INTO updated_claim
  FROM jsonb_array_elements(current_claim) elem
  WHERE elem::integer <> old_team_id;
  updated_claim := updated_claim || jsonb_build_array(new_team_id);
  PERFORM set_claim(affected_user_id, 'team_member', updated_claim);
  RAISE NOTICE 'Set claim (update team): %', updated_claim;
  RETURN NEW;
END;
$$;

Creating individual triggers

CREATE TRIGGER team_member_claim_insert_trigger
AFTER INSERT ON team_member
FOR EACH ROW EXECUTE PROCEDURE team_member_claim_insert();

CREATE TRIGGER team_member_claim_delete_trigger
AFTER DELETE ON team_member
FOR EACH ROW EXECUTE PROCEDURE team_member_claim_delete();

CREATE TRIGGER team_member_claim_update_trigger
AFTER UPDATE ON team_member
FOR EACH ROW EXECUTE PROCEDURE team_member_claim_update();

Conclusion

In this blog post, we explored how triggers can be used to map database relationships in custom claims. By automating the process of updating custom claims when changes occur in the relationship table, we can enhance the efficiency and security of our applications. The trigger function we created allows for the seamless management of claims and ensures that access control remains accurate and up-to-date.

By leveraging triggers and custom claims, you can build robust row-level security mechanisms in your Supabase-powered applications, providing a granular level of access control to your users. Having the claims is very helpful when building policies for storage buckets.

Stay tuned for more exciting topics on leveraging Supabase and custom claims to build powerful and secure applications.

Blog