forgot to add db
This commit is contained in:
@@ -0,0 +1,99 @@
|
|||||||
|
import sqlite3
|
||||||
|
from datetime import datetime
|
||||||
|
import json
|
||||||
|
|
||||||
|
class Database:
|
||||||
|
def __init__(self, db_path="bot.db"):
|
||||||
|
self.db_path = db_path
|
||||||
|
self.init_db()
|
||||||
|
|
||||||
|
def get_connection(self):
|
||||||
|
return sqlite3.connect(self.db_path)
|
||||||
|
|
||||||
|
def init_db(self):
|
||||||
|
with self.get_connection() as conn:
|
||||||
|
conn.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS users (
|
||||||
|
user_id INTEGER PRIMARY KEY,
|
||||||
|
username TEXT,
|
||||||
|
full_name TEXT,
|
||||||
|
language_code TEXT,
|
||||||
|
first_seen TIMESTAMP
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
conn.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS messages (
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
user_id INTEGER,
|
||||||
|
content_type TEXT,
|
||||||
|
text TEXT,
|
||||||
|
json_data TEXT,
|
||||||
|
timestamp TIMESTAMP,
|
||||||
|
FOREIGN KEY (user_id) REFERENCES users(user_id)
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
def log_user(self, user):
|
||||||
|
with self.get_connection() as conn:
|
||||||
|
conn.execute("""
|
||||||
|
INSERT OR REPLACE INTO users (user_id, username, full_name, language_code, first_seen)
|
||||||
|
VALUES (?, ?, ?, ?, ?)
|
||||||
|
""", (user.id, user.username, user.full_name, user.language_code, datetime.now()))
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
def log_message(self, message):
|
||||||
|
self.log_user(message.from_user)
|
||||||
|
with self.get_connection() as conn:
|
||||||
|
content_type = message.content_type
|
||||||
|
text = message.text or message.caption or ""
|
||||||
|
try:
|
||||||
|
json_data = message.model_dump_json()
|
||||||
|
except Exception:
|
||||||
|
# Fallback if model_dump_json fails due to non-serializable types
|
||||||
|
try:
|
||||||
|
import json
|
||||||
|
# Use model_dump to get a dict, then serialize manually
|
||||||
|
# exclude_none=True can help, and we can handle unknown types
|
||||||
|
data_dict = message.model_dump(exclude_none=True)
|
||||||
|
def custom_serializer(obj):
|
||||||
|
return str(obj)
|
||||||
|
json_data = json.dumps(data_dict, default=custom_serializer)
|
||||||
|
except Exception as e:
|
||||||
|
json_data = json.dumps({"error": f"Failed to serialize message: {str(e)}"})
|
||||||
|
|
||||||
|
cursor = conn.execute("""
|
||||||
|
INSERT INTO messages (user_id, content_type, text, json_data, timestamp)
|
||||||
|
VALUES (?, ?, ?, ?, ?)
|
||||||
|
""", (message.from_user.id, content_type, text, json_data, datetime.now()))
|
||||||
|
conn.commit()
|
||||||
|
return cursor.lastrowid
|
||||||
|
|
||||||
|
def get_stats(self):
|
||||||
|
with self.get_connection() as conn:
|
||||||
|
cursor = conn.execute("SELECT COUNT(*) FROM users")
|
||||||
|
user_count = cursor.fetchone()[0]
|
||||||
|
return {"user_count": user_count}
|
||||||
|
|
||||||
|
def get_recent_messages(self, limit=10):
|
||||||
|
with self.get_connection() as conn:
|
||||||
|
cursor = conn.execute("""
|
||||||
|
SELECT m.id, u.full_name, m.text, m.timestamp
|
||||||
|
FROM messages m
|
||||||
|
JOIN users u ON m.user_id = u.user_id
|
||||||
|
ORDER BY m.timestamp DESC
|
||||||
|
LIMIT ?
|
||||||
|
""", (limit,))
|
||||||
|
return cursor.fetchall()
|
||||||
|
|
||||||
|
def get_message_details(self, msg_id):
|
||||||
|
with self.get_connection() as conn:
|
||||||
|
cursor = conn.execute("""
|
||||||
|
SELECT m.id, u.full_name, u.username, u.user_id, m.content_type, m.timestamp, m.text
|
||||||
|
FROM messages m
|
||||||
|
JOIN users u ON m.user_id = u.user_id
|
||||||
|
WHERE m.id = ?
|
||||||
|
""", (msg_id,))
|
||||||
|
return cursor.fetchone()
|
||||||
|
|
||||||
|
db = Database()
|
||||||
Reference in New Issue
Block a user