Merge pull request #1078 from ImageMarkup/use-opensearch-dsl

Use opensearch-dsl for opensearch queries
ImageMarkup · Feb 6, 2025 · 8ec73d8 · 8ec73d8
2 parents 8bc051b + 5c01028
commit 8ec73d8
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 29 deletions.
diff --git a/isic/ingest/models/lesion.py b/isic/ingest/models/lesion.py
@@ -15,6 +15,8 @@
 from django.db.models.query import QuerySet
 from django.db.models.query_utils import Q
 from django.urls import reverse
+from opensearch_dsl import Search
+from opensearch_dsl.query import Q as OpenSearchQ  # noqa: N811
 
 from isic.core.constants import LESION_ID_REGEX
 
@@ -27,39 +29,35 @@ def get_lesion_count_for_user(user: User | AnonymousUser) -> int:
     if user.is_staff:
         return es.count(index=settings.ISIC_ELASTICSEARCH_LESIONS_INDEX)["count"]
 
-    query = {
-        "size": 0,
-        "query": {
-            "bool": {
-                "must_not": {
-                    "nested": {
-                        "path": "images",
-                        "query": {
-                            "bool": {
-                                "must_not": [
-                                    {"bool": {"should": [{"term": {"images.public": True}}]}}
-                                ]
-                            }
-                        },
-                    }
-                }
-            }
-        },
-        "track_total_hits": True,
-    }
+    should = [OpenSearchQ("term", **{"images.public": True})]
 
     if user.is_authenticated:
-        query["query"]["bool"]["must_not"]["nested"]["query"]["bool"]["must_not"][0]["bool"][  # type: ignore[index]
-            "should"
-        ] += [
-            {"term": {"images.contributor_owner_ids": user.pk}},
-            {"term": {"images.shared_to": user.pk}},
+        should += [
+            OpenSearchQ("term", **{"images.contributor_owner_ids": user.pk}),
+            OpenSearchQ("term", **{"images.shared_to": user.pk}),
         ]
 
-    return es.search(
-        index=settings.ISIC_ELASTICSEARCH_LESIONS_INDEX,
-        body=query,
-    )["hits"]["total"]["value"]
+    # find all documents where it's NOT true that the nested images array does NOT
+    # contain any images that match should (OR of should).
+    query = (
+        Search(using=es, index=settings.ISIC_ELASTICSEARCH_LESIONS_INDEX)
+        .query(
+            "bool",
+            must_not=[
+                OpenSearchQ(
+                    "nested",
+                    path="images",
+                    query=OpenSearchQ(
+                        "bool",
+                        must_not=[OpenSearchQ("bool", should=should)],
+                    ),
+                )
+            ],
+        )
+        .extra(track_total_hits=True, size=0)
+    )
+
+    return query.execute().hits.total.value
 
 
 def _default_id():

diff --git a/requirements.txt b/requirements.txt
@@ -72,6 +72,7 @@ mdurl==0.1.2
 numpy==2.2.2
 oauth2client==4.1.3
 oauthlib==3.2.2
+opensearch-dsl==2.1.0
 opensearch-py==2.8.0
 orderly-set==5.2.3
 packaging==24.2

diff --git a/setup.py b/setup.py
@@ -76,6 +76,7 @@
         "numpy",
         "oauth2client",
         "opensearch-py",
+        "opensearch-dsl",
         "pandas",
         "Pillow",
         "psycopg",