Implement searching different indexes

core/lib/opensearch.py

@@ -1,6 +1,6 @@
 from django.conf import settings
 from opensearchpy import OpenSearch
-from opensearchpy.exceptions import RequestError
+from opensearchpy.exceptions import NotFoundError, RequestError
 
 from core.lib.threshold import annotate_num_chans, annotate_num_users, annotate_online
 
@@ -44,14 +44,24 @@ def annotate_results(results_parsed):
 
     for net in nets:
         # Annotate the online attribute from Threshold
-        nicks = [
-            x["nick"] for x in results_parsed if x["src"] == "irc" and x["net"] == net
-        ]
-        channels = [
-            x["channel"]
-            for x in results_parsed
-            if x["src"] == "irc" and x["net"] == net
-        ]
+        nicks = list(
+            set(
+                [
+                    x["nick"]
+                    for x in results_parsed
+                    if x["src"] == "irc" and x["net"] == net and "nick" in x
+                ]
+            )
+        )
+        channels = list(
+            set(
+                [
+                    x["channel"]
+                    for x in results_parsed
+                    if x["src"] == "irc" and x["net"] == net and "channel" in x
+                ]
+            )
+        )
         online_info = annotate_online(net, nicks)
         # Annotate the number of users in the channel
         num_users = annotate_num_users(net, channels)
@@ -137,6 +147,9 @@ def run_main_query(client, user, query, custom_query=False, index=None, size=Non
     except RequestError as err:
         print("OpenSearch error", err)
         return err
+    except NotFoundError as err:
+        print("OpenSearch error", err)
+        return err
     filter_blacklisted(user, response)
     return response
 
@@ -254,16 +267,36 @@ def query_results(request, query_params, size=None):
                     search_query["query"]["bool"]["must_not"] = [item]
         if sort:
             search_query["sort"] = sort
+        if "index" in query_params:
+            if not request.user.is_superuser:
+                message = "How did you get here?"
+                message_class = "danger"
+                return {"message": message, "class": message_class}
+            else:
+                index = query_params["index"]
+                if index == "main":
+                    index = settings.OPENSEARCH_INDEX_MAIN
+                elif index == "meta":
+                    index = settings.OPENSEARCH_INDEX_META
+                elif index == "int":
+                    index = settings.OPENSEARCH_INDEX_INT
+                else:
+                    message = "Index is not valid."
+                    message_class = "danger"
+                    return {"message": message, "class": message_class}
+        else:
+            index = settings.OPENSEARCH_INDEX_MAIN
         results = run_main_query(
             client,
             request.user,  # passed through run_main_query to filter_blacklisted
             search_query,
             custom_query=True,
+            index=index,
             size=size,
         )
         if not results:
             return False
-        if isinstance(results, RequestError):
+        if isinstance(results, Exception):
             message = results.info["error"]["root_cause"][0]["reason"]
             message_class = "danger"
             return {"message": message, "class": message_class}