VkszQP4S

1"""
2Manages the database for all of Augusta's users
3"""
4
5import sqlite3
6import csv
7import os
8
9class Database(object):
10    """
11    Object used to abstract away the underlying SQLite database of users and their information
12    """
13
14    def __init__(self):
15        self.connection = None
16
17    # ------------------------------------ GENERIC DATABASE FUNCTIONS ---------------------------------
18
19    def connect(self, database):
20        """
21        Connects to the specified database
22
23        :param database: the database to connect to
24        :return: None
25        """
26        try:
27            self.connection = sqlite3.connect(database)
28        except sqlite3.Error as e:
29            print("Connection failed.")
30            print("Error:", e.args[0])
31
32    def disconnect(self):
33        """
34        Closes the connection that was previously connected to
35
36        :return: None
37        """
38        try:
39            self.connection.close()
40        except sqlite3.Error as e:
41            print("Connection failed to close...")
42            print("Error: ", e.args[0])
43
44    def create_table(self, table, **kwargs):
45        """
46        Creates a table in the DB
47
48        :param table:   The name of the table
49        :param kwargs:  The column information in the form (name, ATTRB). ATTRB can dictate any attribute of the column
50        :return: (True, "Success") if the table was created successfully. (False, "Message") if otherwise
51        """
52        query = "CREATE TABLE {name} (".format(name=table)
53        for (name, attrb) in kwargs.items():
54            query += "{} {},".format(name, attrb)
55        query = query[:-1] + ')'
56
57        return self.execute(query)
58
59    def drop_table(self, table):
60        """
61        Drops a specified table in the DB
62
63        :param table: the table that needs dropping
64        :return: None
65        """
66        query = "DROP TABLE {}".format(table)
67
68        return self.execute(query)
69
70    def insert(self, table, *data):
71        """
72        Inserts the specified data into the table
73        :param table:   the table to insert data into
74        :param data:    the data
75        :return: (True, "Success") if insertion was successful. (False, "Message") if not. Raises an exception if
76        data already exists
77        """
78        query = "INSERT INTO {table} VALUES (".format(table=table)
79        for value in data:
80            query += "\'{}\',".format(value)
81        query = query[:-1] + ')'
82
83        return self.execute(query)
84
85    def execute(self, query):
86        """
87        Unsafe query executor. Don't use this unless the query has been preprocessed.
88
89        :param query: The query to be executed
90        :return: (True, "Success") if the query executed fine. (False, "Message") otherwise
91        """
92        try:
93            with self.connection:
94                self.connection.execute(query)
95        except sqlite3.Error as e:
96            return (False, e.args[0])
97        return (True, "Success")
98
99    # ------------------------------------ BOT SPECIFIC FUNCTIONS -------------------------------------