# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one # or more contributor license agreements. Licensed under the Elastic License # 2.0; you may not use this file except in compliance with the Elastic License # 2.0. from collections import defaultdict from eql import Walker from .errors import KqlCompileError def boolean(**kwargs): """Wrap a query in a boolean term and optimize while building.""" assert len(kwargs) == 1 [(boolean_type, children)] = kwargs.items() if not isinstance(children, list): children = [children] dsl = defaultdict(list) if boolean_type in ("must", "filter"): # safe to convert and(and(x), y) -> and(x, y) for child in children: if list(child) == ["bool"]: for child_type, child_terms in child["bool"].items(): if child_type in ("must", "filter"): dsl[child_type].extend(child_terms) elif child_type == "should": if "should" not in dsl: dsl[child_type].extend(child_terms) else: dsl[boolean_type].append(boolean(should=child_terms)) elif child_type == "must_not": dsl[child_type].extend(child_terms) elif child_type != "minimum_should_match": raise ValueError("Unknown term {}: {}".format(child_type, child_terms)) else: dsl[boolean_type].append(child) elif boolean_type == "should": # can flatten `should` of `should` for child in children: if list(child) == ["bool"] and set(child["bool"]).issubset({"should", "minimum_should_match"}): dsl["should"].extend(child["bool"]["should"]) else: dsl[boolean_type].append(child) elif boolean_type == "must_not" and len(children) == 1: # must_not: [{bool: {must: x}}] -> {must_not: x} child = children[0] if list(child) == ["bool"] and list(child["bool"]) in (["filter"], ["must"]): negated, = child["bool"].values() dsl = {"must_not": negated} else: dsl = {"must_not": children} else: dsl = dict(kwargs) if "should" in dsl: dsl.update(minimum_should_match=1) dsl = {"bool": dict(dsl)} return dsl class ToDsl(Walker): def _walk_default(self, node, *args, **kwargs): raise KqlCompileError("Unable to convert {}".format(node)) def _walk_exists(self, _): return lambda field: {"exists": {"field": field}} def _walk_wildcard(self, tree): return lambda field: {"query_string": {"fields": [field], "query": tree.value}} def _walk_value(self, tree): return lambda field: {"match": {field: tree.value}} def _walk_field(self, field): return field.name def _walk_field_range(self, tree): operator_map = {"<": "lt", "<=": "lte", ">=": "gte", ">": "gt"} field = self.walk(tree.field) return {"range": {field: {operator_map[tree.operator]: tree.value.value}}} def _walk_not_expr(self, tree): return boolean(must_not=[self.walk(tree.expr)]) def _walk_and_expr(self, tree): return boolean(filter=[self.walk(node) for node in tree.items]) def _walk_or_expr(self, tree): return boolean(should=[self.walk(node) for node in tree.items]) def _walk_and_values(self, tree): children = [self.walk(node) for node in tree.items] return lambda field: boolean(filter=[child(field) for child in children]) def _walk_or_values(self, tree): children = [self.walk(node) for node in tree.items] return lambda field: boolean(should=[child(field) for child in children]) def _walk_not_value(self, tree): child = self.walk(tree.value) return lambda field: boolean(must_not=[child(field)]) def _walk_field_comparison(self, tree): field = self.walk(tree.field) value_fn = self.walk(tree.value) return value_fn(field) @classmethod def convert(cls, tree): return boolean(filter=[cls().walk(tree)])