Python pyeqs.QuerySet类代码示例

def test_post_query_actions(context):
    """
    Search with match_all query with post query actions
    """
    # When create a queryset
    t = QuerySet("localhost", index="foo")

    # And there are records
    add_document("foo", {"bar": 1})
    add_document("foo", {"bar": 2})
    add_document("foo", {"bar": 3})

    # And I have a post query action
    global my_global_var
    my_global_var = 1

    def action(self, results, start, stop):
        global my_global_var
        my_global_var += 1

    t.post_query_actions(action)

    # And I do a search
    t.order_by(Sort("bar", order="asc"))
    results = t[0:10]

    # Then I get a the expected results
    len(results).should.equal(3)
    results[0]["_source"]["bar"].should.equal(1)
    results[1]["_source"]["bar"].should.equal(2)
    results[2]["_source"]["bar"].should.equal(3)
    my_global_var.should.equal(2)

def test_create_queryset_with_filters():
    """
    Create QuerySet with Multiple Filters
    """
    # When create a query block
    t = QuerySet("foobar")

    # And I add a filter
    t.filter(Term("foo", "bar"))
    t.filter(Term("foobar", "foobar"))

    # Then I see the appropriate JSON
    results = {
        "query": {
            "filtered": {
                "query": {"match_all": {}},
                "filter": {
                    "and": [
                        {
                            "term": {
                                "foo": "bar"
                            }
                        },
                        {
                            "term": {
                                "foobar": "foobar"
                            }
                        }
                    ]
                }
            }
        }
    }

    homogeneous(t._query, results)

def test_search_multiple_scoring(context):
    """
    Search with custom scoring and params
    """
    # When create a query block
    t = QuerySet("localhost", index="foo")

    # And there are records
    add_document("foo", {"bar": 1, "baz": 4})
    add_document("foo", {"bar": 1})

    # And I add scoring with params
    score = ScriptScore("s = custom_param + doc['bar'].value", params={"custom_param": 1})
    t.score(score)

    boost = {
        "boost_factor": "10",
        "filter": Exists("baz")
    }
    t.score(boost)
    results = t[0:10]

    # Then my results are scored correctly
    len(results).should.equal(2)
    results[0]["_source"]["baz"].should.equal(4)
    ("baz" in results[1]["_source"].keys()).should.be.false

def test_create_queryset_with_multiple_sorting():
    """
    Create QuerySet with Multiple Sorting
    """
    # When create a query block
    t = QuerySet("foobar")

    # And I add sorting
    s = Sort("_id", order="asc")
    t.order_by(s)

    ss = Sort("_id", order="desc")
    t.order_by(ss)

    # Then I see the appropriate JSON
    results = {
        "sort": [
            {
                "_id": {
                    "order": "asc"
                }
            },
            {
                "_id": {
                    "order": "desc"
                }
            }
        ],
        "query": {
            "match_all": {}
        }
    }

    homogeneous(t._query, results)

def test_wrappers(context):
    """
    Search with wrapped match_all query
    """
    # When create a queryset
    t = QuerySet("localhost", index="foo")

    # And there are records
    add_document("foo", {"bar": 1})
    add_document("foo", {"bar": 2})
    add_document("foo", {"bar": 3})

    # And I do a search
    def wrapper_function(search_results):
        return map(lambda x: x["_source"]["bar"] + 1, search_results)

    t.wrappers(wrapper_function)
    t.order_by(Sort("bar", order="asc"))
    results = t[0:10]

    # Then I get a the expected results
    len(results).should.equal(3)
    results[0].should.equal(2)
    results[1].should.equal(3)
    results[2].should.equal(4)

