-
Notifications
You must be signed in to change notification settings - Fork 0
/
sql.py
139 lines (118 loc) · 4.7 KB
/
sql.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import os
import sqlite3
class SQL(sqlite3.Cursor):
def __init__(self, db: str = "database.db", path: str = ""):
self.con = sqlite3.connect(os.path.join(path, db), timeout=5)
self.con.row_factory = (
sqlite3.Row
) # So we can use column names as keys, cause less issues if we later add a column with index values
sqlite3.Cursor.__init__(self, self.con)
self.cur = self.con.cursor()
def execute(self, command: str, *args):
with self.con:
if args:
tmp = self.con.execute(command, args[0])
else:
tmp = self.con.execute(command)
self.con.commit() # sqlite3.Connection automatically commits
# tmp = self.cur.execute(command)
# self.con.commit()
return tmp
def commit(self):
return self.con.commit()
def add_simple(self, table: str, **kwargs) -> sqlite3.Cursor:
"""Simply add values to the table, no need to specify columns
Args:
table (str): the table to add to
**kwargs: the values to add to the table
Returns:
Cursor: the cursor of the db
"""
with self.con:
# column_names = self.execute(f"PRAGMA table_info({table})").fetchall()
tmp = self.con.execute(
f"INSERT INTO {table} ({', '.join(kwargs)}) VALUES ({', '.join(['?'] * len(kwargs))})",
tuple(kwargs.values()),
)
self.con.commit()
return tmp
def add(self, table: str, data: tuple[tuple[str, type]]) -> sqlite3.Cursor:
"""Insert data into the table with tuple as (column, value)
Args:
table (str): the table to add to
data (tuple): ((column, value), (column, value))
Returns:
sqlite3.Cursor: the cursor of the db
"""
with self.con:
tmp = self.con.execute(
f"INSERT INTO {table} ({', '.join([i[0] for i in data])}) VALUES ({', '.join(['?'] * len(data))})",
tuple([i[1] for i in data]),
)
self.con.commit()
return tmp
def create_table(self, table: str, data: tuple[tuple] | str) -> sqlite3.Cursor:
if isinstance(data, str):
with self.con:
if "PRIMARY KEY" not in data:
exec_ = f"CREATE TABLE IF NOT EXISTS {table} (id INTEGER PRIMARY KEY AUTOINCREMENT, {data})"
else:
exec_ = f"CREATE TABLE IF NOT EXISTS {table} ({data})"
tmp = self.con.execute(exec_)
self.con.commit()
return tmp
types = []
for d in data:
if d[1] == int:
types.append(f"{d[0]} INTEGER")
elif d[1] == str:
types.append(f"{d[0]} TEXT")
elif d[1] == float:
types.append(f"{d[0]} REAL")
elif d[1] == bool:
types.append(f"{d[0]} BOOLEAN")
else:
# raise TypeError(f"Type {d[1]} is not supported")
types.append(f"{d[0]} {d[1]}")
types = ", ".join(types)
if "PRIMARY KEY" not in types:
exec_ = f"CREATE TABLE IF NOT EXISTS {table} (id INTEGER PRIMARY KEY AUTOINCREMENT, {types})"
else:
exec_ = f"CREATE TABLE IF NOT EXISTS {table} ({types})"
with self.con:
self.con.execute(exec_)
def add_safe(self, table: str, data: tuple[tuple]) -> sqlite3.Cursor:
"""If the table does not exist, create it. Insert data into the table with tuple as (column, value)
Args:
table (str): the table to add to
data (tuple): ((column, value), (column, value))
Returns:
sqlite3.Cursor: the cursor of the db
"""
with self.con:
if not self.table_exists(table):
tmp = self.create_table(table, data)
tmp = self.add(table, data)
self.con.commit()
return tmp
def table_exists(self, table: str) -> bool:
"""Check if a table exists
Args:
table (str): the table to check
Returns:
bool: if the table exists
"""
tmp = self.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table,)
)
return bool(tmp.fetchall())
# def insert(conn, command: str):
# with conn.cursor() as c:
# c.execute(...)
# conn.commit()
# def execute_soft(self, command: str):
# return self.cur.execute(command)
def executemany(self, command: str, args: list[tuple]):
tmp = self.cur.executemany(command, args)
self.con.commit()
return tmp