def _call_add_database(self, body: Body) -> Response:
"""Add database and initialize driver for it."""
user = body["user"]
password = body["password"]
host = body["host"]
port = body["port"]
dbname = body["dbname"]
number_workers = body["number_workers"]
if not Driver.validate_connection(user, password, host, port,
dbname): # TODO move to Database
return get_response(400)
if body["id"] in self._databases:
return get_response(400)
db_instance = Database(
body["id"],
user,
password,
host,
port,
dbname,
number_workers,
"tcp://{:s}:{:s}".format(
self._workload_sub_host,
self._workload_pubsub_port,
),
self._default_tables,
self._storage_host,
self._storage_password,
self._storage_port,
self._storage_user,
)
self._databases[body["id"]] = db_instance
return get_response(200)
def _call_delete_data(self, body: Body) -> Response:
folder_name: str = body["folder_name"]
if self._check_if_database_blocked():
return get_error_response(400, "Already loading data")
for database in list(self._databases.values()):
if not database.delete_data(folder_name):
return get_response(400)
return get_response(200)
def _call_delete_database(self, body: Body) -> Response:
id: str = body["id"]
database: Optional[Database] = self._databases.pop(id, None)
if database:
database.close()
del database
return get_response(200)
else:
return get_response(404)
def _call_load_data(self, body: Body) -> Response:
folder_name: str = body["folder_name"]
if self._check_if_database_blocked():
return get_error_response(400, "Already loading data")
for database in list(self._databases.values()):
if not database.load_data(folder_name):
return get_response(
400) # TODO return which DB couldn't import
return get_response(200)
def _call_deactivate_plugin(self, body: Body) -> Response:
id: str = body["id"]
plugin: str = body["plugin"]
if id not in self._databases.keys():
response = get_response(400)
elif self._databases[id].deactivate_plugin(plugin):
response = get_response(200)
else:
response = get_response(423)
return response
def _handle_request(self, request: Request) -> Response:
try:
validate(instance=request, schema=request_schema)
func, schema = self._calls[request["header"]["message"]]
if schema:
validate(request["body"], schema=schema)
return func(request["body"])
except ValidationError:
return get_response(400)
except KeyError:
return get_response(404)
def _call_plugin_setting(self, body: Body) -> Response:
id: str = body["id"]
name: str = body["name"]
value: str = body["value"]
if id not in self._databases.keys():
response = get_response(400)
elif self._databases[id].set_plugin_setting(name, value):
response = get_response(200)
else:
response = get_response(423)
return response
def get(self) -> Union[int, Response]:
"""Return latency information from the stored queries."""
currentts = time_ns()
startts = currentts - 2_000_000_000
endts = currentts - 1_000_000_000
latency: Dict[str, float] = {}
try:
active_databases = _active_databases()
except ValidationError:
return 500
for database in active_databases:
result = storage_connection.query(
'SELECT MEAN("latency") as "latency" FROM successful_queries WHERE time > $startts AND time <= $endts;',
database=database,
bind_params={"startts": startts, "endts": endts},
)
latency_value = list(result["successful_queries", None])
if len(latency_value) > 0:
latency[database] = list(result["successful_queries", None])[0][
"latency"
]
else:
latency[database] = 0
response = get_response(200)
response["body"]["latency"] = latency
return response
def get(self) -> Union[int, Response]:
"""Return throughput information from the stored queries."""
currentts = time_ns()
startts = currentts - 2_000_000_000
endts = currentts - 1_000_000_000
throughput: Dict[str, int] = {}
try:
active_databases = _active_databases()
except ValidationError:
return 500
for database in active_databases:
result = storage_connection.query(
'SELECT COUNT("latency") FROM successful_queries WHERE time > $startts AND time <= $endts;',
database=database,
bind_params={"startts": startts, "endts": endts},
)
throughput_value = list(result["successful_queries", None])
if len(throughput_value) > 0:
throughput[database] = list(result["successful_queries", None])[0][
"count"
]
else:
throughput[database] = 0
response = get_response(200)
response["body"]["throughput"] = throughput
return response
def _call_queue_length(self, body: Body) -> Response:
queue_length = {}
for id, database in self._databases.items():
queue_length[id] = database.get_queue_length()
response = get_response(200)
response["body"]["queue_length"] = queue_length
return response
def _generate_workload(self) -> None:
if self._generate_workload_flag:
workload = self._get_workload(self._workload_type)
queries = workload.generate_workload(self._frequency)
response = get_response(200)
response["body"] = {"querylist": queries}
self._publish_data(response)
def _call_get_plugin_setting(self, body: Body) -> Response:
plugin_settings = [{
"id": id,
"plugin_settings": database.get_plugin_setting()
} for id, database in self._databases.items()]
response = get_response(200)
response["body"]["plugin_settings"] = plugin_settings
return response
def _call_get_plugins(self, body: Body) -> Response:
activated_plugins = [{
"id": id,
"plugins": database.get_plugins()
} for id, database in self._databases.items()]
response = get_response(200)
response["body"]["plugins"] = activated_plugins
return response
def test_returns_404_on_call_not_found(self, isolated_server):
"""Returns 404 when call not found."""
request = {
"header": {
"message": "How do you drown a hipster? Throw him in the mainstream."
},
"body": {},
}
assert get_response(404) == isolated_server._handle_request(request)
def test_returns_400_on_valid_request_with_wrong_schema(self, isolated_server):
"""Returns 400 for valid request with wrong schema."""
specific_request_schema = {
"type": "object",
"required": ["data"],
"properties": {"data": {"type": "string"}},
}
isolated_server._calls = {"call": (call_function, specific_request_schema)}
request = {"header": {"message": "call"}, "body": {"wrong_key": "info"}}
assert get_response(400) == isolated_server._handle_request(request)
def test_starts_a_workload(self, isolated_generator: WorkloadGenerator):
"""Test starting of the workload generation."""
body = {"folder_name": "benchmark_name", "frequency": 100}
response = isolated_generator._call_start_workload(body)
assert isolated_generator._frequency == 100
assert isolated_generator._workload_type == "benchmark_name"
assert isolated_generator._generate_workload_flag
assert list(isolated_generator._workloads.keys()) == ["benchmark_name"]
assert response == get_response(200)
def _call_status(self, body: Body) -> Response:
status = [{
"id": database_id,
"database_blocked_status": database.get_database_blocked(),
"worker_pool_status": database.get_worker_pool_status(),
"loaded_benchmarks": database.get_loaded_benchmarks(),
"loaded_tables": database.get_loaded_tables(),
} for database_id, database in self._databases.items()]
response = get_response(200)
response["body"]["status"] = status
return response
def _call_get_databases(self, body: Body) -> Response:
"""Get list of all databases."""
databases = [{
"id": id,
"host": database.driver.host,
"port": database.driver.port,
"number_workers": database.number_workers,
"dbname": database.driver.dbname,
} for id, database in self._databases.items()]
response = get_response(200)
response["body"]["databases"] = databases
return response
def test_generates_a_workload(self, mocked_publish_data,
isolated_generator: WorkloadGenerator):
"""Test workload generation."""
isolated_generator._generate_workload_flag = True
isolated_generator._workload_type = "benchmark_name"
isolated_generator._frequency = 1
mocked_publish_data.return_value = None
isolated_generator._generate_workload()
expected_response = get_response(200)
expected_response["body"] = {"querylist": fake_workload}
mocked_publish_data.assert_any_call(expected_response)
def _call_start_workload(self, body: Body) -> Response:
frequency: int = body["frequency"]
workload_type: str = body["folder_name"]
try:
self._get_workload(workload_type)
except (
NotExistingWorkloadFolderException,
EmptyWorkloadFolderException,
QueryTypeNotFoundException,
QueryTypesNotSpecifiedException,
) as e:
return get_error_response(400, str(e))
self._workload_type = workload_type
self._frequency = frequency
self._generate_workload_flag = True
return get_response(200)
def get(self) -> Union[int, Response]:
"""Return chunks data information for every database."""
chunks: Dict[str, Dict] = {}
try:
active_databases = _active_databases()
except ValidationError:
return 500
for database in active_databases:
result = storage_connection.query(
'SELECT LAST("chunks_data_meta_information") FROM chunks_data',
database=database,
)
chunks_value = list(result["chunks_data", None])
if len(chunks_value) > 0:
chunks[database] = loads(chunks_value[0]["last"])
else:
chunks[database] = {}
response = get_response(200)
response["body"]["chunks_data"] = chunks
return response
def get(self) -> Union[int, Response]:
"""Return storage metadata from database manager."""
storage: Dict[str, Dict] = {}
try:
active_databases = _active_databases()
except ValidationError:
return 500
for database in active_databases:
result = storage_connection.query(
'SELECT LAST("storage_meta_information") FROM storage',
database=database,
)
storage_value = list(result["storage", None])
if len(storage_value) > 0:
storage[database] = loads(storage_value[0]["last"])
else:
storage[database] = {}
response = get_response(200)
response["body"]["storage"] = storage
return response
def get(self) -> Union[int, Response]:
"""Return cpu and memory information for every database and the number of thread it is using from database manager."""
system: Dict[str, Dict] = {}
try:
active_databases = _active_databases()
except ValidationError:
return 500
for database in active_databases:
result = storage_connection.query(
'SELECT LAST("cpu"), * FROM system_data', database=database,
)
system_value = list(result["system_data", None])
if len(system_value) > 0:
system[database] = {
"cpu": loads(system_value[0]["cpu"]),
"memory": loads(system_value[0]["memory"]),
"database_threads": loads(system_value[0]["database_threads"]),
}
else:
system[database] = {}
response = get_response(200)
response["body"]["system_data"] = system
return response
def post(self) -> Response:
"""Start the workload generator."""
message = Request(header=Header(message="start worker"), body={})
response = _send_message(db_manager_socket, message)
if response["header"]["status"] != 200:
return get_error_response(
400, response["body"].get("error", "Error during starting of worker")
)
message = Request(
header=Header(message="start workload"),
body={
"folder_name": control.payload["folder_name"],
"frequency": control.payload.get("frequency", 200),
},
)
response = _send_message(generator_socket, message)
if response["header"]["status"] != 200:
return get_error_response(
400,
response["body"].get("error", "Error during starting of the workload"),
)
return get_response(200)
def call_function(body: Body):
"""Execute request."""
return get_response(200)
def test_returns_200_on_valid_request(self, isolated_server):
"""Returns 200 for valid request."""
isolated_server._calls = {"call": (call_function, None)}
request = {"header": {"message": "call"}, "body": {}}
assert get_response(200) == isolated_server._handle_request(request)
def test_returns_400_on_bad_request(self, isolated_server):
"""Returns 400 on bad request."""
request = {
"What's the best thing about Switzerland?": "I don't know, but the flag is a big plus."
}
assert get_response(400) == isolated_server._handle_request(request)
def _call_close_worker(self, body: Body) -> Response:
for database in self._databases.values():
if not database.close_worker():
return get_response(400)
return get_response(200)
def test_stops_a_workload(self, isolated_generator: WorkloadGenerator):
"""Test stopping of the workload generation."""
response = isolated_generator._call_stop_workload({})
assert not isolated_generator._generate_workload_flag
assert response == get_response(200)
def _call_stop_workload(self, body: Body) -> Response:
self._generate_workload_flag = False
return get_response(200)
