import psycopg2
# from psycopg2 import sql
from psycopg2.extensions import AsIs

import config


class MyDataBase:
    def __init__(self, tb_name):
        self.table = tb_name
        self.connect = psycopg2.connect(dbname="test", host="localhost", user=config.user, password=config.password)
        self.connect.autocommit = True
        self.cursor = self.connect.cursor()
        self.cursor.execute("DROP TABLE IF EXISTS %s", (AsIs(self.table),))
        self.cursor.execute("CREATE TABLE %s (id integer, msg_count integer)", (AsIs(self.table),))


    def insert(self, user_id, msg_count):
        self.cursor.execute("INSERT INTO %s (id, msg_count) VALUES (%s, %s)", (AsIs(self.table), user_id, msg_count))


    def update(self, user_id, msg_count):
        self.cursor.execute("UPDATE %s SET msg_count=%s WHERE id=%s", (AsIs(self.table), msg_count, user_id))


    def outprint(self):
        self.cursor.execute("SELECT * FROM %s;", (AsIs(self.table),))
        print(self.cursor.fetchall())


    def closedb(self):
        self.cursor.close()
        self.connect.close()