<commit_before>import copy
from urllib import unquote
import rawes
from django.conf import settings
from corehq.pillows.mappings.app_mapping import APP_INDEX
from corehq.pillows.mappings.case_mapping import CASE_INDEX
from corehq.pillows.mappings.domain_mapping import DOMAIN_INDEX
from corehq.pillows.mappings.group_mapping import GROUP_INDEX
from corehq.pillows.mappings.sms_mapping import SMS_INDEX
from corehq.pillows.mappings.tc_sms_mapping import TCSMS_INDEX
from corehq.pillows.mappings.user_mapping import USER_INDEX
from corehq.pillows.mappings.xform_mapping import XFORM_INDEX


def get_es(timeout=30):
    """
    Get a handle to the configured elastic search DB
    """
    return rawes.Elastic('%s:%s' % (settings.ELASTICSEARCH_HOST, 
                                    settings.ELASTICSEARCH_PORT),
                         timeout=timeout)


ES_URLS = {
    "forms": XFORM_INDEX + '/xform/_search',
    "cases": CASE_INDEX + '/case/_search',
    "users": USER_INDEX + '/user/_search',
    "domains": DOMAIN_INDEX + '/hqdomain/_search',
    "apps": APP_INDEX + '/app/_search',
    "groups": GROUP_INDEX + '/group/_search',
    "sms": SMS_INDEX + '/sms/_search',
    "tc_sms": TCSMS_INDEX + '/tc_sms/_search',
}

ADD_TO_ES_FILTER = {
    "forms": [
        {"term": {"doc_type": "xforminstance"}},
        {"not": {"missing": {"field": "xmlns"}}},
        {"not": {"missing": {"field": "form.meta.userID"}}},
    ],
    "users": [
        {"term": {"doc_type": "CommCareUser"}},
        {"term": {"base_doc": "couchuser"}},
        {"term": {"is_active": True}},
    ],
}

DATE_FIELDS = {
    "forms": "received_on",
    "cases": "opened_on",
    "users": "created_on",
    "sms": 'date',
    "tc_sms": 'date',
}

ES_MAX_CLAUSE_COUNT = 1024  #  this is what ES's maxClauseCount is currently set to,
                            #  can change this config value if we want to support querying over more domains


class ESError(Exception):
    pass


def run_query(url, q):
    return get_es().get(url, data=q)


def get_stats_data(domains, histo_type, datespan, interval="day", user_type_mobile=None):
    histo_data = dict([(d['display_name'],
                        es_histogram(histo_type, d["names"], datespan.startdate_display, datespan.enddate_display, interval=interval, user_type_mobile=user_type_mobile))
                        for d in domains])

    def _total_until_date(histo_type, doms=None):
        query = {"in": {"domain.exact": doms}} if doms is not None else {"match_all": {}}
        q = {
            "query": query,
            "filter": {
                "and": [
                    {"range": {DATE_FIELDS[histo_type]: {"lt": datespan.startdate_display}}},
                ],
            },
        }
        q["filter"]["and"].extend(ADD_TO_ES_FILTER.get(histo_type, [])[:])
        if user_type_mobile is not None:
            from corehq.apps.users.models import WebUser, CommCareUser
            q["filter"]["and"].extend([{
                "terms": {
                    "form.meta.userID":
                        [mobile_user._id for mobile_user in CommCareUser.all()
                         if mobile_user.doc_type == "CommCareUser"]
                        if user_type_mobile else
                        [web_user._id for web_user in WebUser.all()
                         if web_user.doc_type == "WebUser"]
                }
            }])

        return es_query(q=q, es_url=ES_URLS[histo_type], size=0)["hits"]["total"]

    return {
        'histo_data': histo_data,
        'initial_values': dict([(dom["display_name"],
                                 _total_until_date(histo_type, dom["names"])) for dom in domains]),
        'startdate': datespan.startdate_key_utc,
        'enddate': datespan.enddate_key_utc,
    }


