[关闭]
@heavysheep 2020-11-06T12:00:51.000000Z 字数 8905 阅读 551

ES封装

代码

  1. class ESQuery(object):
  2.     def __init__(self, index: str = None, doc_type: str = None, bool: dict = None, collapse: dict = None,
  3.                  dsl: dict = None):
  4.         """
  5.         ES 查询dsl封装
  6.         :param index: ES index
  7.         :param doc_type: ES doc_type(如有)
  8.         :param bool: bool dsl
  9.         :param collapse: collapse dsl.折叠dsl
  10.         :param dsl: source dsl.传入时直接生成dsl，和bool collapse等解析参数有所出入
  12.         example:
  13.             # [build condition]
  14.             filter_condition = Query.term(field="is_match", value=True)
  15.             agg_condition = [Aggs.terms(name="group", field="website.keyword", size=500)]
  17.             # [input condition]
  18.             query = ESQuery(index="index_1", doc_type="type_1", bool=bools(filter=filter_condition))
  19.             query.aggs(condition=agg_condition)
  21.             # [modify attr]
  22.             query.size(num=5)
  23.             query.source(include=["title", "content", "website"])
  24.             query.sort(condition="data_time:desc")
  26.             # [get result]
  27.             query_result = query.search()
  28.         """
  29.         if dsl and bool:
  30.             raise ValueError("不允许同时传入'dsl' 'bool'参数")
  31.         if dsl and collapse:
  32.             raise ValueError("不允许同时传入'dsl' 'collapse'参数")
  34.         self.index = index
  35.         self.doc_type = doc_type
  37.         self.params = {}
  39.         self.bool = bool
  41.         self.dsl = dsl if dsl is not None else {}
  42.         if bool:
  43.             self.dsl["query"] = bool
  44.         if collapse:
  45.             self.dsl["collapse"] = collapse
  47.     def bools(self, must: (dict, list) = None, filter: (dict, list) = None, should: (dict, list) = None,
  48.               must_not: (dict, list) = None):
  49.         """
  50.         bool 条件
  51.         :param must: must条件
  52.         :param filter: filter条件
  53.         :param should: should条件
  54.         :param must_not: must not条件
  55.         """
  56.         for name, condition_info in (("must", must), ("filter", filter), ("should", should), ("must_not", must_not)):
  57.             if self.dsl.get("query") is None:
  58.                 self.dsl["query"] = {"bool": {}}
  59.             if self.dsl["query"].get("bool") is None:
  60.                 self.dsl["query"]["bool"] = {}
  61.             if condition_info is not None:
  62.                 condition = _parser_condition(condition_info)
  63.                 if name not in self.dsl["query"]["bool"]:
  64.                     self.dsl["query"]["bool"][name] = condition
  65.                 else:
  66.                     self.dsl["query"]["bool"][name].extend(condition)
  68.     def aggs(self, condition):
  69.         if isinstance(condition, list):
  70.             condition = {k: v for c in condition for k, v in c.items()}
  71.         self.dsl["aggs"] = condition
  73.     def range(self, range_condition: dict):
  74.         """
  75.         区间条件配置方法 对应DSL _source
  76.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/query-dsl-range-query.html]
  77.         :param range_condition: 区间条件
  78.         """
  79.         if self.dsl.get("query") is None:
  80.             self.dsl["query"] = {"bool": {}}
  81.         if self.dsl["query"].get("bool") is None:
  82.             self.dsl["query"]["bool"] = {}
  83.         if self.dsl["query"]["bool"].get("must") is None:
  84.             self.dsl["query"]["bool"]["must"] = range_condition
  85.             return
  87.         for i, condition in enumerate(self.dsl["query"]["bool"]["must"]):
  88.             if "range" in condition.keys():
  89.                 self.dsl["query"]["bool"]["must"][i] = range_condition
  90.                 return
  91.         self.dsl["query"]["bool"]["must"].append(range_condition)
  93.     def source(self, include: (str, list) = None, exclude: (str, list) = None):
  94.         """
  95.         字段配置方法 对应DSL _source
  96.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/search-request-source-filtering.html]
  97.         :param include: 保留字段
  98.         :param exclude: 不保留字段
  99.         """
  100.         self.dsl["_source"] = {}
  101.         if include:
  102.             self.dsl["_source"]["includes"] = include
  103.         if exclude:
  104.             self.dsl["_source"]["excludes"] = exclude
  106.     def pagination(self, page: int = 1, per_page: int = 20):
  107.         """
  108.         分页方法 对应DSL from/size
  109.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/search-request-from-size.html]
  110.         :param page: 页数，从1开始，默认1
  111.         :param per_page: 页容量，默认20
  112.         """
  113.         if page is None:
  114.             page = 1
  115.         if per_page is None:
  116.             per_page = 20
  118.         from_ = (page - 1) * per_page
  119.         self.dsl["from"] = from_
  120.         self.dsl["size"] = per_page
  122.     def sort(self, condition: (str, dict)):
  123.         """
  124.         排序方法，对应DSL sort
  125.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/search-request-sort.html]
  126.         注意此方法是对query对象的排序方法
  127.         :param condition: 排序条件
  128.         """
  129.         if isinstance(condition, dict):
  130.             condition = ",".join(("{}:{}".format(c_k, c_v) for c_k, c_v in condition.items()))
  131.         self.params["sort"] = condition
  133.     def size(self, num: int):
  134.         """
  135.         单容量方法，对应DSL size
  136.         注意和pagination有冲突
  137.         :param num: 容量
  138.         """
  139.         self.dsl.update({"size": num})
  141.     def scroll(self, scroll: str):
  142.         """
  143.         滚动方法 对应DSL scroll
  144.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/search-request-scroll.html]
  145.         :param scroll: 滚动ID
  146.         """
  147.         self.params.update({"scroll": scroll})
  149.     def search(self, debug=False):
  150.         """
  151.         search方法 对应ES request body search方法
  152.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/search-uri-request.html]
  153.         :param debug: 是否打印(info日志)DSL request body
  154.         """
  155.         if debug:
  156.             logger.info(json.dumps(self.dsl, ensure_ascii=False, indent=4))
  157.         if self.dsl.get("size") is None:
  158.             self.size(50)
  159.         result = es.search(body=self.dsl, index=self.index, doc_type=self.doc_type, params=self.params)
  160.         result.update({"pager_info": self._get_page_info(result)})
  161.         return result
  163.     def _get_page_info(self, result: dict):
  164.         """
  165.         分页信息方法，用于返回分页相关情况说明
  166.         :param result: ES查询结果
  167.         :return: 分页相关信息
  168.         """
  169.         pager_info = {}
  171.         if self.dsl.get("from") is not None:
  172.             pager_info["page"] = self.dsl["from"]
  173.             pager_info["per_page"] = self.dsl["size"]
  174.             pager_info["total"] = result["hits"]["total"]
  175.             pager_info["pages"] = math.ceil(pager_info["total"] / pager_info["per_page"])
  177.         if self.params.get("sort") is not None:
  178.             pager_info["order_by"] = self.params["sort"]
  180.         return pager_info
  182.     def clone(self):
  183.         return ESQuery(index=self.index, doc_type=self.doc_type, dsl=deepcopy(self.dsl))
  186. class Aggs(object):
  187.     @staticmethod
  188.     def terms(name: str, field: str, size: int = None, include: (str, dict, list) = None, exclude: (str, list) = None,
  189.               order: dict = None, child: (dict, list) = None):
  190.         """
  191.         基础聚类方法
  192.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/query-dsl-term-query.html]
  193.         :param order:
  194.         :param name: 聚类名称
  195.         :param field: 聚类字段名
  196.         :param size: 聚类返回前size个
  197.         :param include: 计入条件(匹配条件则计入聚类)
  198.         :param exclude: 不计入条件(匹配条件则不计入聚类)
  199.         :param child: aggs子条件
  200.         """
  202.         field_condition = {"field": field}
  204.         if size is not None:
  205.             field_condition.update({"size": size})
  206.         if include is not None:
  207.             field_condition.update({"include": include})
  208.         if exclude is not None:
  209.             field_condition.update({"exclude": exclude})
  210.         if order is not None:
  211.             field_condition.update({"order": order})
  213.         if child is not None:
  214.             aggs = {}
  215.             if isinstance(child, dict):
  216.                 aggs.update(child)
  217.             else:
  218.                 for c in child:
  219.                     aggs.update(c)
  220.             return {name: {"terms": field_condition, "aggs": aggs}}
  221.         return {name: {"terms": field_condition}}
  223.     @staticmethod
  224.     def date_histogram(name: str, field: str, interval: str, format="yyyy-MM-dd", start: str = None, end: str = None):
  225.         """
  226.         时间聚类方法
  228.         当有start和end时，强制返回start和end区间所有时间
  229.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/search-aggregations-bucket-datehistogram-aggregation.html]
  230.         :param name: 聚类名称
  231.         :param field: 聚类字段名
  232.         :param interval: 聚类周期
  233.         :param format: 返回时间格式
  234.         :param start: 返回时间区间起始
  235.         :param end: 返回时间区间结束
  236.         :return 返回本对象，以支持进一步聚类
  237.         """
  238.         field_condition = {"field": field, "interval": interval, "format": format}
  239.         if start and end:
  240.             field_condition.update({"extended_bounds": {"min": start, "max": end}})
  242.         return {name: {"date_histogram": field_condition}}
  244.     @staticmethod
  245.     def cardinality(name: str, field: str):
  246.         return {name: {"cardinality": {"field": field}}}
  248.     @staticmethod
  249.     def top_hits(name: str, includes: list = None, sort: list = None, size: int = None):
  250.         """
  251.         top hits必须是terms的一个child
  252.         """
  253.         field_condition = {}
  254.         if sort:
  255.             field_condition["sort"] = sort
  256.         if includes is not None:
  257.             field_condition["_source"] = {"includes": includes}
  258.         if size is not None:
  259.             field_condition["size"] = size
  260.         return {name: {"top_hits": field_condition}}
  262.     @staticmethod
  263.     def trans_terms_order(order_by: str):
  264.         order_condition = []
  265.         for oc in order_by.split(","):
  266.             oc_split = oc.split(":")
  267.             if len(oc_split) == 1:
  268.                 order_condition.append({oc_split[0]: "asc"})
  269.             elif len(oc_split) == 2:
  270.                 order_condition.append({oc_split[0]: oc_split[1]})
  271.             else:
  272.                 raise KeyError("order_by参数解析失败")
  273.         return order_condition
  276. class Query(object):
  277.     @staticmethod
  278.     def terms(field: str, value: list):
  279.         return {"constant_score": {"filter": {"terms": {field: value}}}}
  281.     @staticmethod
  282.     def match(field: str, value: (str, int), operator: str = "and"):
  283.         """
  284.         查询ES中field包含value的数据
  285.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/query-dsl-match-query.html]
  286.         :param field: 字段名称
  287.         :param value: 字段值
  288.         :param operator: 或or/与and条件
  289.         """
  290.         return {"match": {field: {"query": value, "operator": operator}}}
  292.     @staticmethod
  293.     def match_phrase(field: str, value: (str, int)):
  294.         """
  295.         短语匹配，对应DSL Match Phrase Query
  296.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/query-dsl-match-query-phrase.html]
  297.         :param field: 字段名
  298.         :param value: 字段值
  299.         :return:
  300.         """
  301.         return {"match_phrase": {field: value}}
  303.     @staticmethod
  304.     def term(field: str, value: (str, int)):
  305.         """
  306.         数据筛选，对应DSL term
  307.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/query-dsl-term-query.html]
  308.         :param field: 字段名
  309.         :param value: 字段值
  310.         """
  311.         return {"term": {field: value}}
  313.     @staticmethod
  314.     def range(field: str, start: str, end: str, left: str = "gte", right: str = "lt"):
  315.         """
  316.         筛选ES中符合range条件的数据
  318.         left/right支持:
  319.             gte: 大于等于
  320.             gt: 大于
  321.             lte: 小于等于
  322.             lt: 小于
  323.         ES文档: [https://www.elastic.co/guide/en/elasticsearch/reference/5.5/query-dsl-range-query.html]
  325.         :param field: 字段名称
  326.         :param start: 起始值
  327.         :param end: 结束值
  328.         :param left: 起始值匹配条件
  329.         :param right: 结束值匹配条件
  330.         """
  331.         return {"range": {field: {left: start, right: end}}}
  333.     @staticmethod
  334.     def collapse(field: str, name: str = None, size: int = None, sort: list = None):
  335.         field_condition = {"field": field}
  336.         if name is not None or size is not None or sort is not None:
  337.             inner_hits = {}
  338.             if name:
  339.                 inner_hits["name"] = name
  340.             if size is not None:
  341.                 inner_hits["size"] = size
  342.             if sort:
  343.                 inner_hits["sort"] = sort
  344.             field_condition["inner_hits"] = inner_hits
  345.         return field_condition
  348. def bools(must: (dict, list) = None, filter: (dict, list) = None, should: (dict, list) = None,
  349.           must_not: (dict, list) = None):
  350.     query = {}
  351.     if must:
  352.         query["must"] = _parser_condition(must)
  353.     if filter:
  354.         query["filter"] = _parser_condition(filter)
  355.     if should:
  356.         query["should"] = _parser_condition(should)
  357.     if must_not:
  358.         query["must_not"] = _parser_condition(must_not)
  359.     if not query:
  360.         return {}
  361.     return {"bool": query}
  364. def _parser_condition(condition):
  365.     if condition:
  366.         if isinstance(condition, dict):
  367.             return [condition]
  368.         else:
  369.             return [c for c in condition if c]
添加新批注
在作者公开此批注前，只有你和作者可见。
回复批注