composite_agg example and tests

Author: honzakral

examples/composite_agg.py

@@ -0,0 +1,34 @@
+from __future__ import print_function
+
+from elasticsearch_dsl import connections, Search, A
+
+def scan_aggs(search, source_aggs, inner_aggs={}, size=10):
+    """
+    Helper function used to iterate over all possible bucket combinations of
+    ``source_aggs``, returning results of ``inner_aggs`` for each. Uses the
+    ``composite`` aggregation under the hood to perform this.
+    """
+    def run_search(**kwargs):
+        s = search[:0]
+        s.aggs.bucket('comp', 'composite', sources=source_aggs, size=size, **kwargs)
+        for agg_name, agg in inner_aggs.items():
+            s.aggs['comp'][agg_name] = agg
+        return s.execute()
+
+    response = run_search()
+    while response.aggregations.comp.buckets:
+        for b in response.aggregations.comp.buckets:
+            yield b
+        response = run_search(after=response.aggregations.comp.after_key)
+
+
+if __name__ == '__main__':
+    connections.create_connection()
+    s = Search(index='git')
+
+    for b in scan_aggs(s,
+                       {'files': A('terms', field='files')},
+                       {'first_seen': A('min', field='committed_date')}):
+        print('File %s has been modified %d times, first seen at %s.' % (
+            b.key.files, b.doc_count, b.first_seen.value_as_string
+        ))

test_elasticsearch_dsl/test_integration/test_examples/composite_agg.py

@@ -0,0 +1 @@
+../../../examples/composite_agg.py

test_elasticsearch_dsl/test_integration/test_examples/test_composite_aggs.py

@@ -0,0 +1,10 @@
+from elasticsearch_dsl import Search, A
+
+from .composite_agg import scan_aggs
+
+def test_scan_aggs_exhausts_all_files(data_client):
+    s = Search(index='flat-git')
+    key_aggs = {'files': A('terms', field='files')}
+    file_list = list(scan_aggs(s, key_aggs))
+
+    assert len(file_list) == 26