pypistats.dev/pypistats/tasks/db.py

"""Database tasks."""
import psycopg2
from sqlalchemy import create_engine
from sqlalchemy.exc import ProgrammingError

# from pypistats.extensions import db
from pypistats.secret import postgresql


DBNAME = "pypistats"


def create_databases():
    """Create the databases for each environment."""
    env = "prod"
    url = \
        "postgresql://{username}:{password}@{host}:{port}/{dbname}".format(
            username=postgresql[env]["username"],
            password=postgresql[env]["password"],
            host=postgresql[env]["host"],
            port=postgresql[env]["port"],
            dbname=DBNAME,
        )
    engine = create_engine(url)
    connection = engine.connect()

    for env, config in postgresql.items():
        query = f"""CREATE DATABASE {config["dbname"]}"""
        try:
            connection.execute("commit")
            connection.execute(query)
            connection.execute("commit")
            print(f"Created db: {config['dbname']}.")
        except ProgrammingError:
            print(f"Database {config['dbname']} already exists.")


def get_db_connection(env="dev"):
    """Get a db connection cursor."""
    connection = psycopg2.connect(
        dbname=postgresql[env]['dbname'],
        user=postgresql[env]['username'],
        password=postgresql[env]['password'],
        host=postgresql[env]['host'],
        port=postgresql[env]['port'],
        # sslmode='require',
    )
    cursor = connection.cursor()
    return cursor


if __name__ == "__main__":
    create_databases()