from contextlib import contextmanager
import psycopg2
import psycopg2.extras
from config import Config


def get_connection():
    return psycopg2.connect(
        host=Config.DB_HOST,
        dbname=Config.DB_NAME,
        user=Config.DB_USER,
        password=Config.DB_PASSWORD,
        port=Config.DB_PORT,
        cursor_factory=psycopg2.extras.RealDictCursor,
    )


@contextmanager
def get_cursor(commit=False):
    conn = get_connection()
    cur = conn.cursor()
    try:
        yield cur
        if commit:
            conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        cur.close()
        conn.close()


def fetch_one(query, params=None):
    with get_cursor() as cur:
        cur.execute(query, params or ())
        return cur.fetchone()


def fetch_all(query, params=None):
    with get_cursor() as cur:
        cur.execute(query, params or ())
        return cur.fetchall()


def execute(query, params=None):
    with get_cursor(commit=True) as cur:
        cur.execute(query, params or ())


def execute_returning(query, params=None):
    with get_cursor(commit=True) as cur:
        cur.execute(query, params or ())
        return cur.fetchone()