Custom Functions in Flask-SQLAlchemy with PostgreSQL

Posted on Posted in Articles

I have a Flask-SQLAlchemy model backed by a PostgreSQL database that looks like this:

class Thing(db.Model)

title = db.Column(db.Text(), nullable=False)
narrative = db.Column(db.Text(), nullable=False)
tags = db.Column(ARRAY(db.Text()), index=True)

My customer wants to have a full text search over all three fields simultaneously, but wants the results ordered by where the hits are. They should be ordered first by tag, then title, then narrative.

This is a job for PostgreSQL full text search. So first I’ll need to create a ‘document’ by concatenating the fields and then search across the set of documents for my search string.

Unfortunately it’s not that simple, as the tags type is a text array. array_to_string is not immutable, since it is dependent on locale, so we will need to provide an immutable convert function.

If I were doing this directly in the database I’d need to do something like this:

CREATE OR REPLACE FUNCTION tags_to_string(text[])
RETURNS text
AS
$BODY$
 select array_to_string($1, ' ');
$BODY$
LANGUAGE sql
IMMUTABLE;

Then, we can create a simple three column text search vector simply by:

CREATE OR REPLACE FUNCTION
three_column_ts_vector(text, text, text)
RETURNS tsvector
AS
$BODY$
 select ( $1 || '':1A '' || $2 || '':1B '' || $3 || '':1C'' )::tsvector;
$BODY$
LANGUAGE sql
STRICT 
IMMUTABLE;

And, of course, we want an index on the documents to make our searches speedy:

CREATE INDEX ON thing USING gin( three_column_ts_vector(tags_to_string(tags), title, narrative) );

So, how do we get Flask_SQLAlchemy to generate the function and how do we use it in our index creation?

I use the Declarative model, so we can’t just grab the metadata object and pass execute statements to it. instead, I use listeners to listen for table creation:

create_function_tags_to_string = DDL("CREATE OR REPLACE FUNCTION tags_to_string(text[]) RETURNS text AS $BODY$ select array_to_string($1, ' '); $BODY$ LANGUAGE sql IMMUTABLE;")
create_function_three_column_ts_vector = DDL("create or replace function three_column_ts_vector(text, text, text) returns tsvector strict immutable language sql as 'select ( $1 || '':1A '' || $2 || '':1B '' || $3 || '':1C'' )::tsvector;';")
create_composite_index = DDL("create index on things using gin( three_column_ts_vector(tags_to_string(tags), title, narrative) );")
event.listen(Thing.__table__, 'after_create', create_function_tags_to_string.execute_if(dialect='postgresql'))
event.listen(Thing.__table__, 'after_create', create_function_three_column_ts_vector.execute_if(dialect='postgresql'))
event.listen(Thing.__table__, 'after_create', create_composite_index.execute_if(dialect='postgresql'))

Note that this is only triggered by the initial creation of the table. You may wish to add listeners to other events to capture alter tables or drops.

An alternative might be to add this to your migrate scripts – but it means you’re maintaining model attributes outside your model classes, which is fraught with potentially hard to find bugs.