Detecting queries in Django tests

Putting this here so I can find it next time I need to know it…

Django has a useful test assertion you can use to ensure you make a set number of queries. However, at times this is a bit less useful than it needs to be, because something changes and we do indeed have a different number of queries, but it’s got nothing to do with the actual code under test.

If you are running in DEBUG=True mode, then you can examine the queries that have been made to the database connection, and ensure the raw SQL of a specific query matches (and is not duplicated, for instance).

This does require a little bit of trickery:

from django.db import connection
from django.test import TestCase, override_settings

from foo.factories import FooFactory


class TestTheThing(TestCase):
    def test_no_update_query(self):
        foo = FooFactory()

        # Our Foo instance should be smart enough to notice that nothing
        # has changed, and thus should not emit an UPDATE query.
        with override_settings(DEBUG=True):
            foo.save()
            self.assertFalse([
              x
              for x in connection.queries
              if 'UPDATE' in x['sql']
            ])

This is a bit of a contrived case: in my case today it was a celery task that only updated an object if the incoming data differed to the saved data, so there was by necessity a SELECT query to get the object, and then an UPDATE only if the data had changed.

The trick is that we use a context manager to set DEBUG to true, this means that it will start capturing queries, and drop them out at the end of the context. This also means you can have a bunch of these in the same test case, and each one will have an independent set of queries.

Partial, Deferrable Unique Constraints

I had a bug I came across today that was easiest to reproduce by building up a test case that uses a Factory from factory-boy, that creates a set of related objects.

In order to reproduce the failure conditions, I needed to increment a value from a specific column in each of these. This column is a member of a multi-column UNIQUE constraint:

class MyModel(models.Model):
    team = models.ForeignKey(Team)
    day_part = models.ForeignKey(DayPart, null=True)
    level = models.IntegerField()

    class Meta:
        unique_together = (
            ('team', 'day_part', 'level'),
        )

By default, Django creates UNIQUE constraints as immediate, and not deferable. There is currently no way to change this. This prevents you from doing something like:

MyModel.objects.filter(team=X, day_part=Y).update(level=models.F('level') + 1)

The query it generates should be valid, and should actually work, because it all happens in one transaction. But because the constraints are not deferred, then instead you need to do:

for instance in MyModel.objects.filter(...).order_by('-level'):
    instance.level += 1
    instance.level.save()

Which results in a number of queries twice the size of the number of objects, plus one.

Instead, we want to defer the constraint, and make it initially deferred. The name of the index will be generated by Django, so you’ll need to look in the database, and then create a migration accordingly:

ALTER TABLE myapp_mymodel
DROP CONSTRAINT myapp_mymodel_team_id_day_part_id_level_aaaaaa_uniq;

ALTER TABLE myapp_mymodel
ADD CONSTRAINT myapp_mymodel_team_id_day_part_id_level_aaaaaa_uniq
UNIQUE (team_id, day_part_id, level)
DEFERRABLE INITIALLY DEFERRED;

Because the foreign key to DayPart is nullable, it means that any records with a NULL in this field will not be subject to this constraint. In this case, that is the intention, but the requirements for this model also say that for a given team, there may not be any duplicate levels where there is no linked DayPart. I’d initially modelled this as:

CREATE UNIQUE INDEX myapp_mymodel_no_day_part
ON myapp_mymodel(team_id, level)
WHERE day_part_id IS NULL

But there is no way to apply this index as a deferred constraint. Instead, you need to use an EXCLUDE constraint:

ALTER TABLE myapp_mymodel
ADD CONSTRAINT myapp_mymodel_no_day_part
EXCLUDE USING btree(team_id WITH =, level WITH =)
WHERE (day_part_id IS NULL)
DEFERRABLE INITIALLY DEFERRED

Apple Watch Stand Hours

As of watchOS 6/iOS 13, I’ve noticed some inconsistencies between stand hours and stand minutes in Health/Activity.

Specifically, there are often hours where, inside Health.app, I can see that I have at least one stand minute, but when I view in Activity or Health Stand Hours, there is no stand hour.

This morning, I think I figured it out.

I got up, as I normally do, at 6:55. This is my life hack to ensure I get a stand hour at 6-7 am. I went about my normal routine, making sure I moved enough before 7am to trigger the stand hour. It did.

