import os
from contextlib import contextmanager
from psycopg2.pool import ThreadedConnectionPool

DSN = "host=127.0.0.1 dbname={db} user={user} password={password}"

PG_DB = os.environ.get('PG_DB')
PG_USER = os.environ.get('PG_USER')
PG_PASSWORD = os.environ.get('PG_PASSWORD')

pg_pool = ThreadedConnectionPool(1, 5, DSN.format(db=PG_DB, user=PG_USER, password=PG_PASSWORD))


@contextmanager
def get_cursor(commit=True):
    with pg_pool.getconn() as conn:
        try:
            with conn.cursor() as cur:
                yield cur
        finally:
            if commit:
                conn.commit()
            pg_pool.putconn(conn)