From 3c8d2281ca650b1a72f130b550da8244cd270e98 Mon Sep 17 00:00:00 2001
From: Pablo Martin <pablo.martin@lolamarket.com>
Date: Mon, 22 Aug 2022 13:45:46 +0200
Subject: [PATCH] Typing, docstrings, refactors.

---
 connections.py             |  2 +-
 query_performance_gauge.py | 24 +++++++++++++++++++-----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/connections.py b/connections.py
index a74df0d..e99f333 100644
--- a/connections.py
+++ b/connections.py
@@ -46,7 +46,7 @@ class MySSHTunnel:
         self.tunnel.stop()
 
 
-def get_connection(connection_config: dict) -> Union[trino.dbapi.Connection]:
+def get_connection(connection_config: dict) -> Union[trino.dbapi.Connection, mysql.connector.MySQLConnection]:
     """
     Pick the right way to build a connection and pass it the connection details.
 
diff --git a/query_performance_gauge.py b/query_performance_gauge.py
index 5d00fdf..353fc45 100644
--- a/query_performance_gauge.py
+++ b/query_performance_gauge.py
@@ -21,7 +21,25 @@ def run_measuring_session(config: dict) -> None:
 
     connection = get_connection(config["connection_details"])
 
-    for query_config in config["queries_to_measure"]:
+    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"]
@@ -31,10 +49,6 @@ def run_measuring_session(config: dict) -> None:
             print(f"""Something went wrong with query {query_config["name"]}.""")
             print(f"{traceback.format_exc()}")
 
-    print("Finished the measuring session.")
-
-    clean_up_connection(config["connection_details"])
-
 
 class TestableQuery:
     """