connectionES.py 1.1 KB
Newer Older
ZHOU Shihang's avatar
ZHOU Shihang committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40
from elasticsearch import Elasticsearch


def get_basic_query():
    return {
        "query": {
            "bool": {
                "filter": []
            }
        }
    }


def get_summary(query, full_endpoint):
    es = Elasticsearch(
        [full_endpoint],
        verify_certs=False
    )

    page_size = 10
    query = es.search(index="lsst-batch-summary", body=query, scroll='5m', size=page_size)
    results = list(map(lambda x: x["_source"], query['hits']['hits']))
    total = query['hits']['total']["value"]
    print(total)
    scroll_id = query['_scroll_id']
    for i in range(0, int(total / page_size) + 1):
        query_scroll = es.scroll(scroll_id=scroll_id, scroll='5m')['hits']['hits']
        data = list(map(lambda x: x["_source"], query_scroll))
        results += data
    return results


def add_single_filter(q, field, value):
    q["query"]["bool"]["filter"].append({"term": {str(field): value}})
    return q


def add_time_range_filter(q, start, end):
    q["query"]["bool"]["filter"].append({"range": {"Timestamp": {"gte": start, "lte": end, "format": "epoch_millis"}}})
    return q