diff --git a/aiida/storage/psql_dos/orm/querybuilder/main.py b/aiida/storage/psql_dos/orm/querybuilder/main.py index d7efafe7ef..ada9560824 100644 --- a/aiida/storage/psql_dos/orm/querybuilder/main.py +++ b/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): diff --git a/tests/orm/test_querybuilder/test_as_sql.txt b/tests/orm/test_querybuilder/test_as_sql.txt index d610d3f564..6d3bb3f1bc 100644 --- a/tests/orm/test_querybuilder/test_as_sql.txt +++ b/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}