👌 IMPROVE: Use sqlalchemy.func for JSONB QB filters (#5393)

The QueryBuilder declared a custom `jsonb_typeof` FunctionElement,
for which sqlalchemy disables query caching by default.
This switches to the built-in `sqlalchemy.func` way of accessing 
the postgres functions, which gets rid of this issue along the way.

aiida/storage/psql_dos/orm/querybuilder/main.py

@@ -20,16 +20,14 @@
 from sqlalchemy import not_, or_
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.exc import SAWarning
-from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import aliased
 from sqlalchemy.orm.attributes import InstrumentedAttribute, QueryableAttribute
 from sqlalchemy.orm.query import Query
 from sqlalchemy.orm.session import Session
 from sqlalchemy.orm.util import AliasedClass
-from sqlalchemy.sql.compiler import SQLCompiler, TypeCompiler
+from sqlalchemy.sql.compiler import SQLCompiler
 from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList, Cast, ColumnClause, ColumnElement, Label
 from sqlalchemy.sql.expression import case, text
-from sqlalchemy.sql.functions import FunctionElement
 from sqlalchemy.types import Boolean, DateTime, Float, Integer, String
 
 from aiida.common.exceptions import NotExistent
@@ -38,41 +36,9 @@
 
 from .joiner import SqlaJoiner
 
-
-class jsonb_array_length(FunctionElement):  # pylint: disable=abstract-method,invalid-name
-    name = 'jsonb_array_len'
-
-
-@compiles(jsonb_array_length)
-def compile(element, compiler: TypeCompiler, **kwargs):  # pylint: disable=function-redefined, redefined-builtin
-    """
-    Get length of array defined in a JSONB column
-    """
-    return f'jsonb_array_length({compiler.process(element.clauses, **kwargs)})'
-
-
-class array_length(FunctionElement):  # pylint: disable=abstract-method,invalid-name
-    name = 'array_len'
-
-
-@compiles(array_length)  # type: ignore[no-redef]
-def compile(element, compiler: TypeCompiler, **kwargs):  # pylint: disable=function-redefined
-    """
-    Get length of array defined in a JSONB column
-    """
-    return f'array_length({compiler.process(element.clauses, **kwargs)})'
-
-
-class jsonb_typeof(FunctionElement):  # pylint: disable=abstract-method,invalid-name
-    name = 'jsonb_typeof'
-
-
-@compiles(jsonb_typeof)  # type: ignore[no-redef]
-def compile(element, compiler: TypeCompiler, **kwargs):  # pylint: disable=function-redefined
-    """
-    Get length of array defined in a JSONB column
-    """
-    return f'jsonb_typeof({compiler.process(element.clauses, **kwargs)})'
+jsonb_typeof = sa_func.jsonb_typeof
+jsonb_array_length = sa_func.jsonb_array_length
+array_length = sa_func.array_length
 
 
 class SqlaQueryBuilder(BackendQueryBuilder):

tests/orm/test_querybuilder/test_as_sql.txt

@@ -1 +1 @@
-'SELECT db_dbnode_1.id, db_dbnode_1.uuid \nFROM db_dbnode AS db_dbnode_1 \nWHERE CAST(db_dbnode_1.node_type AS VARCHAR) LIKE %(param_1)s AND CASE WHEN (jsonb_typeof((db_dbnode_1.extras #> %(extras_1)s)) = %(param_2)s) THEN (db_dbnode_1.extras #>> %(extras_1)s) = %(param_3)s ELSE %(param_4)s END' % {'param_1': '%', 'extras_1': ('tag4',), 'param_2': 'string', 'param_3': 'appl_pecoal', 'param_4': False}
+'SELECT db_dbnode_1.id, db_dbnode_1.uuid \nFROM db_dbnode AS db_dbnode_1 \nWHERE CAST(db_dbnode_1.node_type AS VARCHAR) LIKE %(param_1)s AND CASE WHEN (jsonb_typeof((db_dbnode_1.extras #> %(extras_1)s)) = %(jsonb_typeof_1)s) THEN (db_dbnode_1.extras #>> %(extras_1)s) = %(param_2)s ELSE %(param_3)s END' % {'param_1': '%', 'extras_1': ('tag4',), 'jsonb_typeof_1': 'string', 'param_2': 'appl_pecoal', 'param_3': False}