Completed
Pull Request — master (#58)
by
unknown
01:10
created

ElasticsearchStorage.save()   A

Complexity

Conditions 1

Size

Total Lines 6

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 0
Metric Value
cc 1
c 2
b 0
f 0
dl 0
loc 6
rs 9.4285
1
import datetime
2
3
4
class ElasticsearchStorage(object):
5
    def __init__(self, elasticsearch_hosts, elasticsearch_index, elasticsearch_doctype, logger,
6
                 default_machine_id=None):
7
        try:
8
            import elasticsearch
9
        except ImportError as exc:
10
            raise ImportError(exc.args, "Please install elasticsearch or pytest-benchmark[elasticsearch]")
11
        self._elasticsearch_hosts = elasticsearch_hosts
12
        self._elasticsearch_index = elasticsearch_index
13
        self._elasticsearch_doctype = elasticsearch_doctype
14
        self._elasticsearch = elasticsearch.Elasticsearch(self._elasticsearch_hosts)
15
        self.default_machine_id = default_machine_id
16
        self.logger = logger
17
        self._cache = {}
18
        self._create_index()
19
20
    def __str__(self):
21
        return str(self._elasticsearch_hosts)
22
23
    @property
24
    def location(self):
25
        return str(self._elasticsearch_hosts)
26
27
    def query(self, project):
28
        """
29
        Returns sorted records names (ids) that corresponds with project.
30
        """
31
        return [commit_and_time for commit_and_time, _ in self.load(project)]
32
33
    def load(self, project, id_prefix=None):
34
        """
35
        Yield key and content of records that corresponds with project name.
36
        """
37
        r = self._search(project, id_prefix)
38
        groupped_data = self._group_by_commit_and_time(r["hits"]["hits"])
39
        result = [(key, value) for key, value in groupped_data.items()]
40
        result.sort(key=lambda x: datetime.datetime.strptime(x[1]["datetime"], "%Y-%m-%dT%H:%M:%S.%f"))
41
        for key, data in result:
42
            yield key, data
43
44
    def _search(self, project, id_prefix=None):
45
        body = {
46
            "size": 1000,
47
            "sort": [
48
                {
49
                    "datetime": {
50
                        "order": "desc"
51
                    }
52
                }
53
            ],
54
            "query": {
55
                "bool": {
56
                    "filter": {
57
                        "term": {
58
                            "commit_info.project": project
59
                        }
60
                    }
61
                }
62
            }
63
        }
64
        if id_prefix:
65
            body["query"]["bool"]["must"] = {
66
                "prefix": {
67
                    "_id": id_prefix
68
                }
69
            }
70
71
        return self._elasticsearch.search(index=self._elasticsearch_index,
72
                                          doc_type=self._elasticsearch_doctype,
73
                                          body=body)
74
75
    @staticmethod
76
    def _benchmark_from_es_record(source_es_record):
77
        result = {}
78
        for benchmark_key in ("group", "stats", "options", "param", "name", "params", "fullname"):
79
            result[benchmark_key] = source_es_record[benchmark_key]
80
        return result
81
82
    @staticmethod
83
    def _run_info_from_es_record(source_es_record):
84
        result = {}
85
        for run_key in ("machine_info", "commit_info", "datetime", "version"):
86
            result[run_key] = source_es_record[run_key]
87
        return result
88
89
    def _group_by_commit_and_time(self, hits):
90
        result = {}
91
        for hit in hits:
92
            source_hit = hit["_source"]
93
            key = "%s_%s" % (source_hit["commit_info"]["id"], source_hit["datetime"])
94
            benchmark = self._benchmark_from_es_record(source_hit)
95
            if key in result:
96
                result[key]["benchmarks"].append(benchmark)
97
            else:
98
                run_info = self._run_info_from_es_record(source_hit)
99
                run_info["benchmarks"] = [benchmark]
100
                result[key] = run_info
101
        return result
102
103
    def load_benchmarks(self, project):
104
        """
105
        Yield benchmarks that corresponds with project. Put path and
106
        source (uncommon part of path) to benchmark dict.
107
        """
108
        r = self._search(project)
109
        for hit in r["hits"]["hits"]:
110
            yield self._benchmark_from_es_record(hit["_source"])
111
112
    def save(self, document, document_id):
113
        self._elasticsearch.index(
114
            index=self._elasticsearch_index,
115
            doc_type=self._elasticsearch_doctype,
116
            body=document,
117
            id=document_id,
118
        )
119
120
    def _create_index(self):
121
        mapping = {
122
            "mappings": {
123
                "benchmark": {
124
                    "properties": {
125
                        "commit_info": {
126
                            "properties": {
127
                                "dirty": {
128
                                    "type": "boolean"
129
                                },
130
                                "id": {
131
                                    "type": "string",
132
                                    "index": "not_analyzed"
133
134
                                },
135
                                "project": {
136
                                    "type": "string",
137
                                    "index": "not_analyzed"
138
                                }
139
                            }
140
                        },
141
                        "datetime": {
142
                            "type": "date",
143
                            "format": "strict_date_optional_time||epoch_millis"
144
                        },
145
                        "name": {
146
                            "type": "string",
147
                            "index": "not_analyzed"
148
                        },
149
                        "fullname": {
150
                            "type": "string",
151
                            "index": "not_analyzed"
152
                        },
153
                        "version": {
154
                            "type": "string",
155
                            "index": "not_analyzed"
156
                        },
157
                        "machine_info": {
158
                            "properties": {
159
                                "machine": {
160
                                    "type": "string",
161
                                    "index": "not_analyzed"
162
                                },
163
                                "node": {
164
                                    "type": "string",
165
                                    "index": "not_analyzed"
166
                                },
167
                                "processor": {
168
                                    "type": "string",
169
                                    "index": "not_analyzed"
170
                                },
171
                                "python_build": {
172
                                    "type": "string",
173
                                    "index": "not_analyzed"
174
                                },
175
                                "python_compiler": {
176
                                    "type": "string",
177
                                    "index": "not_analyzed"
178
                                },
179
                                "python_implementation": {
180
                                    "type": "string",
181
                                    "index": "not_analyzed"
182
                                },
183
                                "python_implementation_version": {
184
                                    "type": "string",
185
                                    "index": "not_analyzed"
186
                                },
187
                                "python_version": {
188
                                    "type": "string",
189
                                    "index": "not_analyzed"
190
                                },
191
                                "release": {
192
                                    "type": "string",
193
                                    "index": "not_analyzed"
194
                                },
195
                                "system": {
196
                                    "type": "string",
197
                                    "index": "not_analyzed"
198
                                }
199
                            }
200
                        },
201
                        "options": {
202
                            "properties": {
203
                                "disable_gc": {
204
                                    "type": "boolean"
205
                                },
206
                                "max_time": {
207
                                    "type": "double"
208
                                },
209
                                "min_rounds": {
210
                                    "type": "long"
211
                                },
212
                                "min_time": {
213
                                    "type": "double"
214
                                },
215
                                "timer": {
216
                                    "type": "string"
217
                                },
218
                                "warmup": {
219
                                    "type": "boolean"
220
                                }
221
                            }
222
                        },
223
                        "stats": {
224
                            "properties": {
225
                                "hd15iqr": {
226
                                    "type": "double"
227
                                },
228
                                "iqr": {
229
                                    "type": "double"
230
                                },
231
                                "iqr_outliers": {
232
                                    "type": "long"
233
                                },
234
                                "iterations": {
235
                                    "type": "long"
236
                                },
237
                                "ld15iqr": {
238
                                    "type": "double"
239
                                },
240
                                "max": {
241
                                    "type": "double"
242
                                },
243
                                "mean": {
244
                                    "type": "double"
245
                                },
246
                                "median": {
247
                                    "type": "double"
248
                                },
249
                                "min": {
250
                                    "type": "double"
251
                                },
252
                                "outliers": {
253
                                    "type": "string"
254
                                },
255
                                "q1": {
256
                                    "type": "double"
257
                                },
258
                                "q3": {
259
                                    "type": "double"
260
                                },
261
                                "rounds": {
262
                                    "type": "long"
263
                                },
264
                                "stddev": {
265
                                    "type": "double"
266
                                },
267
                                "stddev_outliers": {
268
                                    "type": "long"
269
                                }
270
                            }
271
                        },
272
                    }
273
                }
274
            }
275
        }
276
        self._elasticsearch.indices.create(index=self._elasticsearch_index, ignore=400, body=mapping)
277
278