First Commit
This commit is contained in:
Executable
Executable
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Executable
+89
@@ -0,0 +1,89 @@
|
||||
import discord
|
||||
from utils.sql_commands import DatabaseManager
|
||||
from typing import Dict
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
db = DatabaseManager("economy")
|
||||
|
||||
|
||||
async def create_account(user: discord.Member | discord.User):
|
||||
"""Create a new economy account for a user."""
|
||||
db.insert(
|
||||
"INSERT INTO economy (ID, WALLET, BANK) VALUES (%s, %s, %s)", (user.id, 0, 1000)
|
||||
)
|
||||
|
||||
|
||||
async def bank_data(user: discord.Member | discord.User) -> Dict[str, int]:
|
||||
"""Get the bank data for a user. If the user has no account, create one.
|
||||
|
||||
Args:
|
||||
user: The user to get the bank data for.
|
||||
|
||||
Returns:
|
||||
A dictionary with the keys ``WALLET`` and ``BANK`` containing the user's wallet and bank balances.
|
||||
"""
|
||||
balance = db.fetch_one("SELECT * FROM economy WHERE ID = %s", (user.id,))
|
||||
if balance is None:
|
||||
await create_account(user)
|
||||
return await bank_data(user)
|
||||
return balance
|
||||
|
||||
|
||||
async def record_transaction(
|
||||
user_id: int, transaction_type: str, amount: float
|
||||
) -> None:
|
||||
"""Record a transaction in the database.
|
||||
|
||||
Args:
|
||||
user_id: The ID of the user who performed the transaction.
|
||||
transaction_type: The type of transaction performed.
|
||||
amount: The amount of the transaction.
|
||||
"""
|
||||
db.insert("transactions", (user_id, transaction_type, amount))
|
||||
|
||||
|
||||
async def update_money(
|
||||
user: discord.Member | discord.User, bank: int = 0, wallet: int = 0
|
||||
):
|
||||
"""Update the bank balance of a user.
|
||||
|
||||
Args:
|
||||
user: The user to update the balance for.
|
||||
bank: The amount to add to the user's bank balance.
|
||||
wallet: The amount to add to the user's wallet balance.
|
||||
|
||||
Returns:
|
||||
A dictionary with the keys ``WALLET`` and ``BANK`` containing the updated balances.
|
||||
"""
|
||||
balance = await bank_data(user)
|
||||
if bank or wallet:
|
||||
new_wallet = balance.get("WALLET", 0) + wallet
|
||||
new_bank = balance.get("BANK", 0) + bank
|
||||
db.insert(
|
||||
"INSERT INTO economy (ID, WALLET, BANK) VALUES (%s, %s, %s)",
|
||||
(user.id, new_wallet, new_bank),
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
|
||||
async def reset_bank(user: discord.Member | discord.User) -> None:
|
||||
"""Reset the bank account for a user.
|
||||
|
||||
Args:
|
||||
user (discord.Member | discord.User): The user whose bank account is to be reset.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
db.delete("economy", {"ID": user.id})
|
||||
await create_account(user)
|
||||
|
||||
|
||||
async def update_daily_timestamp(user: discord.User | discord.Member, timestamp: datetime) -> None:
|
||||
"""
|
||||
Updates the DAILY_TIMESTAMP field for the user in the economy table.
|
||||
Stores the timestamp as a float (UNIX time).
|
||||
"""
|
||||
db.execute_query("UPDATE economy SET DAILY = %s WHERE ID = %s",(timestamp.timestamp(), user.id),
|
||||
)
|
||||
Executable
+361
@@ -0,0 +1,361 @@
|
||||
import mysql.connector
|
||||
from mysql.connector import pooling
|
||||
from dotenv import load_dotenv
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
import logging
|
||||
import pandas as pd
|
||||
from copy import copy
|
||||
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
filename="database.log",
|
||||
format="%(asctime)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
def __init__(self, env="development"):
|
||||
# Load environment variables based on environment
|
||||
self.load_env(env)
|
||||
|
||||
self.config = {
|
||||
"host": os.getenv("SQLHOST", "localhost"),
|
||||
"user": os.getenv("SQLUSER", "root"),
|
||||
"password": os.getenv("SQLPASS", ""),
|
||||
"database": os.getenv("SQLDB", "testdb"),
|
||||
"pool_reset_session": bool(os.getenv("POOL_RESET_SESSION", False)),
|
||||
}
|
||||
|
||||
self.pool = pooling.MySQLConnectionPool(
|
||||
pool_name="mypool", pool_size=5, **self.config
|
||||
)
|
||||
logger.info("Database connection pool created.")
|
||||
|
||||
# Ensure required tables exist
|
||||
self.create_table_if_not_exists(
|
||||
"feedback",
|
||||
"""
|
||||
ID INT AUTO_INCREMENT PRIMARY KEY,
|
||||
USER VARCHAR(100),
|
||||
GUILDID BIGINT,
|
||||
TIMESTAMP VARCHAR(32),
|
||||
CONTENT TEXT
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"afk_status",
|
||||
"""
|
||||
USERID BIGINT,
|
||||
GUILDID BIGINT,
|
||||
TIMESTAMP VARCHAR(32),
|
||||
REASON TEXT,
|
||||
PRIMARY KEY (USERID, GUILDID)
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_tickets",
|
||||
"""
|
||||
ID INT AUTO_INCREMENT PRIMARY KEY,
|
||||
USERID BIGINT NOT NULL,
|
||||
TIMESTAMP DATETIME NOT NULL,
|
||||
TICKET_TYPE VARCHAR(32) NOT NULL,
|
||||
group_id VARCHAR(32) DEFAULT NULL
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_results",
|
||||
"""
|
||||
ID INT AUTO_INCREMENT PRIMARY KEY,
|
||||
WINNER_ID BIGINT NOT NULL,
|
||||
AMOUNT INT NOT NULL,
|
||||
DRAW_TIME DATETIME NOT NULL,
|
||||
CLAIMED BOOLEAN DEFAULT 0,
|
||||
WIN_TYPE VARCHAR(16) DEFAULT 'user'
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_state",
|
||||
"""
|
||||
id INT PRIMARY KEY,
|
||||
jackpot INT NOT NULL
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_draw_time",
|
||||
"""
|
||||
id INT PRIMARY KEY,
|
||||
last_draw DATETIME NOT NULL
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_luck",
|
||||
"""
|
||||
USERID BIGINT PRIMARY KEY,
|
||||
LUCK INT NOT NULL DEFAULT 0
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_groups",
|
||||
"""
|
||||
group_id VARCHAR(32) PRIMARY KEY,
|
||||
creator_id BIGINT NOT NULL
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"lottery_group_members",
|
||||
"""
|
||||
group_id VARCHAR(32),
|
||||
user_id BIGINT,
|
||||
PRIMARY KEY (group_id, user_id)
|
||||
""",
|
||||
)
|
||||
self.create_table_if_not_exists(
|
||||
"gamble_rooms",
|
||||
"""
|
||||
channel_id BIGINT PRIMARY KEY,
|
||||
host_id BIGINT NOT NULL,
|
||||
invited TEXT NOT NULL,
|
||||
inactivity INT NOT NULL
|
||||
""",
|
||||
)
|
||||
|
||||
def load_env(self, env):
|
||||
env_file = f".env"
|
||||
load_dotenv(env_file)
|
||||
logger.info(f"Loaded environment variables from {env_file}")
|
||||
|
||||
def get_connection(self):
|
||||
return self.pool.get_connection()
|
||||
|
||||
def execute_query(self, query, params=None, retries=3, delay=1):
|
||||
cursor = None
|
||||
connection = None
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
connection = self.get_connection()
|
||||
cursor = connection.cursor(dictionary=True, buffered=True)
|
||||
cursor.execute(query, params or ())
|
||||
logger.info(f"Executed query: {query} with params: {params}")
|
||||
connection.commit()
|
||||
return copy(cursor)
|
||||
except mysql.connector.Error as err:
|
||||
logger.warning(f"Attempt {attempt + 1} failed: {err}")
|
||||
time.sleep(delay * (2**attempt))
|
||||
finally:
|
||||
if cursor:
|
||||
cursor.close()
|
||||
if connection:
|
||||
connection.close()
|
||||
|
||||
logger.error(f"All {retries} attempts failed for query: {query}")
|
||||
return None
|
||||
|
||||
def insert(self, query: str, params: tuple, overwrite: bool = True) -> None:
|
||||
"""
|
||||
Inserts data into the database using the given query and parameters.
|
||||
|
||||
Args:
|
||||
query (str): The SQL query to execute.
|
||||
params (tuple): The parameters to pass into the query.
|
||||
overwrite (bool, optional): Whether to perform an upsert operation. Defaults to False.
|
||||
|
||||
Raises:
|
||||
ValueError: If no parameters are provided.
|
||||
"""
|
||||
if not params:
|
||||
raise ValueError("Params must be provided for the insert operation.")
|
||||
|
||||
try:
|
||||
if overwrite:
|
||||
columns = [
|
||||
col.split("=")[0].strip()
|
||||
for col in query.split("VALUES")[0]
|
||||
.split("(")[1]
|
||||
.split(")")[0]
|
||||
.split(",")
|
||||
]
|
||||
update_set = ", ".join(f"{col} = VALUES({col})" for col in columns)
|
||||
query = f"{query} ON DUPLICATE KEY UPDATE {update_set}"
|
||||
|
||||
cursor = self.execute_query(query, params)
|
||||
if cursor:
|
||||
logger.info(f"Insert completed with query: {query}.")
|
||||
except mysql.connector.Error as err:
|
||||
logger.error(f"Insert failed with query: {query}. Error: {err}")
|
||||
|
||||
def bulk_insert(self, query, params=None):
|
||||
if not params:
|
||||
logger.warning("No data provided for bulk insert.")
|
||||
return
|
||||
|
||||
# Assuming params is a list of dictionaries
|
||||
if not isinstance(params, list) or not all(isinstance(d, dict) for d in params):
|
||||
raise ValueError("Params must be a list of dictionaries for bulk insert.")
|
||||
|
||||
keys = params[0].keys()
|
||||
placeholders = ", ".join(["%s"] * len(keys))
|
||||
query = f"{query} ({', '.join(keys)}) VALUES ({placeholders})"
|
||||
values = [tuple(data.values()) for data in params]
|
||||
|
||||
try:
|
||||
connection = self.get_connection()
|
||||
cursor = connection.cursor()
|
||||
cursor.executemany(query, values)
|
||||
connection.commit()
|
||||
logger.info(
|
||||
f"Bulk insert completed for {len(params)} records with query: {query}."
|
||||
)
|
||||
except mysql.connector.Error as err:
|
||||
logger.error(f"Bulk insert failed: {err}")
|
||||
connection.rollback() # Roll back on error
|
||||
finally:
|
||||
cursor.close()
|
||||
connection.close()
|
||||
|
||||
def delete(self, table_name: str, condition: dict) -> None:
|
||||
"""Deletes a record from the specified table based on the condition provided."""
|
||||
condition_column, condition_value = next(iter(condition.items()))
|
||||
query = f"DELETE FROM {table_name} WHERE {condition_column} = %s"
|
||||
self.execute_query(query, (condition_value,))
|
||||
|
||||
def fetch_one(self, query, params=None):
|
||||
cursor = self.execute_query(query, params)
|
||||
return cursor.fetchone() if cursor else {}
|
||||
|
||||
def fetch_all(self, query, params=None):
|
||||
cursor = self.execute_query(query, params)
|
||||
return cursor.fetchall() if cursor else []
|
||||
|
||||
def fetch_as_dataframe(self, query, params=None):
|
||||
cursor = self.execute_query(query, params)
|
||||
if cursor:
|
||||
try:
|
||||
# Ensure cursor has a result to fetch
|
||||
if cursor.with_rows:
|
||||
results = cursor.fetchall()
|
||||
return pd.DataFrame(results) if results else pd.DataFrame()
|
||||
else:
|
||||
logger.warning("No result set to fetch from.")
|
||||
return pd.DataFrame()
|
||||
finally:
|
||||
cursor.close()
|
||||
return pd.DataFrame()
|
||||
|
||||
def create_table_if_not_exists(self, table_name, schema):
|
||||
query = f"CREATE TABLE IF NOT EXISTS {table_name} ({schema})"
|
||||
self.execute_query(query)
|
||||
logger.info(f"Ensured table {table_name} exists with schema: {schema}")
|
||||
|
||||
def cleanup_old_transactions(self, days=365):
|
||||
cutoff_date = datetime.now() - timedelta(days=days)
|
||||
self.execute_query("DELETE FROM transactions WHERE TIME < %s", (cutoff_date,))
|
||||
logger.info(f"Old transactions older than {days} days have been cleaned up.")
|
||||
|
||||
@staticmethod
|
||||
def random_timestamp(start_date, end_date):
|
||||
time_between_dates = end_date - start_date
|
||||
random_date = start_date + timedelta(
|
||||
days=random.randint(0, time_between_dates.days)
|
||||
)
|
||||
return random_date
|
||||
|
||||
def build_query(self, base_query, filters):
|
||||
conditions = " AND ".join([f"{k} = %s" for k in filters.keys()])
|
||||
return f"{base_query} WHERE {conditions}", list(filters.values())
|
||||
|
||||
|
||||
# SQL scripts to create tables
|
||||
create_feedback_table = """
|
||||
CREATE TABLE IF NOT EXISTS feedback (
|
||||
ID INT AUTO_INCREMENT PRIMARY KEY,
|
||||
USER VARCHAR(100),
|
||||
GUILDID BIGINT,
|
||||
TIMESTAMP VARCHAR(32),
|
||||
CONTENT TEXT
|
||||
);
|
||||
"""
|
||||
|
||||
create_afk_status_table = """
|
||||
CREATE TABLE IF NOT EXISTS afk_status (
|
||||
USERID BIGINT,
|
||||
GUILDID BIGINT,
|
||||
REASON TEXT,
|
||||
PRIMARY KEY (USERID, GUILDID)
|
||||
);
|
||||
"""
|
||||
|
||||
create_lottery_tickets_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_tickets (
|
||||
ID INT AUTO_INCREMENT PRIMARY KEY,
|
||||
USERID BIGINT NOT NULL,
|
||||
TIMESTAMP DATETIME NOT NULL,
|
||||
TICKET_TYPE VARCHAR(32) NOT NULL,
|
||||
group_id VARCHAR(32) DEFAULT NULL
|
||||
);
|
||||
"""
|
||||
|
||||
create_lottery_results_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_results (
|
||||
ID INT AUTO_INCREMENT PRIMARY KEY,
|
||||
WINNER_ID BIGINT NOT NULL,
|
||||
AMOUNT INT NOT NULL,
|
||||
DRAW_TIME DATETIME NOT NULL,
|
||||
CLAIMED BOOLEAN DEFAULT 0,
|
||||
WIN_TYPE VARCHAR(16) DEFAULT 'user'
|
||||
);
|
||||
"""
|
||||
|
||||
create_lottery_state_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_state (
|
||||
id INT PRIMARY KEY,
|
||||
jackpot INT NOT NULL
|
||||
);
|
||||
-- Insert initial row if not exists:
|
||||
INSERT IGNORE INTO lottery_state (id, jackpot) VALUES (1, 0);
|
||||
"""
|
||||
|
||||
create_lottery_draw_time_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_draw_time (
|
||||
id INT PRIMARY KEY,
|
||||
last_draw DATETIME NOT NULL
|
||||
);
|
||||
-- Insert initial row if not exists:
|
||||
INSERT IGNORE INTO lottery_draw_time (id, last_draw) VALUES (1, NOW());
|
||||
"""
|
||||
|
||||
create_lottery_luck_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_luck (
|
||||
USERID BIGINT PRIMARY KEY,
|
||||
LUCK INT NOT NULL DEFAULT 0
|
||||
);
|
||||
"""
|
||||
|
||||
create_lottery_groups_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_groups (
|
||||
group_id VARCHAR(32) PRIMARY KEY,
|
||||
creator_id BIGINT NOT NULL
|
||||
);
|
||||
"""
|
||||
|
||||
create_lottery_group_members_table = """
|
||||
CREATE TABLE IF NOT EXISTS lottery_group_members (
|
||||
group_id VARCHAR(32),
|
||||
user_id BIGINT,
|
||||
PRIMARY KEY (group_id, user_id)
|
||||
);
|
||||
"""
|
||||
|
||||
create_gamble_rooms_table = """
|
||||
CREATE TABLE IF NOT EXISTS gamble_rooms (
|
||||
channel_id BIGINT PRIMARY KEY,
|
||||
host_id BIGINT NOT NULL,
|
||||
invited TEXT NOT NULL,
|
||||
inactivity INT NOT NULL
|
||||
);
|
||||
"""
|
||||
Reference in New Issue
Block a user