Docstrings, typing, a bit of renaming and minor stuff.

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))

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

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."
+    print(
+        f"Query '{query_to_measure.name}' took {int(time.time() - start_time)} seconds to run."
     )
-
-    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"],
-    )
-
-
-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