Building PostgreSQL Extensions: Dropping...

I recently created a Postgres extension which utilizes the pg_cron extension to schedule recurring activities using the cron.schedule(). Everything worked great. The only problem was when I dropped my extension, it left the cron job scheduled, which resulted in regular errors:

2024-04-06 16:00:00.026 EST [1548187] LOG:  cron job 2 starting: SELECT bridge_stats.update_stats('55 minutes', false)
2024-04-06 16:00:00.047 EST [1580698] ERROR:  schema "bridge_stats" does not exist at character 8
2024-04-06 16:00:00.047 EST [1580698] STATEMENT:  SELECT bridge_stats.update_stats('55 minutes', false)

If you look in the cron.job table, you can see the SQL for the cron job is still present, even though the extension/schema isn’t:

select schedule, command, jobname from cron.job;

schedule  |                        command                        |             jobname
-----------+-------------------------------------------------------+----------------------------------
0 0 * * 0 | SELECT bridge_stats.weekly_stats_update()             | bridge-stats-weekly-maintenance
0 * * * * | SELECT bridge_stats.update_stats('55 minutes', false) | bridge-stats-hourly-snapshot
(2 rows)

This got me thinking: how can you create a Postgres extension that can clean up after itself for cases like this?

How Extension Creation/Cleanup works

If you’ve created or used an extension in Postgres (such as pg_partman, PostGIS, pg_kaboom, etc) you may know that every extension in PostgreSQL has a SQL file that gets run as part of the creation.

This SQL file may create database objects for you, such as schemas, tables, functions, etc. When database objects are created in the context of a CREATE EXTENSION command, they have an object dependency created against the underlying pg_extension object. (These are stored in the pg_depend system catalog, if you are interested in the more fine-grained details.)

When Postgres removes an extension (via the DROP EXTENSION command), it will also remove any dependent objects that were created for this extension. (This is true for any dependencies, all of which are tracked in a similar way.)

This is how a simple command like DROP EXTENSION can remove dozens or hundreds of associated objects.

Why didn’t this cleanup?

You may be asking why this didn’t clean up the underlying cron jobs, since Postgres is clearly able to track the individual database objects associated with a given extension?

This is because the dependencies are tracked at the database object level (basically tracking the entries in the system tables that depend on each other). It is not general-purpose for cleanup.

So how to clean up?

We would like to be able to clean up these rows that were created by our extension. We don’t want to spam the user’s logs with unnecessary errors, particularly since we know exactly what we did to create the external rows.

In an ideal world, the extension itself could register a function that could be called when it’s being cleaned up. However, we do not live in an ideal world. (Not to mention there is probably a 125-email thread on the pgsql-hackers mailing list as to why that’s a bad idea; leaving finding that as an exercise to the reader…)

Since we don’t have that capacity, the general advice on the interwebs and in the Postgres docs is to use an EVENT TRIGGER.

Attempt 1: `CREATE EVENT TRIGGER`

Event triggers are a function that runs around special “events” that occur in a database. The current event trigger types are ddl_command_start, ddl_command_end, sql_drop, and rewrite_table. These let you take special action inside the database and run code when a given event occurs.

Since we are trying to run some code when this extension is dropped, clearly we want the sql_drop event trigger type.

Let’s take an initial stab at our cleanup function, created in our extension’s SQL file:

CREATE FUNCTION bridge_stats.cleanup() RETURNS event_trigger AS $$
DECLARE
    obj record;
BEGIN
    FOR obj IN SELECT * FROM pg_event_trigger_dropped_objects() LOOP
        IF obj.object_identity = 'bridge_stats' AND obj.object_type = 'extension' THEN
            PERFORM cron.unschedule('bridge-stats-weekly-maintenance');
            PERFORM cron.unschedule('bridge-stats-hourly-snapshot');
        END IF;
    END LOOP;
END;
$$ LANGUAGE plpgsql;

CREATE EVENT TRIGGER bridge_stats_cleanup ON sql_drop
WHEN TAG IN ('DROP EXTENSION')
EXECUTE FUNCTION bridge_stats.cleanup();

This seems like a straightforward attempt. We have created a function and an event trigger pair that end up being run any time a DROP EXTENSION is run. Our bridge_stats.cleanup() function in turn verifies that the extension itself is in the list of the dropped objects (returned by the pg_event_trigger_dropped_objects() function), and if it is, then we run the appropriate commands to unschedule our cron jobs. Easy-peasy.

Let’s go ahead and verify

“That was easy,” I say to myself, closing my text editor of choice (emacs, of course), and open my terminal to verify:

postgres=# create extension bridge_stats;
CREATE EXTENSION

postgres=# drop extension bridge_stats;
DROP EXTENSION

postgres=# select schedule, command, jobname from cron.job;

schedule  |                        command                        |             jobname
-----------+-------------------------------------------------------+----------------------------------
0 0 * * 0 | SELECT bridge_stats.weekly_stats_update()             | bridge-stats-weekly-maintenance
0 * * * * | SELECT bridge_stats.update_stats('55 minutes', false) | bridge-stats-hourly-snapshot
(2 rows)

The sweet smell of succ—oh wait. That didn’t work.

Adding logging (a la RAISE NOTICE 'BLARGH'), it appears that my event trigger was not even being called.

After considering a bit, it occurred to me that this wasn’t working because the event trigger must have been deleted as part of the extension’s schema, so it did not exist in the system when the sql_drop event trigger was called.

Perhaps the sql_drop event was run too late in the process? What about another one of the event trigger types?