def es_histogram(histo_type, domains=None, startdate=None, enddate=None, tz_diff=None, interval="day", q=None, user_type_mobile=None):
    q = q or {"query": {"match_all":{}}}

    if domains is not None:
        q["query"] = {"bool": {"must": [q["query"], {"in": {"domain.exact": domains}}]}}

    date_field = DATE_FIELDS[histo_type]

    q.update({
        "facets": {
            "histo": {
                "date_histogram": {
                    "field": date_field,
                    "interval": interval
                },
                "facet_filter": {
                    "and": [{
                        "range": {
                            date_field: {
                                "from": startdate,
                                "to": enddate
                            }}}]}}},
        "size": 0
    })

    if user_type_mobile is not None:
        from corehq.apps.users.models import WebUser, CommCareUser
        q["facets"]["histo"]["facet_filter"]["and"].extend([{
            "terms": {
                "form.meta.userID":
                    [mobile_user._id for mobile_user in CommCareUser.all()
                     if mobile_user.doc_type == "CommCareUser"]
                    if user_type_mobile else
                    [web_user._id for web_user in WebUser.all()
                     if web_user.doc_type == "WebUser"]
            }
        }])

    if tz_diff:
        q["facets"]["histo"]["date_histogram"]["time_zone"] = tz_diff

    q["facets"]["histo"]["facet_filter"]["and"].extend(ADD_TO_ES_FILTER.get(histo_type, []))

    es = get_es()
    ret_data = es.get(ES_URLS[histo_type], data=q)
    return ret_data["facets"]["histo"]["entries"]


SIZE_LIMIT = 1000000
def es_query(params=None, facets=None, terms=None, q=None, es_url=None, start_at=None, size=None, dict_only=False,
             fields=None, facet_size=None):
    if terms is None:
        terms = []
    if q is None:
        q = {}
    else:
        q = copy.deepcopy(q)
    if params is None:
        params = {}

    q["size"] = size if size is not None else q.get("size", SIZE_LIMIT)
    q["from"] = start_at or 0

    def get_or_init_anded_filter_from_query_dict(qdict):
        and_filter = qdict.get("filter", {}).pop("and", [])
        filter = qdict.pop("filter", None)
        if filter:<commit_msg>add user type filter only for the specified histogram types<commit_after>import copy
from urllib import unquote
import rawes
from django.conf import settings
from corehq.pillows.mappings.app_mapping import APP_INDEX
from corehq.pillows.mappings.case_mapping import CASE_INDEX
from corehq.pillows.mappings.domain_mapping import DOMAIN_INDEX
from corehq.pillows.mappings.group_mapping import GROUP_INDEX
from corehq.pillows.mappings.sms_mapping import SMS_INDEX
from corehq.pillows.mappings.tc_sms_mapping import TCSMS_INDEX
from corehq.pillows.mappings.user_mapping import USER_INDEX
from corehq.pillows.mappings.xform_mapping import XFORM_INDEX


def get_es(timeout=30):
    """
    Get a handle to the configured elastic search DB
    """
    return rawes.Elastic('%s:%s' % (settings.ELASTICSEARCH_HOST, 
                                    settings.ELASTICSEARCH_PORT),
                         timeout=timeout)


ES_URLS = {
    "forms": XFORM_INDEX + '/xform/_search',
    "cases": CASE_INDEX + '/case/_search',
    "users": USER_INDEX + '/user/_search',
    "domains": DOMAIN_INDEX + '/hqdomain/_search',
    "apps": APP_INDEX + '/app/_search',
    "groups": GROUP_INDEX + '/group/_search',
    "sms": SMS_INDEX + '/sms/_search',
    "tc_sms": TCSMS_INDEX + '/tc_sms/_search',
}

ADD_TO_ES_FILTER = {
    "forms": [
        {"term": {"doc_type": "xforminstance"}},
        {"not": {"missing": {"field": "xmlns"}}},
        {"not": {"missing": {"field": "form.meta.userID"}}},
    ],
    "users": [
        {"term": {"doc_type": "CommCareUser"}},
        {"term": {"base_doc": "couchuser"}},
        {"term": {"is_active": True}},
    ],
}

DATE_FIELDS = {
    "forms": "received_on",
    "cases": "opened_on",
    "users": "created_on",
    "sms": 'date',
    "tc_sms": 'date',
}

ES_MAX_CLAUSE_COUNT = 1024  #  this is what ES's maxClauseCount is currently set to,
                            #  can change this config value if we want to support querying over more domains


class ESError(Exception):
    pass


