DjangoORM: Resolve F Expression in Custom DB Function

Question

I am trying to write a custom PostgreSQL function in Django that will coerce datetimes to a specified timezone inside of a queryset. My first pass at the db function looks like this:

from django.db.models.expressions import Func

class DateTimeInTimezone(Func):
    template="%(expressions)s AT TIME ZONE %(tz_info)s"

This function works in the simple case where I pass a timezone string into the function directly like so:

MyModel.objects.filter(timefield__gte=DateTimeInTimezone(Now(), tz_info='EST'))

However it doesn't work in the more complex case, where the timezone is defined on some field on the model. Consider the following contrived example:

class User(models.Model):
    time_zone = models.CharField()

class Meeting(models.Model):
    users = models.ManyToManyField(User, related_name='meetings')
    start_time = models.DateTimeField()  # in UTC
    end_time = models.DateTimeField()  # in UTC

To answer the question "What users will be in a meeting at 12pm local time today?", I'd need some variation of this queryset:

noon_utc = ...
User.objects.filter(
    meetings__start_time__lte=DateTimeInTimezone(noon_utc, tz_info=F('time_zone')),
    meetings__end_time__gt=DateTimeInTimezone(noon_utc, tz_info=F('time_zone'))
)

As currently written, however, DateTimeInTimezone will simply inject the string F('time_zone') into the sql rather than resolving the field.

Is it possible to add support for F Expressions to this function? Is there some other approach that I should consider?

score 5 · Accepted Answer · answered Oct 03 '17 at 09:24

A simple solution is possible with parameter arg_joiner:

class DateTimeInTimezone(Func):
    function = ''
    arg_joiner = ' AT TIME ZONE '

    def __init__(self, timestamp, tz_info):
        super(DateTimeInTimezone, self).__init__(timestamp, tz_info)

The method __init__ is used only for the purpose of readability with clear names of parameters. Then arity is not important if parameters are declared by __init__.

A oneliner function is useful for fast development if readability is not important:

...filter(
    meetings__start_time__lte=Func(noon_utc, tz_info=F('time_zone'), arg_joiner=' AT TIME ZONE ', function=''),
)

Verified:

>>> qs = User.objects.filter(...)
>>> print(str(qs.query))
SELECT ... WHERE ("app_meeting"."start_time" <= ((2017-10-03 08:18:12.663640 AT TIME ZONE "app_user"."time_zone")) AND ...)

Hey, this is great! Thanks for suggesting this solution. I made your response the accepted answer because it is definitely preferable to monkey-patching (as you mentioned in your comment). — chukkwagon, Oct 03 '17 at 22:07

score 1 · Answer 2 · answered Apr 18 '17 at 19:53

Found an acceptable solution. I overrode the as_sql method for function like so, allowing the django internals to resolve the F expression and then separating it back out into a kwarg I could use in a different part of the template.

class DateTimeInTimezone(Func):
'''
Coerce a datetime into a specified timezone
'''
template="%(expressions)s AT TIME ZONE %(tz_info)s"
arity = 2

def as_sql(self, compiler, connection, function=None, template=None, arg_joiner=None, **extra_context):
    connection.ops.check_expression_support(self)
    sql_parts = []
    params = []
    for arg in self.source_expressions:
        arg_sql, arg_params = compiler.compile(arg)
        sql_parts.append(arg_sql)
        params.extend(arg_params)
    data = self.extra.copy()
    data.update(**extra_context)
    # Use the first supplied value in this order: the parameter to this
    # method, a value supplied in __init__()'s **extra (the value in
    # `data`), or the value defined on the class.
    if function is not None:
        data['function'] = function
    else:
        data.setdefault('function', self.function)
    template = template or data.get('template', self.template)
    arg_joiner = arg_joiner or data.get('arg_joiner', self.arg_joiner)
    data['expressions'] = data['field'] = arg_joiner.join(sql_parts)
    parts = data['expressions'].split(', ')
    data['expressions'] = parts[0]
    data['tz_info'] = parts[1]
    return template % data, params

I added the three lines between the assignment of data['expressions'] and the final return template % data, params. This isn't a great long term solution as the django internals for this method could change in the next version, but it suits my needs for the time being.

+1 it works, but monkey patching is the last resort because maintenance could be laborious after Django upgrade. — hynekcer, Oct 03 '17 at 09:38

DjangoORM: Resolve F Expression in Custom DB Function

2 Answers2