es.py 1.24 KB
from elasticsearch import Elasticsearch as Es
from elasticsearch import helpers


def es_index_adapt(index_prefix, doc_type, rw=None):
    """get the adapted index name
    """
    assert rw in [None, "read", "write"]
    index = "-".join((index_prefix, doc_type))
    if rw:
        index = "-".join((index, rw))
    return index


def get_es():
    init_args = {
        "sniff_on_start": False,
        "sniff_on_connection_fail": False,
    }
    new_hosts = [{
        "host": "172.16.31.17",
        "port": 9000,
    }, {
        "host": "172.16.31.11",
        "port": 9000,
    }, {
        "host": "172.16.31.13",
        "port": 9000,
    }]
    new_es = Es(hosts=new_hosts, **init_args)
    return new_es


def es_query(doc, body, offset=0, size=100, es=None, rw="read"):
    if es is None:
        es = get_es()
    index = es_index_adapt(index_prefix="gm-dbmw", doc_type=doc, rw=rw)
    res = es.search(index=index, timeout="10s", body=body, from_=offset, size=size)
    return res


def es_scan(doc, body, es=None, rw="read"):
    if es is None:
        es = get_es()
    index = es_index_adapt(index_prefix="gm-dbmw", doc_type=doc, rw=rw)
    return helpers.scan(es, index=index, query=body, request_timeout=100, scroll="300m", raise_on_error=False)