QueryBuilder __str__ does not replace all templates

[ltalirz]

The query string synthesized by the QueryBuilder sometimes contains template variables that are replaced somewhere down the line before the query is handed over to postgresql.

It seems that there are cases, where calling print(qb) on a QueryBuilder instance results in a string, where not all of the variables have been replaced -- meaning, when copy-pasting this string into a psql shell, the query will fail.

A very simple example is this (see %(extras_1)s)

In [12]: qb=QueryBuilder()

In [13]: qb.append(Node, filters={"extras.tag4": "appl_pecoal"})
Out[13]: <aiida.orm.querybuilder.QueryBuilder at 0x11d9a9a58>

In [14]: print(qb)
SELECT db_dbnode_1.id, db_dbnode_1.uuid, db_dbnode_1.node_type, db_dbnode_1.process_type, db_dbnode_1.label, db_dbnode_1.description, db_dbnode_1.ctime, db_dbnode_1.mtime, db_dbnode_1.attributes, db_dbnode_1.extras, db_dbnode_1.user_id, db_dbnode_1.dbcomputer_id
FROM db_dbnode AS db_dbnode_1
WHERE CAST(db_dbnode_1.node_type AS VARCHAR) LIKE '%%' AND CASE WHEN (jsonb_typeof(db_dbnode_1.extras #> %(extras_1)s) = 'string') THEN (db_dbnode_1.extras #>> '{tag4}') = 'appl_pecoal' ELSE false END

[giovannipizzi]

Minimal working example

In an effort to have a minimal reproduction of the issue, I got this working example (working only for a SQLAlchemy backend for simplicity):

from sqlalchemy.orm import aliased
from sqlalchemy.sql.expression import FunctionElement
from sqlalchemy import case, and_
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.orm import Query
from sqlalchemy.dialects import postgresql as mydialect

from aiida.backends.sqlalchemy.models.node import DbNode
from aiida.manage.manager import get_manager
backend = get_manager().get_backend()
session = backend.get_session()

class jsonb_typeof(FunctionElement):
    name = 'jsonb_typeof'

@compiles(jsonb_typeof)
def compile(element, compiler, **_kw):  # pylint: disable=function-redefined
    """
    Get length of array defined in a JSONB column
    """
    return 'jsonb_typeof(%s)' % compiler.process(element.clauses)

alias = aliased(DbNode)

column_name = 'extras'
column = getattr(alias, column_name)

json_path = column[('key',)]

expr = case([(jsonb_typeof(json_path) == 'string', json_path.astext == 'value')], else_=False)

#query = session.query(alias).filter(json_path.astext == 'value')

query = session.query(alias).filter(expr)

print(str(query.statement.compile(compile_kwargs={'literal_binds': True}, dialect=mydialect.dialect())))

This returns indeed a problematic

SELECT db_dbnode_1.id, db_dbnode_1.uuid, db_dbnode_1.node_type, db_dbnode_1.process_type, db_dbnode_1.label, db_dbnode_1.description, db_dbnode_1.ctime, db_dbnode_1.mtime, db_dbnode_1.attributes, db_dbnode_1.extras, db_dbnode_1.dbcomputer_id, db_dbnode_1.user_id 
FROM db_dbnode AS db_dbnode_1 
WHERE CASE WHEN (jsonb_typeof(db_dbnode_1.extras #> %(extras_1)s) = 'string') THEN (db_dbnode_1.extras #>> '{key}') = 'value' ELSE false END

(with %(extras_1)s unsubstituted).

Solution

The problem is in the @compiles(jsonb_typeof) part, that does not pass the kwargs (_kw) down to compile.process.

The solution is to changing the return clause with return 'jsonb_typeof(%s)' % compiler.process(element.clauses, **_kw).
This indeed returns correctly:

SELECT db_dbnode_1.id, db_dbnode_1.uuid, db_dbnode_1.node_type, db_dbnode_1.process_type, db_dbnode_1.label, db_dbnode_1.description, db_dbnode_1.ctime, db_dbnode_1.mtime, db_dbnode_1.attributes, db_dbnode_1.extras, db_dbnode_1.dbcomputer_id, db_dbnode_1.user_id 
FROM db_dbnode AS db_dbnode_1 
WHERE CASE WHEN (jsonb_typeof((db_dbnode_1.extras #> '{key}')) = 'string') THEN (db_dbnode_1.extras #>> '{key}') = 'value' ELSE false END

Comment

This needs to be applied to all functions decorated with @compiles, both in aiida/orm/implementation/django/querybuilder.py and in aiida/orm/implementation/sqlalchemy/querybuilder.py.

@lekah are you willing to do it? Should be quick. Otherwise please unassigned yourself.

[lekah]

@giovannipizzi Sorry, it's my last 2 weeks here and I'm very busy finishing up things, so I can't take care of this.