def test_create_queryset_with_scoring():
    """
    Create QuerySet with Scoring, Minimum Score and Track Scores
    """
    # When create a query block
    t = QuerySet("foobar")

    # And I add scoring
    s = ScriptScore("foo = 0.0")
    t.score(s, min_score=0, track_scores=True)

    # Then I see the appropriate JSON
    results = {
        "min_score": 0,
        "track_scores": True,
        "query": {
            "function_score": {
                "functions": [
                    {
                        "script_score": {
                            "script": "foo = 0.0"
                        }
                    }
                ],
                "query": {"match_all": {}},
                "boost_mode": "replace",
                "score_mode": "multiply"
            }
        }
    }

    homogeneous(t._query, results)

def test_search_global_aggregations(context):
    """
    Search with global aggregations
    """
    # With a specific query
    # q = QueryBuilder(query_string=QueryString(query={"match": {"foo_attr": "yes"}}))

    # When create a query block
    t = QuerySet("localhost", index="foo")

    # And there are records
    add_document("foo", {"bar": "baz", "foo_attr": "yes"})
    add_document("foo", {"bar": "bazbaz", "foo_attr": "no"})
    add_document("foo", {"bar": "bazbaz"})

    # And I do a global
    t.aggregate(aggregation=Aggregations("foo_attrs", "bar", "terms", global_name="all_bar"))
    t[0:10]

    # I get the expected results
    t.aggregations().should.have.key("all_bar").being.equal({
        u'foo_attrs': {u'buckets': [
            {u'key': u'bazbaz', u'doc_count': 2},
            {u'key': u'baz', u'doc_count': 1}
        ]},
        u'doc_count': 3})

def test_search_nested_terms_aggregations_with_size(context):
    """
    Search with nested terms aggregation and a specified size
    """
    # When create a query block
    t = QuerySet("localhost", index="foo")

    # And there are nested records
    add_document("foo", {"child": [{"stuff": "yep", "bazbaz": "type0"}], "foo": "foo"})
    add_document("foo", {"child": [{"stuff": "yep", "bazbaz": "type0"}], "foo": "foo"})
    add_document("foo", {"child": [{"stuff": "nope", "bazbaz": "type1"}], "foo": "foofoo"})
    add_document("foo", {"child": [{"stuff": "nope", "bazbaz": "type2"}], "foo": "foofoo"})

    # And I do a nested
    t.aggregate(aggregation=Aggregations("baz_types", "bazbaz", "terms",
                                         nested_path="child", size=1))
    t[0:10]

    # The I get the expected results
    t.aggregations().should.have.key("child").being.equal({
        u'baz_types': {
            u'buckets': [{u'doc_count': 2, u'key': u'type0'}]
        },
        u'doc_count': 4
    })

def test_queryset_iteration_with_no_results():
    """
    Fetch results with QuerySet via __iter__ with no results
    """

    # When I create a query block
    t = QuerySet("foobar", index="bar")
    wrapper = lambda y: list(map(lambda x: x['_id'], y))
    t.wrappers(wrapper)

    # And I have no records
    response = {
       "took": 12,
       "hits": {
          "total": 0,
          "max_score": 0,
          "hits": []
       }
    }
    httpretty.register_uri(httpretty.GET, "http://foobar:9200/bar/_search",
                       body=json.dumps(response),
                       content_type="application/json")

    results = []
    for result in t:
        results.append(result)
    len(results).should.equal(0)
    t.count().should.equal(0)

def test_create_queryset_with_multiple_scoring():
    """
    Create QuerySet with Multiple Scoring
    """
    # When create a query block
    t = QuerySet("foobar")

    # And I add scoring
    s = ScriptScore("foo = 0.0")
    t.score(s)

    # And I add more scoring
    ss = ScriptScore("foo = 1.0")
    t.score(ss)

    # Then I see the appropriate JSON
    results = {
        "query": {
            "function_score": {
                "query": {"match_all": {}},
                "script_score": {
                    "script": "foo = 1.0"
                },
                "boost_mode": "replace"
            }
        }
    }

    homogeneous(t._query, results)

