Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Feb 25, 2022
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
20b4518
fix: enable sqlachemy query caching in QB
ltalirz Feb 25, 2022
de7d0e2
simplify code
ltalirz Feb 25, 2022
4da53dc
reintrodyce type: ignore
ltalirz Feb 25, 2022
8c8cc16
update test data
ltalirz Feb 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 4 additions & 38 deletions aiida/storage/psql_dos/orm/querybuilder/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion tests/orm/test_querybuilder/test_as_sql.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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}