import psycopg2
from psycopg2._psycopg import connection
import pandas as pd


def get_connection(username: str,
                   password: str,
                   host: str = "reporting-db.smeanalyticsapps.com",
                   dbname="smaredshiftdb",
                   port: int = 5432) -> connection:
    return psycopg2.connect(username=username,
                            password=password,
                            host=host,
                            dbname=dbname,
                            port=port)

def query_to_df(sqlquery: str, connection: connection) -> pd.DataFrame:
    """
    Returns a SQL query via a connection

    :param sqlquery: the SQL query to perform
    :param connection: the connection over which to query (use get_connection to create)
    :return: Pandas DataFrame object with the results and correct column names
    """
    crs = connection.cursor()
    _raw = crs.execute(sqlquery)

    col_names = [elt[0] for elt in crs.description]

    df = pd.DataFrame(crs.fetchall(), columns=col_names)

    crs.close()
    connection.close()

    return df