From 6bf7f9a74e9cd7edd8e328f400be5e86b630c73c Mon Sep 17 00:00:00 2001
From: "Chart.js" <>
Date: Mon, 3 Oct 2022 14:52:06 -0400
Subject: [PATCH] feat: JSONB Contains

---
 README.md                                     |  9 ++++
 src/ormar_postgres_extensions/fields/jsonb.py | 36 +++++++++++++
 tests/fields/test_jsonb.py                    | 54 +++++++++++++++++++
 3 files changed, 99 insertions(+)

diff --git a/README.md b/README.md
index 363a445..9f37c90 100644
--- a/README.md
+++ b/README.md
@@ -49,6 +49,15 @@ class JSONBTestModel(ormar.Model):
     id: int = ormar.Integer(primary_key=True)
     data: dict = ormar_pg_ext.JSONB()
 ```
+
+##### jsonb_contains
+
+The maps to the [`contains`](https://docs.sqlalchemy.org/en/14/dialects/postgresql.html#sqlalchemy.dialects.postgresql.JSONB.Comparator.contains) operator in Postgres.
+
+```python
+await JSONBTestModel.objects.filter(data__jsonb_contains=dict(key="value")).all()
+```
+
 #### Array
 
 Array field requires a bit more setup to pass the type of the array into the field
diff --git a/src/ormar_postgres_extensions/fields/jsonb.py b/src/ormar_postgres_extensions/fields/jsonb.py
index 2c41560..9bc7f29 100644
--- a/src/ormar_postgres_extensions/fields/jsonb.py
+++ b/src/ormar_postgres_extensions/fields/jsonb.py
@@ -4,6 +4,42 @@
 from sqlalchemy.dialects import postgresql
 
 
+def jsonb_contains(self, other: Any) -> ormar.queryset.clause.FilterGroup:
+    """
+    works as postgresql `column @> VALUE::jsonb`
+    :param other: value to check against operator
+    :type other: Any
+    :return: FilterGroup for operator
+    :rtype: ormar.queryset.clause.FilterGroup
+    """
+    return self._select_operator(op="jsonb_contains", other=other)
+
+
+# Need to patch the filter objects to support JSONB specifc actions
+FIELD_ACCESSOR_MAP = [
+    ("jsonb_contains", jsonb_contains),
+]
+
+
+for (method_name, method) in FIELD_ACCESSOR_MAP:
+    setattr(ormar.queryset.FieldAccessor, method_name, method)
+
+
+# These lines allow Ormar to lookup the new filter methods and map
+# it to the correct PGSQL functions
+ACCESSOR_MAP = [
+    ("jsonb_contains", "contains"),
+]
+
+for (ormar_operation, pg_operation) in ACCESSOR_MAP:
+    ormar.queryset.actions.filter_action.FILTER_OPERATORS[
+        ormar_operation
+    ] = pg_operation
+    ormar.queryset.actions.filter_action.METHODS_TO_OPERATORS[
+        ormar_operation
+    ] = ormar_operation
+
+
 class JSONB(ormar.JSON):
     """
     Custom JSON field uses a native PG JSONB type
diff --git a/tests/fields/test_jsonb.py b/tests/fields/test_jsonb.py
index 201a36f..147ee29 100644
--- a/tests/fields/test_jsonb.py
+++ b/tests/fields/test_jsonb.py
@@ -46,3 +46,57 @@ async def test_create_model_with_nullable_jsonb(db):
 
     found = await NullableJSONBTestModel.objects.get()
     assert found.data is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "value1, value2, value3",
+    [
+        # Test with some JSON primitives
+        ("value1", "value2", "value3"),
+        (1, 2, 3),
+        (True, False, None),
+        # Verifies that matches with NULL work
+        (True, None, False),
+    ],
+)
+async def test_contains(db, value1, value2, value3):
+    await JSONBTestModel(data=json.dumps(dict(key=value1))).save()
+    await JSONBTestModel(data=json.dumps(dict(key=value2))).save()
+
+    found = await JSONBTestModel.objects.filter(
+        data__jsonb_contains=dict(key=value2)
+    ).all()
+    assert len(found) == 1
+
+    found = await JSONBTestModel.objects.filter(
+        data__jsonb_contains=dict(key=value3)
+    ).all()
+    assert len(found) == 0
+
+
+@pytest.mark.asyncio
+async def test_contains_array(db):
+    await JSONBTestModel(data=json.dumps([1, 2])).save()
+    await JSONBTestModel(data=json.dumps([1, 3])).save()
+
+    found = await JSONBTestModel.objects.filter(data__jsonb_contains=[1]).all()
+    assert len(found) == 2
+
+    found = await JSONBTestModel.objects.filter(data__jsonb_contains=[1, 2]).all()
+    assert len(found) == 1
+
+    found = await JSONBTestModel.objects.filter(data__jsonb_contains=[4]).all()
+    assert len(found) == 0
+
+
+@pytest.mark.asyncio
+async def test_contains_array_text(db):
+    await JSONBTestModel(data=json.dumps([1, 2])).save()
+    await JSONBTestModel(data=json.dumps([1, 3])).save()
+
+    found = await JSONBTestModel.objects.filter(data__jsonb_contains="1").all()
+    assert len(found) == 2
+
+    found = await JSONBTestModel.objects.filter(data__jsonb_contains="4").all()
+    assert len(found) == 0