neptune/core/lib/meta.py

from math import ceil

from django.conf import settings
from numpy import array_split

from core.lib.opensearch import client, run_main_query


def construct_query(net, nicks):
    # Construct the query
    query_nicks = [{"match": {"nick": x}} for x in nicks]
    query_should = query_nicks
    # Get the initial query
    query = {
        "size": settings.META_QUERY_SIZE,
        "query": {
            "bool": {
                "must": [
                    {"match": {"net": net}},
                    {"match": {"type": "who"}},
                    {
                        "bool": {
                            "should": query_should,
                        }
                    },
                ]
            }
        },
    }
    return query


def get_meta(request, net, nicks, iter=True):
    """
    Get all related nicknames of the given nickname by tracking nickname changes.
    """

    # Split query into chunks
    split_nicks = array_split(nicks, ceil(len(nicks) / settings.META_MAX_CHUNK_SIZE))
    meta = []
    for nicks_chunked in split_nicks:
        if len(nicks_chunked) == 0:
            break
        meta_tmp = []
        query = construct_query(net, nicks_chunked)
        results = run_main_query(
            client,
            request.user,
            query,
            custom_query=True,
            index=settings.OPENSEARCH_INDEX_META,
        )
        if "hits" in results.keys():
            if "hits" in results["hits"]:
                for item in results["hits"]["hits"]:
                    element = item["_source"]
                    element["id"] = item["_id"]

                    # Split the timestamp into date and time
                    ts = element["ts"]
                    ts_spl = ts.split("T")
                    date = ts_spl[0]
                    time = ts_spl[1]
                    element["date"] = date
                    element["time"] = time
                    meta_tmp.append(element)
        for x in meta_tmp:
            if x not in meta:
                meta.append(x)

    # Run the search again, passing in all the users we found

    # Nicknames we find from the repeated search
    # nicks_searched = []
    # if iter:
    #     nicks_l2 = []
    #     loop = 0
    #     while loop < settings.META_MAX_ITERATIONS:
    #         loop += 1
    #         nicks_not_searched = [x for x in nicks if x not in nicks_searched]
    #         nicks_l2 = get_nicks(request, net, nicks, False)

    #         # Add all the nicks we just searched for to the list
    #         for x in nicks_not_searched:
    #             if x not in nicks_not_searched:
    #                 nicks_searched.append(x)

    #         # If all of the nicks we received now, we already know about
    #         if set(nicks_l2).issubset(set(nicks)):
    #             break
    #         for x in nicks_l2:
    #             if x not in nicks:
    #                 nicks.append(x)

    return meta