query-performance-gauge/query_performance_gauge.py

import time
import traceback
from typing import Union

import mysql.connector.connection
import trino.dbapi

from connections import get_connection


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_details"])

    for query_config in config["queries_to_measure"]:
        try:
            query = TestableQuery(
                name=query_config["name"], query_string=query_config["query_string"]
            )
            measure_query_runtime(connection, query)
        except Exception as e:
            print(f"""Something went wrong with query {query_config["name"]}.""")
            print(f"{traceback.format_exc()}")

    print("Finished the measuring session.")


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: 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_to_measure.query_string)
    rows = cur.fetchall()
    print(
        f"Query '{query_to_measure.name}' took {int(time.time() - start_time)} seconds to run."
    )