Writing a Postgres Extension With Pgrx for Visual Query Plans

When I worked at SingleStoreDB, we built this SQL syntax for instantly jumping into a visual query plan from the CLI:

sql> EXPLAIN UI select customer, sum(cost_cents) as revenue from notifications n
inner join offers o on o.offer_id = n.offer_id
group by ts, customer
order by revenue desc;

********************* 1 row *********************
EXPLAIN: https://...

This command would print out a URL in the shell that users could just click to hop into their web browser where the query plan would be displayed. This feature is still supported in SingleStore today.

Recently, I've been exploring the Postgres ecosystem and I missed this feature. So, I decided to build it with an extension. And I used pgrx, which is a framework for developing Postgres extensions in Rust.

Using pgrx

I followed this blog post which guides us through creating an extension using pgrx. I highly recommend it! However, it may be slightly out of date given the most recent pgrx udpates. So, I'll go through everything from scratch.

1. Start with cargo pgrx new pg_my_extension which creates a new directory called pg_my_extension/ with all the code you need.
2. The generated src/lib.rs includes all of the code. (You can ignore the src/bin/pgrx_embed.rs file)
3. To test the extension, run cargo pgrx new pg17. This starts a new Postgres 17 instance and you can then run create extension pg_my_extension in order to load the extension and be able to play with it locally.
4. As for how to write the extension's actual code, again, I'll just refer to this blog post which does a great job of explaining how Rust functions will be mapped to Postgres functions. But I promise it's really easy, you can also follow along a very simple example such as this extension.

Implementing pg_explain_ui

To build pg_explain_ui, I had to issue a POST request to Dalibo's API in order to create a new plan file which could be browsed online. For that, I used Rust's reqwest:

// Construct URL-encoded form data manually
let form_data = format!(
    "query={}&title={}&plan={}",
    urlencoding::encode(&query_text),
    urlencoding::encode(""),
    urlencoding::encode(&plan_json_str)
);

// Create a client
let client = Client::new();

// Perform the POST request
let res = client
    .post("https://explain.dalibo.com/new")
    .header("Content-Type", "application/x-www-form-urlencoded")
    .body(form_data)
.send()?;

This is extremely simple! There's really not that much more code to this.

Using the extension

Making use of this couldn't be easier. Just create a schema:

psql> CREATE TABLE authors (
    author_id SERIAL PRIMARY KEY,
    name VARCHAR(255) NOT NULL
);

CREATE TABLE publishers (
    publisher_id SERIAL PRIMARY KEY,
    name VARCHAR(255) NOT NULL
);

CREATE TABLE categories (
    category_id SERIAL PRIMARY KEY,
    name VARCHAR(255) NOT NULL
);

CREATE TABLE books (
    book_id SERIAL PRIMARY KEY,
    title VARCHAR(255) NOT NULL,
    author_id INT NOT NULL REFERENCES authors(author_id),
    publisher_id INT NOT NULL REFERENCES publishers(publisher_id),
    publication_date DATE
);

CREATE TABLE book_categories (
    book_id INT NOT NULL REFERENCES books(book_id),
    category_id INT NOT NULL REFERENCES categories(category_id),
    PRIMARY KEY (book_id, category_id)
);

And then, call SELECT explain_ui(...):

psql> SELECT explain_ui($$SELECT
    b.title AS "Book Title",
    a.name AS "Author",
    p.name AS "Publisher",
    array_agg(c.name) AS "Categories",
    b.publication_date AS "Publication Date",
    COUNT(c.category_id) AS "Number of Categories"
FROM
    books b
INNER JOIN authors a ON b.author_id = a.author_id
INNER JOIN publishers p ON b.publisher_id = p.publisher_id
INNER JOIN book_categories bc ON b.book_id = bc.book_id
INNER JOIN categories c ON bc.category_id = c.category_id
WHERE
    p.name != 'Test'
GROUP BY
    b.book_id, a.name, p.name, b.publication_date
HAVING
    COUNT(c.category_id) > 1
ORDER BY
    b.publication_date DESC;
$$);
                    explain_ui
--------------------------------------------------
 https://explain.dalibo.com/plan/ccg2e5fedd913bb7
(1 row)

And that's it! Now, you can click the link and analyze the query plan visually:

I will need to figure out how to get this extension installed across a bunch of managed Postgres providers (Neon, Supabase, RDS, Crunchy Data, Azure Postgres, etc.). However, I think that will only really possible if the extension becomes popular enough.

Feel free to reach out on Twitter!