tmpl-flask/flaskapp/db.py

"""Works with MySQL database objects"""

from typing import List, Optional

from flask import current_app

from pymysql import Connection
from pymysql.cursors import Cursor

from . import exts

TABLE = '${REPO_NAME_SNAKE}'
db_obj: List[Optional[Connection]] = [None]


def init_db() -> None:
    """Initializes MySQL database
    from schema.sql file"""

    cur = get_cursor()

    with current_app.open_resource('../db/schema.sql') as f:

        schema: bytes = f.read()  # type: ignore
        schema_str = schema.decode('utf-8')

        for query in schema_str.split(';'):

            query = query.strip()
            if query == '' or query.startswith('--'):
                continue

            cur.execute(query)

    cur.close()


def create_db(name: str) -> None:
    """Create the database if not exists

    Args:
        name (str): Database name
    """

    sql = get_db()
    cur = sql.cursor(Cursor)
    cur.execute(
        f'create database if not exists {name}'
    )
    cur.close()

    sql.select_db(name)


def get_db() -> Connection:
    """Get MySQL connection object

    Returns:
        MySQL database
    """

    if db_obj[0] is None:
        db_obj[0] = exts.sql.connect()
    return db_obj[0]  # type: ignore


def get_cursor() -> Cursor:
    """Get MySQL database cursor object
    for executing commands.
    Equivalent to:
    ```
    conn = db.get_db()
    cur = conn.cursor()
    ```
    """

    return get_db().cursor()