aboutsummaryrefslogtreecommitdiff
path: root/qc_app/db_utils.py
blob: 75b6b73e69259f5f949615a7311965b26472dbc9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
"""module contains all db related stuff"""
import logging
import traceback
import contextlib
from typing import Any, Tuple, Optional, Iterator, Callable

from urllib.parse import urlparse
import MySQLdb as mdb
from flask import current_app as app

def parse_db_url(db_url) -> Tuple:
    """
    Parse SQL_URI configuration variable.
    """
    parsed_db = urlparse(db_url)
    return (parsed_db.hostname, parsed_db.username,
            parsed_db.password, parsed_db.path[1:], parsed_db.port)


@contextlib.contextmanager
def database_connection(db_url: Optional[str] = None) -> Iterator[mdb.Connection]:
    """function to create db connector"""
    host, user, passwd, db_name, db_port = parse_db_url(
        db_url or app.config["SQL_URI"])
    connection = mdb.connect(
        host, user, passwd, db_name, port=(db_port or 3306))
    try:
        yield connection
        connection.commit()
    except mdb.Error as _mdb_err:
        logging.error(traceback.format_exc())
        connection.rollback()
    finally:
        connection.close()

def with_db_connection(func: Callable[[mdb.Connection], Any]) -> Any:
    """Call `func` with a MySQDdb database connection."""
    with database_connection(app.config["SQL_URI"]) as conn:
        return func(conn)