def test_queryset_count():
    """
    Get QuerySet Count
    """
    # When I create a query block
    t = QuerySet("foobar")

    # Then I get an appropriate Count
    t.count().should.equal(None)

def test_queryset_max_score():
    """
    Get QuerySet Max Score
    """
    # When I create a query block
    t = QuerySet("foobar")

    # Then I get an appropriate max score
    t.max_score().should.equal(None)

def test_queryset_getitem_with_post_query_action():
    """
    Fetch from QuerySet with __getitem__ and post query action
    """
    # When I create a query block
    t = QuerySet("localhost", index="bar")

    # And I have a post query action
    global my_global_var
    my_global_var = 1

    def action(self, results, start, stop):
        global my_global_var
        my_global_var += 1

    t.post_query_actions(action)

    # And I have records
    response = {
        "took": 12,
        "timed_out": False,
        "_shards": {
            "total": 5,
            "successful": 5,
            "failed": 0
        },
        "hits": {
            "total": 1,
            "max_score": 10,
            "hits": [
                {
                    "_index": "bar",
                    "_type": "baz",
                    "_id": "1",
                    "_score": 10,
                    "_source": {
                        "foo": "bar"
                    },
                    "sort": [
                        1395687078000
                    ]
                }
            ]
        }
    }
    httpretty.register_uri(httpretty.GET, "http://localhost:9200/bar/_search",
                       body=json.dumps(response),
                       content_type="application/json")

    results = t[0:1]
    len(results).should.equal(1)
    t.count().should.equal(1)

    # Then I see the correct results
    results[0]['_id'].should.equal('1')
    my_global_var.should.equal(2)

def test_search_with_filter_and_scoring_and_sorting_and_fields(context):
    """
    Search with match_all query, filter, scoring, sorting, and fields
    """
    # When create a queryset
    t = QuerySet("localhost", index="foo")

    # And there are records
    add_document("foo", {"bar": "baz", "scoring_field": 0, "sorting_field": 30})
    add_document("foo", {"bar": "baz", "scoring_field": 1, "sorting_field": 20})
    add_document("foo", {"bar": "baz", "scoring_field": 2, "sorting_field": 10})
    add_document("foo", {"bar": "bazbaz", "scoring_field": 3, "sorting_field": 0})

    # And I do a search
    t.filter(Term("bar", "baz"))
    score = ScriptScore("final_score = 0 + doc['scoring_field'].value;")
    t.score(score)
    sorting = Sort("sorting_field", order="desc")
    t.order_by(sorting)
    t.only(["bar"])
    results = t[0:10]

    # Then I get a the expected results
    len(results).should.equal(3)
    results[0]['fields'].should.equal({"bar": ["baz"]})
    results[1]['fields'].should.equal({"bar": ["baz"]})
    results[2]['fields'].should.equal({"bar": ["baz"]})

def test_search_filtered_aggregations(context):
    """
    Search with filtered aggregations
    """
    # When create a query block
    t = QuerySet("localhost", index="foo")

    # And there are records
    add_document("foo", {"bar": 0, "foo": "baz"})
    add_document("foo", {"bar": 0, "foo": "baz"})
    add_document("foo", {"bar": 2})
    add_document("foo", {"bar": 3, "foo": "bazbaz"})
    add_document("foo", {"bar": 5, "foo": "foobar"})
    add_document("foo", {"bar": 5})
    add_document("foo", {"bar": 5, "foo": "foobaz"})
    add_document("foo", {"bar": 9})

    # And I do a filtered
    f = Filter().filter(Range("bar", gte=5))
    t.aggregate(aggregation=Aggregations("missing_foo", "foo", "missing",
                                         filter_val=f, filter_name="high_bar"))
    t[0:10]

    # Then I get the expected results
    t.aggregations().should.have.key("high_bar")
    t.aggregations()["high_bar"].should.have.key("missing_foo").being.equal(
        {u'doc_count': 2})
    t.aggregations()["high_bar"]["doc_count"].should.equal(4)