-
Notifications
You must be signed in to change notification settings - Fork 34
/
db.py
113 lines (111 loc) · 6.21 KB
/
db.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
from psycopg_pool import AsyncConnectionPool
from os import getenv
from dotenv import load_dotenv
from typing import List, Tuple
load_dotenv()
pool = AsyncConnectionPool(conninfo=getenv("DATABASE_URL"), timeout = 10, max_lifetime=600, check=AsyncConnectionPool.check_connection, open = False)
class DataBase:
async def open_pool():
await pool.open()
await pool.wait()
async def is_user(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT user_id FROM users WHERE user_id = '{user_id}'")
result = await cursor.fetchone()
return result
async def insert_user(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute("INSERT INTO users(user_id, chatgpt, dall_e, stable_diffusion) VALUES (%s, %s, %s, %s)", (user_id,3000,3,3))
await conn.commit()
async def get_chatgpt(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT chatgpt FROM users WHERE user_id = '{user_id}'")
result = int((await cursor.fetchone())[0])
return result
async def set_chatgpt(user_id: int, result: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"UPDATE users SET chatgpt = {result} WHERE user_id = '{user_id}'")
await conn.commit()
async def get_dalle(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT dall_e FROM users WHERE user_id = '{user_id}'")
result = int((await cursor.fetchone())[0])
return result
async def set_dalle(user_id: int, result: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"UPDATE users SET dall_e = {result} WHERE user_id = '{user_id}'")
await conn.commit()
async def get_stable(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT stable_diffusion FROM users WHERE user_id = '{user_id}'")
result = int((await cursor.fetchone())[0])
return result
async def set_stable(user_id: int, result: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"UPDATE users SET stable_diffusion = {result} WHERE user_id = '{user_id}'")
await conn.commit()
async def get_userinfo(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT chatgpt, dall_e, stable_diffusion FROM users WHERE user_id = '{user_id}'")
result = await cursor.fetchone()
return result
async def new_order(invoice_id: int, user_id: int, product: str):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute("INSERT INTO orders(invoice_id, user_id, product) VALUES (%s, %s, %s)", (invoice_id, user_id, product))
await conn.commit()
async def get_orderdata(invoice_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT user_id, product FROM orders WHERE invoice_id = {invoice_id}")
result = await cursor.fetchone()
return result
async def update_chatgpt(user_id: int, invoice_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"UPDATE users SET chatgpt = chatgpt + 100000 WHERE user_id = '{user_id}'")
await cursor.execute(f"DELETE FROM orders WHERE invoice_id = {invoice_id}")
await conn.commit()
async def update_dalle(user_id: int, invoice_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"UPDATE users SET dall_e = dall_e + 50 WHERE user_id = '{user_id}'")
await cursor.execute(f"DELETE FROM orders WHERE invoice_id = {invoice_id}")
await conn.commit()
async def update_stable(user_id: int, invoice_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"UPDATE users SET stable_diffusion = stable_diffusion + 50 WHERE user_id = '{user_id}'")
await cursor.execute(f"DELETE FROM orders WHERE invoice_id = {invoice_id}")
await conn.commit()
async def save_message(user_id: int, role: str, message: str, tokens: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
conn.commit()
await cursor.execute("INSERT INTO messages(user_id, role, content, tokens) VALUES (%s, %s, %s, %s)", (user_id, role, message, tokens))
await conn.commit()
async def delete_message(message_ids: List[int]):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"DELETE FROM messages WHERE id IN ({', '.join(map(str, message_ids))})")
await conn.commit()
async def delete_messages(user_id: int):
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"DELETE FROM messages WHERE user_id = {user_id}")
await conn.commit()
async def get_messages(user_id: int) -> List[Tuple[int, str, str, int]]:
async with pool.connection() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT id, role, content, tokens FROM messages WHERE user_id = {user_id}")
result = await cursor.fetchall()
return result