1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
# Copyright (C) 2022-present MongoDB, Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the Server Side Public License, version 1,
# as published by MongoDB, Inc.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# Server Side Public License for more details.
#
# You should have received a copy of the Server Side Public License
# along with this program. If not, see
# <http://www.mongodb.com/licensing/server-side-public-license>.
#
# As a special exception, the copyright holders give permission to link the
# code of portions of this program with the OpenSSL library under certain
# conditions as described in each individual source file and distribute
# linked combinations including the program with the OpenSSL library. You
# must comply with the Server Side Public License in all respects for
# all of the code used other than as permitted herein. If you modify file(s)
# with this exception, you may extend this exception to your version of the
# file(s), but you are not obligated to do so. If you do not wish to do so,
# delete this exception statement from your version. If you delete this
# exception statement from all source files in the program, then also delete
# it in the license file.
#
"""Workload Execution. Run the queries to collect data for calibration."""
from __future__ import annotations
from dataclasses import asdict, dataclass
from typing import Sequence
import bson.json_util as json
from bson.objectid import ObjectId
from data_generator import CollectionInfo
from database_instance import DatabaseInstance, Pipeline
from config import WorkloadExecutionConfig, WriteMode
__all__ = ['execute']
@dataclass
class Query:
"""Query pipleline and related model input parameters."""
pipeline: Pipeline
keys_length_in_bytes: int = 0
number_of_fields: int = 0
note: any = None
@dataclass
class QueryParameters:
"""Model input parameters specific for a workload query executed on some collection and used for calibration."""
keys_length_in_bytes: int
average_document_size_in_bytes: float
note: any
number_of_fields: int = 0
def to_json(self) -> str:
"""Serialize the parameters to JSON."""
return json.dumps(asdict(self))
@staticmethod
def from_json(json_str: str) -> QueryParameters:
"""Deserialize from JSON."""
return QueryParameters(**json.loads(json_str))
async def execute(database: DatabaseInstance, config: WorkloadExecutionConfig,
collection_infos: Sequence[CollectionInfo], queries: Sequence[Query]):
"""Run the given queries and write the collected explain into collection."""
if not config.enabled:
return
collector = WorkloadExecution(database, config)
await collector.async_init()
print('>>> running queries')
await collector.collect(collection_infos, queries)
class WorkloadExecution:
"""Runs a number of queries to generate and collect execution statistics."""
def __init__(self, database: DatabaseInstance, config: WorkloadExecutionConfig):
self.database = database
self.config = config
async def async_init(self):
"""Initialize the database settings."""
await self.database.enable_sbe(True)
await self.database.enable_cascades(True)
if self.config.write_mode == WriteMode.REPLACE:
await self.database.drop_collection(self.config.output_collection_name)
async def collect(self, collection_infos: Sequence[CollectionInfo], queries: Sequence[Query]):
"""Run the given piplelines on the given collection to generate and collect execution statistics."""
measurements = []
for coll_info in collection_infos:
print(f'\n>>>>> running queries on collection {coll_info.name}')
for query in queries:
print(f'>>>>>>> running query {query.pipeline}')
await self._run_query(coll_info, query, measurements)
await self.database.insert_many(self.config.output_collection_name, measurements)
async def _run_query(self, coll_info: CollectionInfo, query: Query, result: Sequence):
# warm up
for _ in range(self.config.warmup_runs):
await self.database.explain(coll_info.name, query.pipeline)
run_id = ObjectId()
avg_doc_size = await self.database.get_average_document_size(coll_info.name)
parameters = QueryParameters(keys_length_in_bytes=query.keys_length_in_bytes,
number_of_fields=query.number_of_fields,
average_document_size_in_bytes=avg_doc_size, note=query.note)
for _ in range(self.config.runs):
explain = await self.database.explain(coll_info.name, query.pipeline)
if explain['ok'] == 1:
result.append({
'run_id': run_id, 'collection': coll_info.name,
'pipeline': json.dumps(query.pipeline), 'explain': json.dumps(explain),
'query_parameters': parameters.to_json()
})
|