def run_query(url, q):
    return get_es().get(url, data=q)


def get_stats_data(domains, histo_type, datespan, interval="day", user_type_mobile=None):
    histo_data = dict([(d['display_name'],
                        es_histogram(histo_type, d["names"], datespan.startdate_display, datespan.enddate_display, interval=interval, user_type_mobile=user_type_mobile))
                        for d in domains])

    def _total_until_date(histo_type, doms=None):
        query = {"in": {"domain.exact": doms}} if doms is not None else {"match_all": {}}
        q = {
            "query": query,
            "filter": {
                "and": [
                    {"range": {DATE_FIELDS[histo_type]: {"lt": datespan.startdate_display}}},
                ],
            },
        }
        q["filter"]["and"].extend(ADD_TO_ES_FILTER.get(histo_type, [])[:])
        if user_type_mobile is not None:
            from corehq.apps.users.models import WebUser, CommCareUser
            if histo_type == 'forms':
                q["filter"]["and"].extend([{
                    "terms": {
                        "form.meta.userID":
                            [mobile_user._id for mobile_user in CommCareUser.all()
                             if mobile_user.doc_type == "CommCareUser"]
                            if user_type_mobile else
                            [web_user._id for web_user in WebUser.all()
                             if web_user.doc_type == "WebUser"]
                    }
                }])

        return es_query(q=q, es_url=ES_URLS[histo_type], size=0)["hits"]["total"]

    return {
        'histo_data': histo_data,
        'initial_values': dict([(dom["display_name"],
                                 _total_until_date(histo_type, dom["names"])) for dom in domains]),
        'startdate': datespan.startdate_key_utc,
        'enddate': datespan.enddate_key_utc,
    }


def es_histogram(histo_type, domains=None, startdate=None, enddate=None, tz_diff=None, interval="day", q=None, user_type_mobile=None):
    q = q or {"query": {"match_all":{}}}

    if domains is not None:
        q["query"] = {"bool": {"must": [q["query"], {"in": {"domain.exact": domains}}]}}

    date_field = DATE_FIELDS[histo_type]

    q.update({
        "facets": {
            "histo": {
                "date_histogram": {
                    "field": date_field,
                    "interval": interval
                },
                "facet_filter": {
                    "and": [{
                        "range": {
                            date_field: {
                                "from": startdate,
                                "to": enddate
                            }}}]}}},
        "size": 0
    })

    if user_type_mobile is not None:
        from corehq.apps.users.models import WebUser, CommCareUser
        if histo_type == 'forms':
            q["facets"]["histo"]["facet_filter"]["and"].extend([{
                "terms": {
                    "form.meta.userID":
                        [mobile_user._id for mobile_user in CommCareUser.all()
                         if mobile_user.doc_type == "CommCareUser"]
                        if user_type_mobile else
                        [web_user._id for web_user in WebUser.all()
                         if web_user.doc_type == "WebUser"]
                }
            }])

    if tz_diff:
        q["facets"]["histo"]["date_histogram"]["time_zone"] = tz_diff

    q["facets"]["histo"]["facet_filter"]["and"].extend(ADD_TO_ES_FILTER.get(histo_type, []))

    es = get_es()
    ret_data = es.get(ES_URLS[histo_type], data=q)
    return ret_data["facets"]["histo"]["entries"]


SIZE_LIMIT = 1000000
def es_query(params=None, facets=None, terms=None, q=None, es_url=None, start_at=None, size=None, dict_only=False,
             fields=None, facet_size=None):
    if terms is None:
        terms = []
    if q is None:
        q = {}
    else:
        q = copy.deepcopy(q)
    if params is None:
        params = {}

    q["size"] = size if size is not None else q.get("size", SIZE_LIMIT)
    q["from"] = start_at or 0

    def get_or_init_anded_filter_from_query_dict(qdict):
        and_filter = qdict.get("filter", {}).pop("and", [])
        filter = qdict.pop("filter", None)
        if filter:
{'fraction_non_alphanumeric': 0.1374415228262623, 'fraction_numerical': 0.0025810614615260526, 'mean_word_length': 2.1662410623084782, 'pattern_counts': {'":': 116, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6297224', 'n_tokens_mistral': 4047, 'n_tokens_neox': 3900, 'n_words': 796}