At around 7:15, I happened to look at my watch again. I still only had one stand hour. I know I had moved around in the previous 15 minutes, so I went into stand minutes. Yep, there was a stand minute after 7am.

At about 7:35, I glanced at my watch again, and noticed there was a second stand hour.

I live in a place that has a non-integer number of hours offset from UTC (+09:30, +10:30 DST). I believe that stand hours are calculated based on UTC, and then offset for the current timezone, wheras it used to actually use the hour within the timezone.

My workaround is to make sure I have a stand minute within both the first and second half-hours of every hour (which can be helped by having half-hourly chimes enabled on the watch).

Speeding up Postgres using a RAM disk?

I happened to come across Speeding up Django unit tests with SQLite, keepdb and /dev/shm today. I can’t quite do that, because we use a bunch of Postgres specific things in our project, but it did make me think “Can we run Postgres on a RAM disk?”

Following up on this, it turns out that using a TABLESPACE in postgres that is on a RAM disk is a really bad idea, but it should be possible to run a seperate database cluster (on a different port) that could have the whole cluster on the database.

It’s possible to create a RAM disk on macOS, and init a cluster there:

#! /bin/sh

PG_CTL=/Applications/Postgres.app/Contents/Versions/10/bin/pg_ctl

new_disk=$(hdid -nomount ram://1280000)
newfs_hfs $new_disk
mkdir /tmp/ramdisk
mount -t hfs $new_disk /tmp/ramdisk
mkdir /tmp/ramdisk/pg

PG_CTL initdb -D /tmp/ramdisk/pg

This needs to be run once, after bootup (and probably before Postgres.app starts it’s servers). You’d also need to do the setup in Postgres.app to show where it is (alternatively, you could just start it on the port you want in the script above).

But, is it faster? Let’s look at the results. The “Initial run” is the first run of tests after deleting the test database.

Database location Initial run Keepdb runs
Disk 6m39s 0m16s
0m19s
0m15s
RAM 6m11s 0m26s
0m17s
0m15s

So, it’s probably not actually worth it. I’m guessing that the variance is more to do with just how busy my machine was at each of the times I ran the test.

Postgres ENUM types in Django

Postgres has the ability to create custom types. There are several kinds of CREATE TYPE statement:

  • composite types
  • domain types
  • range types
  • base types
  • enumerated types

I’ve used a metaclass that is based on Django’s Model classes to do Composite Types in the past, and it’s been working fairly well. The current stuff I have been working on made sense to use an Enumerated Type, because there are four possible values, and having a human readable version of them is going to be nicer than using a lookup table.

In the first iteration, I used just a TEXT column to store the data. However, when I then started to use an enum.Enum class for handling the values in python, I discovered that it was actually storing str(value) in the database, rather than value.value.

So, I thought I would implement something similar to my Composite Type class. Not long after starting, I realised that I could make a cleaner implementation (and easier to declare) using a decorator:

@register_enum(db_type='change_type')
class ChangeType(enum.Enum):
    ADDED = 'added'
    CHANGED = 'changed'
    REMOVED = 'removed'
    CANCELLED = 'cancelled'


ChangeType.choices = [
    (ChangeType.ADDED, _('hours added')),
    (ChangeType.REMOVED, _('hours subtracted')),
    (ChangeType.CHANGED, _('start/finish changed with no loss of hours')),
    (ChangeType.CANCELLED, _('shift cancelled')),
]

Because I’m still on an older version of Python/Django, I could not use the brand new Enumeration types, so in order to make things a bit easier, I then annotate onto the class some extra helpers. It’s important to do this after declaring the class, because otherwise the attributes you define will become “members” of the enumeration. When I move to Django 3.0, I’ll probably try to update this register_enum decorator to work with those classes.

So, let’s get down to business with the decorator. I spent quite some time trying to get it to work using wrapt, before realising that I didn’t actually need to use it. In this case, the decorator is only valid for decorating classes, and we just add things onto the class (and register some things), so it can just return the new class, rather than having to muck around with docstrings and names.

from psycopg2.extensions import (
    new_array_type,
    new_type,
    QuotedString,
    register_adapter,
    register_type,
)
known_types = set()


CREATE_TYPE = 'CREATE TYPE {0} AS ENUM ({1})'
SELECT_OIDS = 'SELECT %s::regtype::oid AS "oid", %s::regtype::oid AS "array_oid"'


class register_enum(object):
    def __init__(self, db_type, managed=True):
        self.db_type = db_type
        self.array_type = '{}[]'.format(db_type)
        self.managed = managed

    def __call__(self, cls):
        # Tell psycopg2 how to turn values of this class into db-ready values.
        register_adapter(cls, lambda value: QuotedString(value.value))

        # Store a reference to this instance's "register" method, which allows
        # us to do the magic to turn database values into this enum type.
        known_types.add(self.register)

        self.values = [
            member.value
            for member in cls.__members__.values()
        ]

        # We need to keep a reference to the new class around, so we can use it later.
        self.cls = cls

        return cls

    def register(self, connection):
        with connection.cursor() as cursor:
            try:
                cursor.execute(SELECT_OIDS, [self.db_type, self.array_type])
                oid, array_oid = cursor.fetchone()
            except ProgrammingError:
                if self.managed:
                    cursor.execute(self.create_enum(connection), self.values)
                else:
                    return

        custom_type = new_type(
            (oid,),
            self.db_type,
            lambda data, cursor: data and self.cls(data) or None
        )
        custom_array = new_array_type(
            (array_oid,),
            self.array_type,
            custom_type
        )
        register_type(custom_type, cursor.connection)
        register_type(custom_array, cursor.connection)

    def create_enum(self, connection):
        qn = connection.ops.quote_name
        return CREATE_TYPE.format(
            qn(self.db_type),
            ', '.join(['%s' for value in self.values])
        )

I’ve extracted out the create_enum method, because it’s then possible to use this in a migration (but I’m not totally happy with the code that generates this migration operation just yet). I also have other code that dynamically creates classes for a ModelField and FormField as attributes on the Enum subclass, but that complicates it a bunch.

Testing WizardView in Django

Sometimes, you just have to resort to a multi-step view to simplify data entry for a user. Django had an “included” package for this: django.contrib.formtools, until some years ago, when it moved into it’s own repository. It’s still very useful though, although there is a bit of a learning curve.

I have some multi-step wizard views, for things where there are a bunch of different (sometimes dependent, but not necessarily) things we need to get from the user, and commit in an all-or-nothing way.

Writing tests for these views is quite cumbersome, but I’ve recently come up with a couple of things that can make this less painful. The reason it’s somewhat difficult is that the wizard, and the current step form both have form data, so to ensure there are no clashes between name fields, the forms are all prefixed. However, prefixes on forms mean that the data you need to pass to the form must also be prefixed in the correct way.

class TestWizard(TestCase):
    def test_wizard_view(self):
        # user has been created...

        self.client.force_login(user)

        # url is the url for this wizard view.

        response = self.client.get(url)
        self.assertEqual(200, response.status_code)

        response = self.client.post(url, {
            'wizard-current_step': 'first',
            'first-field_a': 'value',
            'first-field_b': 'value',
            'first-field_c': 'value',
        })
        self.assertEqual(200, response.status_code)
        # Because django by default returns a 200 on a form validation error,
        # we need another way to check this. If this is a re-rendering of the
        # same form it will be bound, otherwise it will not be.
        self.assertFalse(response.context['form'].is_bound)
        self.assertEqual(
            response.context['wizard']['management_form']['current_step'].value(),
            'second'
        )

        # This next one has a formset instead of a form!
        response = self.client.post(url, {
          'wizard-current_step': 'second',
          'second-0-field_a': 'a',
          'second-0-field_b': 'a',
          'second-1-field_a': 'a',
          'second-1-field_b': 'a',
          'second-INITIAL_FORMS': '0',
          'second-TOTAL_FORMS': '2'
        })
        self.assertEqual(200, response.status_code)
        self.assertFalse(response.context['form'].is_bound)
        # Okay, I'm bored with this already!

It would be much nicer if we could just do something like:

class TestWizard(TestCase):
    def test_wizard_view(self):
        # ...
        self.client.force_login(user)

        test_wizard(self, url, 'wizard', [
          ('first', {'field_a': 'value', 'field_b': 'value', ...}),
          ('second', [
            {'field_a': 'a', 'field_b': 'b'},
            {'field_a', 'a', 'field_b': 'a'},
          ])
        ])

The last part, the formset handling stuff is something I have touched on before, and I’m not sure if I’ll get to that bit now, but the first part is still useful.

def test_wizard(test, url, name, data):
    """
    Using the supplied TestCase, execute the wizard view installed at
    the "url", having the given name, with each item of data.
    """

    test.assertEqual(200, test.client.get(url).status_code)
    for step, step_data in data:
      step_data = {
        '{}-{}'.format(step, key): value
        for key, value in step_data.items()
      }
      step_data['{}-current_step'.format(name)] = step
      response = test.client.post(url, step_data, follow=True)
      test.assertEqual(200, response.status_code)
      if 'form' in response.context:
          test.assertFalse(response.context['form'].errors)

We might come back to the formset thing later.


But what about when we have a bunch of steps where we already have initial data, and we mostly just want to click next in each page of the wizard, possibly altering just a value here or there?

def just_click_next(test, url, override_data):
    response = test.client.get(url)
    while response.status_code != 302:
        form = response.context['form']
        test.assertFalse(form.errors)

        wizard = response.context['wizard']
        current_step = wizard['steps'].current

        # If we have a formset, then we need to do custom handling.
        if hasattr(form, 'forms'):
            current_step_data = step_data.get(current_step, [])
            data = {
                '{}-{}'.format(current_step, key): value
                for key, value in form.management_form.initial.items()
            }
            for i, _form in enumerate(form.forms):
                current_form_data = current_step_data[i] if len(current_step_data) > i else {}
                data.update({
                    '{}-{}'.format(_form.prefix, field): current_form_data.get(
                        field,
                        _form.initial.get(field)
                    )
                    for field in _form.fields
                    if field in current_form_data or field in _form.initial
                })
        else:
            current_step_data = step_data.get(current_step, {})
            data = {
                '{}-{}'.format(form.prefix, field): current_step_data.get(
                    field,
                    form.initial.get(field)
                )
                for field in form.fields
                if field in current_step_data or field in form.initial
            }

        # Add the wizard management form step data.
        data['{}-current_step'.format(wizard['management_form'].prefix)] = current_step

        response = test.client.post(url, data)

    return response

Those paying attention may notice that this one actually does handle formsets in a fairly nice way.

So how do we use it?

class TestViews(TestCase):
    def test_new_revision_wizard(self):
        # Setup required for wizard goes here.
        just_click_next(self, url, {
            'date': {'start_date': '2019-01-01'},
            'confirm': {'confirm': 'on'}
        })
        # Make assertions about the wizard having been run.

This replaces some 80 lines of a test: which actually becomes important when we need to test this with a bunch of subtly different preconditions.

Subquery and Subclasses

Being able to use correlated subqueries in the Django ORM arrived in 1.11, and I also backported it to 1.8.

Quite commonly, I am asked questions about how to use these, so here is an attempt to document them further.

There are three classes that are supplied with Django, but it’s easy to write extensions using subclassing.

Let’s first look at an example of how you might want to use the included classes. We’ll consider a set of temperature sensors, each with a name and a code, both of which are unique. These sensors will log their current temperature at some sort of interval: maybe it’s regular, maybe it varies between devices. We want to keep every reading, but want to only allow one reading for a given sensor+timestamp.

class Sensor(models.Model):
    location = models.TextField(unique=True)
    code = models.TextField(unique=True)


class Reading(models.Model):
    sensor = models.ForeignKey(Sensor, related_name='readings')
    timestamp = models.DateTimeField()
    temperature = models.DecimalField(max_digits=6, decimal_places=3)

    class Meta:
        unique_together = (('sensor', 'timestamp'),)

Some of the things we might want to do for a given sensor:

  • Get the most recent temperature
  • Get the average temperature over a given period
  • Get the maximum temperature over a given period
  • Get the minimum temperature over a given period

If we start with a single sensor instance, we can do each of these without having to use Subquery and friends:

from django.db.models import Avg, Min, Max

most_recent_temperature = sensor.readings.order_by('-timestamp').first().temperature
period_readings = sensor.readings.filter(
    timestamp__gte=start,
    timestamp__lte=finish,
).aggregate(
    average=Avg('temperature'),
    minimum=Min('temperature'),
    maximum=Max('temperature'),
)

We could also get the minimum or maximum using ordering, like we did with the most_recent_temperature.

If we want to do the same for a set of sensors, mostly we can still achieve this (note how similar the code is to the block above):

sensor_readings = Reading.objects.filter(
  timestamp__gte=start,
  timestamp__lte=finish
).values('sensor').annotate(
  average=Avg('temperature'),
  minimum=Min('temperature'),
  maximum=Max('temperature'),
)

We might get something like:

[
    {
        'sensor': 1,
        'average': 17.5,
        'minimum': 11.3,
        'maximum': 25.9
    },
    {
        'sensor': 2,
        'average': 19.63,
        'minimum': 13.6,
        'maximum': 24.33
    },
]

However, it’s not obvious how we would get all of the sensors, and their current temperature in a single query.

Subquery to the rescue!

from django.db.models.expressions import Subquery, OuterRef

current_temperature = Reading.objects.filter(sensor=OuterRef('pk'))\
                                     .order_by('-timestamp')\
                                     .values('temperature')[:1]

Sensor.objects.annotate(
    current_temperature=Subquery(current_temperature)
)

What’s going on here as that we are filtering the Reading objects inside our subquery to only those associated with the sensor in the outer query. This uses the special OuterRef class, that will, when the queryset is “resolved”, build the association. It does mean that if we tried to inspect the current_temperature queryset, we would get an error that it is unresolved.

We then order the filtered readings by newest timestamp first; this, coupled with the slice at the end will limit us to a single row. This is required because the database will reject a query that results in multiple rows being returned for a subquery.

Additionally, we may only have a single column in our subquery: that’s achieved by the .values('temperature').

But maybe there is a problem here: we actually want to know when the reading was taken, as well as the temperature.

We can do that a couple of ways. The simplest is to use two Subqueries:

current_temperature = Reading.objects.filter(sensor=OuterRef('pk'))\
                                     .order_by('-timestamp')[:1]

Sensor.objects.annotate(
    current_temperature=Subquery(current_temperature.values('temperature')),
    last_reading_at=Subquery(current_temperature.values('timestamp')),
)

However, this will do two subqueries at the database level. Since these subqueries will be performed seperately for each row, each additional correlated subquery will result in more work for the database, with possible performance implications.

What about if we are using Postgres, and are okay with turning the temperature and timestamp pair into a JSONB object?

from django.db.models.expressions import Func, F, Value, OuterRef, Subquery
from django.contrib.postgres.fields import JSONField


class JsonBuildObject(Func):
    function = 'jsonb_build_object'
    output_field = JSONField()


last_temperature = Reading.objects.filter(sensor=OuterRef('pk'))\
                                  .order_by('-timestamp')\
                                  .annotate(
                                      json=JsonBuildObject(
                                          Value('timestamp'), F('timestamp'),
                                          Value('temperature'), F('temperature'),
                                      )
                                   ).values('json')[:1]

Sensor.objects.annotate(
    last_temperature=Subquery(last_temperature, output_field=JSONField())
)

Now, your Sensor instances would have an attribute last_temperature, which will be a dict with the timestamp and temperature of the last reading.


There is also a supplied Exists subquery that can be used to force the database to emit an EXISTS statement. This could be used to set a boolean field on our sensors to indicate they have data from within the last day:

recent_readings = Reading.objects.filter(
    sensor=OuterRef('pk'),
    timestamp__gte=datetime.datetime.utcnow() - datetime.timedelta(1)
)
Sensor.objects.annotate(
    has_recent_readings=Exists(recent_readings)
)

Sometimes we’ll have values from multiple rows that we will want to annotate on from the subquery. This can’t be done directly: you will need to aggregate those values in some way. Postgres has a neat feature where you can use an ARRAY() constructor and wrap a subquery in that:

SELECT foo,
       bar,
       ARRAY(SELECT baz
               FROM qux
              WHERE qux.bar = base.bar
              ORDER BY fizz
              LIMIT 5) AS baz
  FROM base

We can build this type of structure using a subclass of Subquery.

from django.contrib.postgres.fields import ArrayField
from django.core.exceptions import FieldError
from django.db.models.expressions import Subquery

class SubqueryArray(Subquery):
    template = 'ARRAY(%(subquery)s)'

    @property
    def output_field(self):
        output_fields = [x.output_field for x in self.get_source_expressions()]

        if len(output_fields) > 1:
            raise FieldError('More than one column detected')

        return ArrayField(base_field=output_fields[0])

And now we can use this where we’ve used a Subquery, but we no longer need to slice to a single row:

json_reading = JsonBuildObject(
    Value('timestamp'), F('timestamp'),
    Value('temperature'), F('temperature'),
)

last_five_readings = Reading.objects.filter(
    sensor=OuterRef('pk')
).order_by('-timestamp').annotate(
    json=json_reading
).values('json')[:5]

Sensor.objects.annotate(last_five_readings=SubqueryArray(last_five_readings))

Each sensor instance would now have up to 5 dicts in a list in it’s attribute last_five_readings.

We could get this data in a slightly different way: let’s say instead of an array, we want a dict keyed by a string representation of the timestamp:

sensor.last_five_readings = {
    '2019-01-01T09:12:35Z': 15.35,
    '2019-01-01T09:13:35Z': 14.33,
    '2019-01-01T09:14:35Z': 14.90,
    ...
}

There is a Postgres aggregate we can use there to do that, too:

class JsonObjectAgg(Subquery):
    template = '(SELECT json_object_agg("_j"."key", "_j"."value") FROM (%(subquery)s) "_j")'
    output_field = JSONField()


last_five_readings = Reading.objects.filter(
    sensor=OuterRef('pk')
).order_by('-timestamp').annotate(
    key=F('timestamp'),
    value=F('temperature'),
).values('key', 'value')[:5]

Sensor.objects.annotate(last_five_readings=JsonObjectAgg(last_five_readings))

Indeed, we can wrap any aggregate in a similar way: to get the number of values of a subquery:

class SubqueryCount(Subquery):
    template = '(SELECT count(*) FROM (%(subquery)s) _count)'
    output_field = models.IntegerField()

Since other aggregates need to operate on a single field, we’ll need something that ensures there is a single value in our .values(), and extract that out and use that in the query.

class SubquerySum(Subquery):
    template = '(SELECT SUM(%(field)s) FROM (%(subquery)s) _sum)'

    def as_sql(self, compiler, connection, template=None, **extra_context):
        if 'field' not in extra_context and 'field' not in self.extra:
            if len(self.queryset._fields) > 1:
                raise FieldError('You must provide the field name, or have a single column')
            extra_context['field'] = self.queryset._fields[0]
        return super(SubquerySum, self).as_sql(
          compiler, connection, template=template, **extra_context
        )

As I mentioned, it’s possible to write a subclass like that for any aggregate function, although it would be far nicer if there was a way to write that purely in the ORM. Maybe one day…

Expression Exclusion Constraints

Today I was working with a junior developer, and was lucky enough to be able to explain exclusion constraints to them. I got partway through it before I realised that the Django model we were working on did not have a range field, but instead had a start and a finish.

class Leave(models.Model):
    person = models.ForeignKey(
        'person.Person',
        related_name='approved_leave',
        on_delete=models.CASCADE,
    )
    start = models.DateTimeField()
    finish = models.DateTimeField()

It turns out that this is not a problem. You can use any expression in a constraint:

ALTER TABLE leave_leave
ADD CONSTRAINT prevent_overlapping_leave
EXCLUDE USING gist(person_id WITH =, TSTZRANGE(start, finish) WITH &&)

Whilst we have application-level validation in place to prevent this, there is a code path that allows it (hence the desire to implement this). Because this is an exclusion constraint, we won’t be able to use the NOT VALID syntax, but will instead have to either fix the invalid data, or use a WHERE clause to only apply the constraint to “new” data.

ALTER TABLE leave_leave
ADD CONSTRAINT prevent_overlapping_leave
EXCLUDE USING gist(person_id WITH =, TSTZRANGE(start, finish) WITH &&)
WHERE start > '2019-07-19';

The other benefit of this is that it creates an index that includes TSTZRANGE(start, finish), which could be used for querying, but also will ensure that start <= finish for all rows.

Infinite Recursion in Postgres

I’m not sure how useful it is, but it turns out it is possible to create a view with an infinite number of rows in postgres:

WITH year AS (
  SELECT 2000 AS year

  UNION ALL

  SELECT year + 1 FROM year
)
SELECT *
  FROM year;

As this stands it really isn’t useful, because it probably won’t start returning rows to the user. However, if you don’t know how many rows you will be generating, you could do something like:

WITH year AS (
  SELECT 2000 AS year

  UNION ALL

  SELECT year + 1 FROM year
)
SELECT *
  FROM year
 LIMIT %s;

Again, this may not seem to be that useful, as you could just use a generate_series(2000, 2000 + %s, 1). But I’m currently working on something that doesn’t always have a fixed count or interval (implementing RRULE repeats in SQL), and I think that maybe this might just be useful…

Merging Adjacent Ranges in Postgres

Previously, I detailed a solution to split/trim/replace overlapping items in a table. Subsequently, I decided I needed to merge all adjacent items that could be merged. In this case, that was with two other fields (only one of which was subject to the exclusion constraint) being identical in adjacent periods.

CREATE EXTENSION IF NOT EXISTS btree_gist;

CREATE TABLE team_membership (
  membership_id SERIAL,
  player_id INTEGER,
  team_id INTEGER,
  period DATERANGE,
  CONSTRAINT prevent_overlapping_memberships EXCLUDE USING gist(player_id WITH =, period WITH &&)
);

Before we can implement the plpgsql trigger function, we need to tell Postgres how to aggregate ranges:

CREATE AGGREGATE sum(anyrange) (
  stype = anyrange,
  sfunc = range_union
);

We should note at this point that range_union, or + (hence the reason I’ve called it SUM) will fail with an error if the two ranges that are being combined do not overlap or touch. We must make sure that in any queries where we are going to use it, that all of the ranges will overlap (and I believe they also must be in “order”, so that as we perform the union on each range in a “reduce” manner we never end up with non-contiguous ranges).

So, let’s look at the trigger function. Initially, I wrote this as two queries:

CREATE OR REPLACE FUNCTION merge_adjacent()
  RETURNS TRIGGER AS $$
  BEGIN
    NEW.period = (SELECT SUM(period) FROM (SELECT NEW.period UNION ALL ...));
    DELETE FROM team_membership ...;
    RETURN NEW;
  END;
  $$ LANGUAGE plpgsql STRICT;

This required me to duplicate the WHERE clauses, and was messy.

Then I remembered you can use the RETURNING clause, and use a CTE, with a SELECT INTO:

CREATE OR REPLACE FUNCTION merge_adjacent()
  RETURNS TRIGGER AS $$

  BEGIN
    WITH matching AS (
      DELETE FROM team_membership mem
            WHERE mem.player_id = NEW.player_id
              AND mem.team_id = NEW.team_id
              AND (mem.period -|- NEW.period OR mem.period && NEW.period)
              AND mem.membership_id <> NEW.membership_id
        RETURNING period
    )
    SELECT INTO NEW.period (
      SELECT SUM(period) FROM (
        SELECT NEW.period
         UNION ALL
        SELECT period FROM matching
        ORDER BY period
    ) _all
    );
    RETURN NEW;
  END;

  $$ LANGUAGE plpgsql STRICT;

CREATE TRIGGER merge_adjacent
BEFORE INSERT OR UPDATE ON team_membership
FOR EACH ROW EXECUTE PROCEDURE merge_adjacent();

The other thing to note about this construct is that it will only work on “already merged” data: if you had ranges:

[2019-01-01, 2019-01-04)
[2019-01-04, 2019-02-02)
# Note there is a gap here...
[2019-05-01, 2019-05-11)
[2019-05-11, 2020-01-01)

and you added in a value to the missing range:

INSERT INTO range (period) VALUES ('[2019-02-02, 2019-05-01)')

You would not merge all of the ranges, only those immediately adjacent. That is, you would wind up with rows:

[2019-01-01, 2019-01-04)
[2019-01-04, 2019-05-11)
[2019-05-11, 2020-01-01)

However, if this trigger is active on the table you would never get to the stage where your data was adjacent but not merged.


This post was updated on 2021-01-04 to ensure that updates to a single row will not result in an error.