QueryBuilder: Implementation of contains Filter Operator applied to JSON Fields for SQLite backend

src/aiida/storage/sqlite_zip/orm.py

@@ -285,8 +285,7 @@ def _cast_json_type(comparator: JSON.Comparator, value: Any) -> Tuple[ColumnElem
             return case((type_filter, casted_entity.ilike(value, escape='\\')), else_=False)
 
         if operator == 'contains':
-            # to-do, see: https://github.com/sqlalchemy/sqlalchemy/discussions/7836
-            raise NotImplementedError('The operator `contains` is not implemented for SQLite-based storage plugins.')
+            return func.json_contains(database_entity, json.dumps(value))
 
         if operator == 'has_key':
             return (

src/aiida/storage/sqlite_zip/utils.py

@@ -48,12 +48,43 @@
     cursor.close()
 
 
+def _contains(lhs: Union[dict, list], rhs: Union[dict, list]):
+    if isinstance(lhs, dict) and isinstance(rhs, dict):
+        for key in rhs:
+            if key not in lhs or not _contains(lhs[key], rhs[key]):
+                return False
+        return True
+    elif isinstance(lhs, list) and isinstance(rhs, list):
+        for item in rhs:
+            if not any(_contains(element, item) for element in lhs):
+                return False
+        return True
+    else:
+        return lhs == rhs
+
+
+def _json_contains(lhs: Union[str, bytes, bytearray, dict, list], rhs: Union[str, bytes, bytearray, dict, list]):
+    try:
+        if isinstance(lhs, (str, bytes, bytearray)):
+            lhs = json.loads(lhs)
+        if isinstance(rhs, (str, bytes, bytearray)):
+            rhs = json.loads(rhs)
+    except json.JSONDecodeError:
+        return 0
+    return int(_contains(lhs, rhs))
+
+
+def register_json_contains(dbapi_connection, _):
+    dbapi_connection.create_function('json_contains', 2, _json_contains)
+
+
 def create_sqla_engine(path: Union[str, Path], *, enforce_foreign_keys: bool = True, **kwargs) -> Engine:
     """Create a new engine instance."""
     engine = create_engine(f'sqlite:///{path}', json_serializer=json.dumps, json_deserializer=json.loads, **kwargs)
     event.listen(engine, 'connect', sqlite_case_sensitive_like)
     if enforce_foreign_keys:
         event.listen(engine, 'connect', sqlite_enforce_foreign_keys)
+    event.listen(engine, 'connect', register_json_contains)
     return engine
 
 

tests/storage/sqlite/test_orm.py

@@ -129,6 +129,166 @@ def test_qb_json_filters(filters, matches):
     assert qbuilder.count() == matches
 
 
+class TestJsonFilters:
+    @pytest.mark.parametrize(
+        'data,filters,is_match',
+        (
+            # contains different types of element
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'contains': [1]}}, True),
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'contains': ['2']}}, True),
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'contains': [None]}}, True),
+            # contains multiple elements of various types
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'contains': [1, None]}}, True),
+            # contains non-exist elements
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'contains': [114514]}}, False),
+            # contains empty set
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'contains': []}}, True),
+            ({'arr': []}, {'attributes.arr': {'contains': []}}, True),
+            # nested arrays
+            ({'arr': [[1, 0], [0, 2]]}, {'attributes.arr': {'contains': [[1, 0]]}}, True),
+            ({'arr': [[2, 3], [0, 1], []]}, {'attributes.arr': {'contains': [[1, 0]]}}, True),  # order doesn't matter
+            ({'arr': [[2, 3], [1]]}, {'attributes.arr': {'contains': [[4]]}}, False),
+            # TODO: the test below is supposed to pass but currently doesn't
+            # ({'arr': [[2, 3], [1]]}, {'attributes.arr': {'contains': [[2]]}}, False),
+            # negations
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'!contains': [1]}}, False),
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'!contains': []}}, False),
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'!contains': [114514]}}, True),
+            ({'arr': [1, '2', None]}, {'attributes.arr': {'!contains': [1, 114514]}}, True),
+            # TODO: these pass, but why? are these behaviors expected?
+            # non-exist `attr_key`s
+            ({'foo': []}, {'attributes.arr': {'contains': []}}, False),
+            # ({'foo': []}, {'attributes.arr': {'!contains': []}}, False),
+        ),
+        ids=json.dumps,
+    )
+    @pytest.mark.usefixtures('aiida_profile_clean')
+    @pytest.mark.requires_psql
+    def test_json_filters_contains_arrays(self, data, filters, is_match):
+        """Test QueryBuilder filter `contains` for JSON array fields"""
+        profile = SqliteTempBackend.create_profile(debug=False)
+        backend = SqliteTempBackend(profile)
+        Dict(data, backend=backend).store()
+        qb = QueryBuilder(backend=backend).append(Dict, filters=filters)
+        assert qb.count() in {0, 1}
+        found = qb.count() == 1
+        assert found == is_match
+
+    @pytest.mark.parametrize(
+        'data,filters,is_match',
+        (
+            # contains different types of values
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'contains': {'k1': 1}}},
+                True,
+            ),
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'contains': {'k1': 1, 'k2': '2'}}},
+                True,
+            ),
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'contains': {'k3': None}}},
+                True,
+            ),
+            # contains empty set
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'contains': {}}},
+                True,
+            ),
+            # doesn't contain non-exist entries
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'contains': {'k1': 1, 'k': 'v'}}},
+                False,
+            ),
+            # negations
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'!contains': {'k1': 1}}},
+                False,
+            ),
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'!contains': {'k1': 1, 'k': 'v'}}},
+                True,
+            ),
+            (
+                {
+                    'dict': {
+                        'k1': 1,
+                        'k2': '2',
+                        'k3': None,
+                    }
+                },
+                {'attributes.dict': {'!contains': {}}},
+                False,
+            ),
+            # TODO: these pass, but why? are these behaviors expected?
+            # non-exist `attr_key`s
+            ({'map': {}}, {'attributes.dict': {'contains': {}}}, False),
+            # ({'map': {}}, {'attributes.dict': {'!contains': {}}}, False),
+        ),
+        ids=json.dumps,
+    )
+    @pytest.mark.usefixtures('aiida_profile_clean')
+    @pytest.mark.requires_psql
+    def test_json_filters_contains_object(self, data, filters, is_match):
+        """Test QueryBuilder filter `contains` for JSON object fields"""
+        profile = SqliteTempBackend.create_profile(debug=False)
+        backend = SqliteTempBackend(profile)
+        Dict(data, backend=backend).store()
+        qb = QueryBuilder(backend=backend).append(Dict, filters=filters)
+        assert qb.count() in {0, 1}
+        found = qb.count() == 1
+        assert found == is_match
+
+
 @pytest.mark.parametrize(
     'filters,matches',
     (