93 lines
2.7 KiB
Python
93 lines
2.7 KiB
Python
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
|