Attempt 2: `CREATE EVENT TRIGGER 2: the what the heckening`

Looking at other options in the event trigger space, what are we left with?

ddl_command_start - run at the start of a DDL command
ddl_command_end - run at the end of a DDL command
rewrite_table - run when a table is rewritten

Clearly rewrite_table is off the, uh, err—you know—menu. Reading the docs for ddl_command_start and ddl_command_end shows that they are triggered before and after a DDL command is run.

“Ahh,” I exclaim, quickly transforming my existing event trigger into one based around the ddl_command_start event, since ddl_command_end runs after even sql_drop, so that one was out:

CREATE FUNCTION bridge_stats.cleanup() RETURNS event_trigger AS $$
DECLARE
    obj record;
BEGIN
    FOR obj IN SELECT * FROM pg_event_trigger_ddl_commands() LOOP
        IF obj.object_identity = 'bridge_stats' AND obj.object_type = 'extension' THEN
            PERFORM cron.unschedule('bridge-stats-weekly-maintenance');
            PERFORM cron.unschedule('bridge-stats-hourly-snapshot');
        END IF;
    END LOOP;
END;
$$ LANGUAGE plpgsql;

CREATE EVENT TRIGGER bridge_stats_cleanup ON ddl_command_start
WHEN TAG IN ('DROP EXTENSION')
EXECUTE FUNCTION bridge_stats.cleanup();

You can see that I’ve changed a couple of things relative to the previous version:

I am using pg_event_trigger_ddl_commands() instead of pg_event_trigger_dropped_objects(); simple API change for this specific filter.
I changed the ON action of the CREATE EVENT TRIGGER statement to be ddl_command_start

Verification, part deux

And now, on to verification:

postgres=# create extension bridge_stats;
CREATE EXTENSION

postgres=# drop extension bridge_stats;
ERROR:  pg_event_trigger_ddl_commands() can only be called in an event trigger function
CONTEXT:  PL/pgSQL function bridge_stats.cleanup() line 5 at FOR over SELECT rows

Queue the reaction gif where I am puzzled at the turn of events.

This function is now clearly getting called, since it is giving me an error message related to the specific function I’m calling. It is also clearly an event trigger, since it’s literally a function returning event_trigger and it’s been executed by the event trigger created by CREATE EVENT TRIGGER.

Well, for whatever reason, it would empirically appear that there is something odd going on with using ddl_start_command in this way; perhaps something with running this on a DROP command? In any case, rather than trying to debug this clearly odd behavior, I started thinking about a different approach.

Attempt 3: A Hero’s Journey

So if we recall my explanation about the dependencies inside Postgres and the objects created by extensions, we can see that the DROP EXTENSION was preemptively deleting my event trigger and the underlying function, meaning that it didn’t exist at the time the sql_drop event was issued. What if there was some way to somehow break that dependency so the event trigger would still exist to be fired, then it could clean itself up after it was done?

This lead me down the path to ALTER EXTENSION.

ALTER EXTENSION lets you dynamically add or remove dependencies between a specific extension and other database objects. While database objects created during CREATE EXTENSION are automatically associated with the creating extension, perhaps we could use this to our advantage.

With blazing eyes and a new tool in my hand, I made the following adjustments to my original attempt:

CREATE FUNCTION bridge_stats.cleanup() RETURNS event_trigger AS $$
DECLARE
    obj record;
BEGIN
    FOR obj IN SELECT * FROM pg_event_trigger_dropped_objects() LOOP
        IF obj.object_identity = 'bridge_stats' AND obj.object_type = 'extension' THEN
            PERFORM cron.unschedule('bridge-stats-weekly-maintenance');
            PERFORM cron.unschedule('bridge-stats-hourly-snapshot');
        END IF;
    END LOOP;
    DROP SCHEMA bridge_stats CASCADE;  -- the only new line in this function!
END;
$$ LANGUAGE plpgsql;

CREATE EVENT TRIGGER bridge_stats_cleanup ON sql_drop
WHEN TAG IN ('DROP EXTENSION')
EXECUTE FUNCTION bridge_stats.cleanup();

ALTER EXTENSION bridge_stats DROP EVENT TRIGGER bridge_stats_cleanup;
ALTER EXTENSION bridge_stats DROP FUNCTION bridge_stats.cleanup();
ALTER EXTENSION bridge_stats DROP SCHEMA bridge_stats;

As you can see, I have added the ALTER EXTENSION command to exclude the event trigger, the underlying function, and the owning schema from being owned by the extension.

I also added a DROP SCHEMA inside the cleanup() function to ensure that the objects that I manually detached from the extension’s schema would still get clean up.

Since everything else in the bridge_stats schema would get cleaned up by the DROP EXTENSION command, this would serve to finish the job, since a function can successfully delete itself in Postgres. (It’s true!)

Final verification

So of course, we need to verify that everything works as expected:

postgres=# create extension bridge_stats;
CREATE EXTENSION

postgres=# drop extension bridge_stats;
DROP EXTENSION

postgres=# select schedule, command, jobname from cron.job;
 schedule | command | jobname
----------+---------+---------
(0 rows)

Success!

TL;DR;

The top-down takeaway here is if you want to run some sort of cleanup action within a Postgres extension, you will have to:

Create your event trigger and associated function
ALTER EXTENSION DROP the event trigger, the function, and the schema
Ensure the cleanup function removes the objects you detached after doing whatever other cleanup job.

I hope that my experience of figuring out “just write an event trigger” helps someone else!

Latest Articles

Building PostgreSQL Extensions: Dropping Extensions and Cleanup