CODEWITHSUNDEEP
pythonpostgresqlsqlalchemyflask-sqlalchemypsycopg2
constemi
constemi

Reputation: 75

Sqlalchemy representation for custom postgres range type

I have the need for a custom range type. I am trying to represent an hourly range. For each day of the week a range (datetime.time, datetime.time) rather than separate TIME columns, I would like to have access to Postgres/sqlalchemy range operators if possible.

I'm looking for something like TSRANGE but hours instead of the normal (datetime.datetime, datetime.datetime)

In postgres itself this works wonderfully. For example.

create type timerange as range (subtype = time);

create table schedule 
(
  id integer not null primary key,
  time_range timerange
);

insert into schedule 
values
(1, timerange(time '08:00', time '10:00', '[]')),
(2, timerange(time '10:00', time '12:00', '[]'));

select *
from schedule
where time_range @> time '09:00'

So here is the question. How do I represent this custom type I have created in Postgres in SQLAlchemy? Subclass TSRANGE, TypeDecorator on TIME, or possibly create a new SQLALchemy UserDefinedType. I am not quite sure which way to go. Any suggestions would be much appreciated. Thanks!

Upvotes: 7

Views: 2355

Answers (1)

Ilja Everilä
Ilja Everilä

Reputation: 52949

In order to use custom range types you need to dig a bit deeper:

When instantiating models that use these column types, you should pass whatever data type is expected by the DBAPI driver you’re using for the column type. For psycopg2 these are NumericRange, DateRange, DateTimeRange and DateTimeTZRange or the class you’ve registered with register_range().

In other words you have to both register your custom range type with your DBAPI – usually psycopg2 – and create the SQLAlchemy types to match the registered types. register_range() takes the name of the PostgreSQL range type, a (strict) subclass of Range and a connection/cursor used to fetch the oids. It can register the new range type either globally, or locally to the given connection or cursor:

In [2]: import psycopg2.extras

The value type you should use when creating instances of models:

In [3]: class TimeRange(psycopg2.extras.Range):
   ...:     pass
   ...:

Use raw_connection() in SQLAlchemy to fetch a proxy to the underlying psycopg2 connection. Registering should perhaps be done in a setup function in an actual implementation:

In [4]: conn = engine.raw_connection()

In [5]: cur = conn.cursor()

In [6]: psycopg2.extras.register_range('timerange', TimeRange, cur, globally=True)
Out[6]: <psycopg2._range.RangeCaster at 0x7f1c980dbe80>

In [7]: cur.close()

In [8]: conn.close()

Next create the SQLAlchemy range type to match the registered TimeRange. A TypeDecorator is ill suited, since you're not using an existing type. UserDefinedType should be the base of all entirely new types. For range operators include the RangeOperators mixin:

It is used by all the range types provided in the postgres dialect and can likely be used for any range types you create yourself.

The rest is pretty much copied directly from the predefined range types:

In [11]: from sqlalchemy.dialects import postgresql

In [13]: from sqlalchemy import types as sqltypes

In [14]: class TIMERANGE(postgresql.ranges.RangeOperators, sqltypes.UserDefinedType):
    ...:     def get_col_spec(self, **kw):
    ...:         return 'timerange'

This is only required for reflection.

In [16]: postgresql.base.ischema_names['timerange'] = TIMERANGE

And then just create your tables and use as you would normally:

In [17]: schedule = Table('schedule', metadata, autoload=True, autoload_with=engine)

In [18]: schedule
Out[18]: Table('schedule', MetaData(bind=Engine(postgresql:///sopython)), Column('id', INTEGER(), table=<schedule>, primary_key=True, nullable=False), Column('time_range', TIMERANGE(), table=<schedule>), schema=None)

In [19]: session.query(schedule).all()
Out[19]: 
[(1, TimeRange(datetime.time(8, 0), datetime.time(10, 0), '[]')),
 (2, TimeRange(datetime.time(10, 0), datetime.time(12, 0), '[]'))]

In [20]: session.query(schedule).\
    ...:     filter(schedule.c.time_range.contains(time(9, 0))).\
    ...:     all()
2017-04-11 10:01:23,864 INFO sqlalchemy.engine.base.Engine SELECT schedule.id AS schedule_id, schedule.time_range AS schedule_time_range 
FROM schedule 
WHERE schedule.time_range @> %(time_range_1)s
2017-04-11 10:01:23,864 INFO sqlalchemy.engine.base.Engine {'time_range_1': datetime.time(9, 0)}
Out[20]: [(1, TimeRange(datetime.time(8, 0), datetime.time(10, 0), '[]'))]

Upvotes: 9

Related Questions