neptune/core/views/dynamic/search.py

import json

from django.conf import settings
from django.contrib.auth.mixins import LoginRequiredMixin
from django.http import HttpResponse, HttpResponseForbidden, JsonResponse
from django.shortcuts import render
from django.views import View

from core.lib.opensearch import initialise_opensearch, run_main_query
from core.lib.threshold import annotate_num_chans, annotate_num_users, annotate_online

client = initialise_opensearch()


def query_results(request, post_params, api=False):
    fields = None
    if "fields" in request.POST:
        fields = request.POST.getlist("fields")
    if "size" in request.POST:
        size = request.POST["size"]
    if "query" in request.POST:
        query = request.POST["query"]
        results = run_main_query(client, request.user, query, fields, size)
        if not results:
            return False
        results_parsed = []
        if "hits" in results.keys():
            if "hits" in results["hits"]:
                for item in results["hits"]["hits"]:
                    element = item["_source"]
                    element["id"] = item["_id"]
                    ts = element["ts"]
                    ts_spl = ts.split("T")
                    date = ts_spl[0]
                    time = ts_spl[1]
                    element["date"] = date
                    element["time"] = time
                    results_parsed.append(element)
        # Figure out items with net (not discord)
        nets = set()
        for x in results_parsed:
            if "net" in x:
                nets.add(x["net"])

        for net in nets:
            # Annotate the online attribute from Threshold
            online_info = annotate_online(
                net, [x["nick"] for x in results_parsed if x["src"] == "irc"]
            )
            # Annotate the number of users in the channel
            num_users = annotate_num_users(
                net, [x["channel"] for x in results_parsed if x["src"] == "irc"]
            )
            # Annotate the number channels the user is on
            num_chans = annotate_num_chans(
                net, [x["nick"] for x in results_parsed if x["src"] == "irc"]
            )
            for item in results_parsed:
                if item["nick"] in online_info:
                    item["online"] = online_info[item["nick"]]
                if item["channel"] in num_users:
                    item["num_users"] = num_users[item["channel"]]
                if item["nick"] in num_chans:
                    item["num_chans"] = num_chans[item["nick"]]

        context = {
            "query": query,
            "results": results_parsed,
            "card": results["hits"]["total"]["value"],
            "took": results["took"],
            "redacted": results["redacted"],
            "exemption": results["exemption"],
            "fields": settings.OPENSEARCH_MAIN_SEARCH_FIELDS,
            "sizes": settings.OPENSEARCH_MAIN_SIZES,
            "timescales": settings.OPENSEARCH_MAIN_TIMESCALES,
        }
        return context


class Search(LoginRequiredMixin, View):
    # parser_classes = [JSONParser]
    template_name = "ui/results.html"
    plan_name = "drilldown"

    def post(self, request):
        if not request.user.has_plan(self.plan_name):
            return HttpResponseForbidden()

        context = query_results(request, request.POST)
        if not context:
            return HttpResponseForbidden()
        context["data"] = json.dumps(
            [
                {
                    "text": item.get("msg", None) or item.get("id"),
                    "nick": item.get("nick", None),
                    "value": item.get("sentiment", None) or None,
                    "date": item.get("ts"),
                }
                for item in context["results"]
            ]
        )
        if context:
            return render(request, self.template_name, context)
        else:
            return HttpResponse("No results")


class APISearch(LoginRequiredMixin, View):
    # parser_classes = [JSONParser]
    template_name = "ui/results.html"
    plan_name = "drilldown"

    def post(self, request):
        if not request.user.has_plan(self.plan_name):
            return JsonResponse({"success": False})

        context = query_results(request, request.POST)
        return JsonResponse(context)