A few months ago, we decided to build a referral program for Mergify. This is a well-known, classical way of bringing more people on your product.
To build this program, we add to define a data model that allowed us to store a list of referrers and referees, with some constraints. Such a list boils down to building a direct acyclic graph (DAG) of your users. A DAG's particularity is to be a directed tree that, therefore, cannot have any loop in it.
There are multiple solutions to detect a loop in a DAG. The straightforward method is to use a depth-first search algorithm: start from a node, keep browsing the tree from node to node, writing down each node encountered. If one of the discovered nodes is already in your stack, then you have a loop.
It is easy enough to find plenty of solutions around this issue on the web, depending on your data structure and programming language.
In our case, we knew we were going to store this data in PostgreSQL, meaning we needed to make sure the database did the check directly on insertion.
To make this work in PostgreSQL, we needed to split the problem in three parts:
- The data structure;
- Write a function that can detect if a cycle is present in your data structure;
- Create a trigger that calls the loop detection function on insertion and update.
We went ahead and implemented a straightforward data structure using a single table containing 2 IDs — the vertices of our DAG. Hence, each row in the table stores the link between our vertices (the referrer and its referee).
The PostgreSQL table can be created with this command:
CREATE TABLE public.referral ( referee_id integer NOT NULL PRIMARY KEY, referrer_id integer NOT NULL, CONSTRAINT check_referrer_not_referee CHECK ((referrer_id <> referee_id)) );
This table already has a simple built-in loop detection: the
check_referrer_not_referee check makes sure that a user cannot refer itself, which, in tree jargon, mean no vertex can link back to itself.
This check is not enough: while A can't link to A, A could link to B which could link back to A, creating the loop A -> B -> A. Not great.
Another case that is prevented with this table is that since the
referee_id is a primary key, it can't be used multiple times, avoiding the scenario where a referee would have multiple referrer.
Loop Detection Function
The initial table creation solved simple cases where the links between vertices are easy to check. We now need to write a function that walks the entire tree when inserting links in the general case.
This is where you need a PostgreSQL function to detect the loop. The PL/pgSQL language is powerful enough to write such a function:
CREATE OR REPLACE FUNCTION detect_cycle() RETURNS TRIGGER LANGUAGE plpgsql AS $func$ BEGIN IF EXISTS ( WITH RECURSIVE list_referrers(referrer) AS ( -- Get the parent from the NEW referrer referral to see -- if the referrer is not already a parent SELECT r.referrer_id FROM referral AS r WHERE r.referee_id = NEW.referrer_id UNION ALL SELECT r.referrer_id FROM referral AS r, list_referrers as lr WHERE r.referee_id = lr.referrer ) SELECT * FROM list_referrers WHERE list_referrers.referrer = NEW.referee_id LIMIT 1 ) THEN RAISE EXCEPTION 'Loop detected'; ELSE RETURN NEW; END IF; END $func$;
The function works in a pretty simple way. It will be used as an
UPDATE trigger, which means it can use the special
NEW table to get access to the newly inserted row.
Then, it uses the
WITH RECURSIVE clause to run a recursive SQL query that can refer to its own output. In that case, PostgreSQL first runs:
SELECT r.referrer_id FROM referral AS r WHERE r.referee_id = NEW.referrer_id
This first query returns the parent, i.e., the vertex that links back to the newly referee (child) node. If the referrer has no parent, it returns nothing, and the query finishes with 0 results.
Otherwise, the second query uses this list to run that same query over and over again:
SELECT r.referrer_id FROM referral AS r, list_referrers as lr WHERE r.referee_id = lr.referrer
This selects the parent of the parent (if any) and continues to do so until there is no more new result. At this stage, all the results are combined together and sent to the new query:
SELECT * FROM list_referrers WHERE list_referrers.referrer = NEW.referee_id LIMIT 1
This query selects the row in the previous list of parents where the parent is the referee. If such a result exists, then it means a loop exists, as the newly added referee is in the referrer's parent list. The
LIMIT 1 statement is an optimization to indicate that one result is enough since our outer check is
IF EXISTS / THEN / RAISE statement takes care of raising an exception if a loop is detected, preventing the new row from being added.
Setting up the trigger
Once you get this function right for you, there's only a need to add it to your table with this:
DROP TRIGGER IF EXISTS detect_referral_cycle ON referral; CREATE CONSTRAINT TRIGGER detect_referral_cycle AFTER INSERT OR UPDATE ON referral FOR EACH ROW EXECUTE PROCEDURE detect_cycle();
Pretty easy! With all of that setup, there is no need to do any kind of check in your application, and you can simply insert the rows as you wish. Catching the PostgreSQL exception that can be raised by the trigger allows you to report back to the user that they're trying to do something invalid.
PostgreSQL 14 and native cycle detection
Starting with version 14, PostgreSQL added a feature that adds the
CYCLE clauses to recursive queries to be able to do produce breadth- or depth-first search orders and detect cycles. Hubert Lubaczewski has a great blog post that shows a few examples on how to use those features, which should make writing such loop checks even easier!