From 7dba389eb5ed7916f1e1c51e54c90bb4836ed7ee Mon Sep 17 00:00:00 2001 From: Pablo Martin Date: Thu, 21 Jul 2022 12:24:35 +0200 Subject: [PATCH] Docstrings, typing, a bit of renaming and minor stuff. --- cli.py | 4 +- connections.py | 93 ++++++++++++++++++++++++++++++++++++++ query_performance_gauge.py | 91 ++++++++++++------------------------- 3 files changed, 125 insertions(+), 63 deletions(-) create mode 100644 connections.py diff --git a/cli.py b/cli.py index fbdcc6e..c21acfa 100644 --- a/cli.py +++ b/cli.py @@ -2,10 +2,10 @@ import json import click -from query_performance_gauge import main +from query_performance_gauge import run_measuring_session @click.command() @click.option("--config", required=True, type=click.File()) def measure_performance(config): - main(json.load(config)) + run_measuring_session(json.load(config)) diff --git a/connections.py b/connections.py new file mode 100644 index 0000000..8d7e34c --- /dev/null +++ b/connections.py @@ -0,0 +1,93 @@ +from typing import Union, Callable + +import mysql.connector +import trino.dbapi +from trino.auth import BasicAuthentication +from trino.dbapi import connect + + +def get_connection(connection_config: dict) -> Union[trino.dbapi.Connection]: + """ + Pick the right way to build a connection and pass it the connection details. + + :param connection_config: confi + :return: + """ + connection_builder = pick_connection_builder(connection_config["engine"]) + connection = connection_builder(connection_config) + return connection + + +def get_possible_connection_builders() -> dict: + """ + Get the complete list of connection builders. + + :return: a dict where the keys are the strings that identify each + connection engine and the values are the callable function that will return + a valid connection object. + """ + return { + "trino": get_connection_to_trino, + "mysql": get_connection_to_mysql, + } + + +def pick_connection_builder(connection_engine_name: str) -> Callable: + """ + Get a connection builder from a string name. + + :param connection_engine_name: the string defining the connection engine. + :return: the connection builder function. + """ + possible_connection_builders = get_possible_connection_builders() + + try: + connection_builder = possible_connection_builders[connection_engine_name] + except KeyError: + raise ValueError( + f"Connection type {connection_engine_name} is unknown. Please review config." + ) + + return connection_builder + + +def get_connection_to_trino(connection_config: dict) -> trino.dbapi.Connection: + """ + Build a connection to Trino from the passed config. + + :param connection_config: specifies host, port, etc. + :return: the connection object + """ + + return connect( + host=connection_config["host"], + port=connection_config["port"], + user=connection_config["user"], + auth=BasicAuthentication( + connection_config["user"], + connection_config["password"], + ), + http_scheme=connection_config["http_scheme"], + catalog=connection_config["catalog"], + schema=connection_config["schema"], + ) + + +def get_connection_to_mysql( + connection_config, +) -> mysql.connector.connection.MySQLConnection: + """ + Build a connection to MySQL from the passed config. + + :param connection_config: specifies host, port, etc. + :return: the connection object + """ + connection = mysql.connector.connect( + host=connection_config["host"], + port=connection_config["port"], + user=connection_config["user"], + password=connection_config["password"], + database=connection_config["schema"], + ) + + return connection diff --git a/query_performance_gauge.py b/query_performance_gauge.py index f762dad..288ed43 100644 --- a/query_performance_gauge.py +++ b/query_performance_gauge.py @@ -1,17 +1,23 @@ import time import traceback -from typing import Union, Callable +from typing import Union +import mysql.connector.connection import trino.dbapi -from trino.dbapi import connect -from trino.auth import BasicAuthentication -import mysql.connector + +from connections import get_connection -def main(config: dict) -> None: +def run_measuring_session(config: dict) -> None: + """ + Complete session flow. Connect, test all queries. + + :param config: the full config for the measuring session. + :return: None + """ print("Starting the measuring session.") - connection = get_connection(config) + connection = get_connection(config["connection_details"]) for query_config in config["queries_to_measure"]: try: @@ -27,67 +33,30 @@ def main(config: dict) -> None: class TestableQuery: + """ + Simple object to hold the details of a query that will be measured. + """ + def __init__(self, name: str, query_string: str): self.name = name self.query_string = query_string -def measure_query_runtime(connection: trino.dbapi.Connection, query: TestableQuery): +def measure_query_runtime( + connection: Union[trino.dbapi.Connection, mysql.connector.MySQLConnection], + query_to_measure: TestableQuery, +) -> None: + """ + Execute a query against the given connection and print the time it took. + + :param connection: a connection object capable of generating cursors. + :param query_to_measure: the query that will be measured. + :return: None + """ start_time = time.time() cur = connection.cursor() - cur.execute(query.query_string) + cur.execute(query_to_measure.query_string) rows = cur.fetchall() - print(f"Query '{query.name}' took {int(time.time() - start_time)} seconds to run.") - - -def get_connection(config: dict) -> Union[trino.dbapi.Connection]: - connection_builder = pick_connection_builder(config["connection_details"]["engine"]) - connection = connection_builder(config) - return connection - - -def get_possible_connection_builders() -> dict: - return { - "trino": get_connection_to_trino, - "mysql": get_connection_to_mysql, - } - - -def pick_connection_builder(connection_engine_name: str) -> Callable: - possible_connection_builders = get_possible_connection_builders() - - try: - connection_builder = possible_connection_builders[connection_engine_name] - except KeyError: - raise ValueError( - f"Connection type {connection_engine_name} is unknown. Please review config." - ) - - return connection_builder - - -def get_connection_to_trino(config): - return connect( - host=config["connection_details"]["host"], - port=config["connection_details"]["port"], - user=config["connection_details"]["user"], - auth=BasicAuthentication( - config["connection_details"]["user"], - config["connection_details"]["password"], - ), - http_scheme=config["connection_details"]["http_scheme"], - catalog=config["connection_details"]["catalog"], - schema=config["connection_details"]["schema"], + print( + f"Query '{query_to_measure.name}' took {int(time.time() - start_time)} seconds to run." ) - - -def get_connection_to_mysql(config) -> mysql.connector.connection.MySQLConnection: - connection = mysql.connector.connect( - host=config["connection_details"]["host"], - port=config["connection_details"]["port"], - user=config["connection_details"]["user"], - password=config["connection_details"]["password"], - database=config["connection_details"]["schema"], - ) - - return connection