import pymysql
from pprint import pprint as pprint
__author__ = "Liquidgenius"

class Query:

    def __init__(self, host, port, user, password, schema):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.schema = schema
        self.table_name = None
        self.column_names = list()
        self.query = None
        self.results = None

    def _get_db(self):
        db = (pymysql.connect(host=self.host,
                              port=self.port,
                              user=self.user,
                              password=self.password,
                              database=self.schema,
                              charset="utf8"))
        return db

    def table(self, table):
        self.column_names = list()
        self.table_name = table
        return self

    def columns(self, cols):
        if not isinstance(cols, list):
            cols = [cols]
        self.column_names = self.column_names + cols
        return self

    def execute(self):
        if (self.table_name is not None) and (len(self.column_names) > 0):
            db = (pymysql.connect(host=self.host,
                                  port=self.port,
                                  user=self.user,
                                  password=self.password,
                                  database=self.schema,
                                  charset="utf8"))
            cursor = db.cursor()
            cols = ', '.join(map(str, self.column_names))
            sql = f"SELECT {cols} FROM {self.table_name}"
            self.query = sql
            cursor.execute(sql)
            self.results = [item for item in cursor.fetchall()]
            db.close()
        return self.results

# MySQL Credentials
host =
port =
user =
password =
schema =

# Instantiate the object
q = Query(host, port, user, password, schema)

# Set parameters
table = "a_database_table"
cols = ["id", "created_at"]

# Execute with chained methods
results = q.table(table).columns(cols).execute()

# Print results
pprint(results)