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, clean_up_connection
from _version import __version__


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(f"Query Performance Gauge - Version: {__version__}")
    print("Starting the measuring session.")

    connection = get_connection(config["connection_details"])

    measure_queries(config["queries_to_measure"], connection)

    print("Finished the measuring session.")

    clean_up_connection(config["connection_details"])


def measure_queries(
    queries_config: dict,
    connection: Union[trino.dbapi.Connection, mysql.connector.MySQLConnection],
) -> None:
    """
    Measure several queries through a connection.

    :param queries_config: the configuration for the queries to measure.
    :param connection: the connection to the queriable server.
    :return: None
    """
    for query_config in queries_config:
        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()}")


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()
    obtained_rows = len(rows)
    print(
        f"Query '{query_to_measure.name}' took {int(time.time() - start_time)} seconds to run and returned {obtained_rows} rows."
    )