You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

604 lines
25 KiB

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

"""
数据库操作模块
"""
import os
import sqlite3
import hashlib
from datetime import datetime, timezone, timedelta
from config import config
def beijing_now_str(fmt='%Y-%m-%d %H:%M:%S'):
"""返回北京时间字符串"""
return datetime.now(timezone(timedelta(hours=8))).strftime(fmt)
class Database:
def __init__(self):
self.db_path = config.get('database')
db_dir = os.path.dirname(self.db_path)
if db_dir and not os.path.exists(db_dir):
os.makedirs(db_dir, exist_ok=True)
self.init_database()
def get_connection(self):
conn = sqlite3.connect(self.db_path, check_same_thread=False)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA foreign_keys=ON")
return conn
def init_database(self):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username VARCHAR(50) UNIQUE NOT NULL,
password VARCHAR(100) NOT NULL,
nickname VARCHAR(50),
avatar VARCHAR(200) DEFAULT 'default.png',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS friendships (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL,
friend_id INTEGER NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (user_id) REFERENCES users(id),
FOREIGN KEY (friend_id) REFERENCES users(id),
UNIQUE(user_id, friend_id)
)''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
sender_id INTEGER NOT NULL,
receiver_id INTEGER,
group_id INTEGER,
content TEXT NOT NULL,
msg_type VARCHAR(20) DEFAULT 'text',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (sender_id) REFERENCES users(id)
)''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS groups (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name VARCHAR(100) NOT NULL,
creator_id INTEGER NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (creator_id) REFERENCES users(id)
)''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS group_members (
id INTEGER PRIMARY KEY AUTOINCREMENT,
group_id INTEGER NOT NULL,
user_id INTEGER NOT NULL,
join_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (group_id) REFERENCES groups(id),
FOREIGN KEY (user_id) REFERENCES users(id),
UNIQUE(group_id, user_id)
)''')
# 默认管理员与测试用户
default_users = [
('admin', '123456', '系统管理员'),
('user1', '123456', '用户一'),
('user2', '123456', '用户二'),
('user3', '123456', '用户三'),
]
for username, password, nickname in default_users:
cursor.execute("SELECT id FROM users WHERE username=?", (username,))
if not cursor.fetchone():
pw = hashlib.md5(password.encode()).hexdigest()
cursor.execute(
"INSERT INTO users (username, password, nickname) VALUES (?,?,?)",
(username, pw, nickname))
conn.commit()
# 为新数据库填充测试数据:好友关系、群组、聊天记录
self._seed_test_data(conn)
conn.close()
def _seed_test_data(self, conn):
"""为新数据库填充测试数据(仅在数据为空时执行)"""
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM messages")
if cursor.fetchone()[0] > 0:
return # 已有数据,跳过
now = beijing_now_str()
cursor.execute("SELECT id, username FROM users")
users = {r[1]: r[0] for r in cursor.fetchall()}
uid = users.get
# 好友关系
friendships = [
('admin', 'user1'), ('admin', 'user2'),
('user1', 'user2'), ('user1', 'user3'), ('user2', 'user3'),
]
for a, b in friendships:
if a in users and b in users:
cursor.execute("INSERT OR IGNORE INTO friendships (user_id, friend_id) VALUES (?,?)",
(users[a], users[b]))
cursor.execute("INSERT OR IGNORE INTO friendships (user_id, friend_id) VALUES (?,?)",
(users[b], users[a]))
# 群组
cursor.execute("INSERT INTO groups (name, creator_id) VALUES (?,?)", ('技术交流群', users['admin']))
gid = cursor.lastrowid
for name in ['admin', 'user1', 'user2', 'user3']:
if name in users:
cursor.execute("INSERT OR IGNORE INTO group_members (group_id, user_id) VALUES (?,?)",
(gid, users[name]))
# 私聊消息
private_msgs = [
('admin', 'user1', '你好 user1欢迎使用 SimpleChat'),
('user1', 'admin', '你好管理员,这个系统看起来很不错!'),
('admin', 'user1', '是的,支持私聊、群聊、好友管理等功能'),
('user1', 'admin', '太棒了,我试试看群聊功能'),
('admin', 'user2', 'user2最近在忙什么呢'),
('user2', 'admin', '在研究 Python 网络编程,这个聊天系统是用什么写的?'),
('admin', 'user2', '用 Python + PyQt5 写的TCP 通信'),
('user2', 'admin', '厉害!我也想学习一下'),
('user1', 'user2', 'hi user2一起加个群聊吧'),
('user2', 'user1', '好啊,我正想找人聊天呢'),
('user1', 'user3', 'user3好久不见'),
('user3', 'user1', '是啊好久不见,最近怎么样?'),
('user1', 'user3', '挺好的,这个聊天系统不错'),
('user3', 'user1', '确实,界面也挺好看的'),
]
for i, (sen, rec, content) in enumerate(private_msgs):
if sen in users and rec in users:
ts = f'2026-05-14 {9 + i // 6:02d}:{30 + i % 6 * 8:02d}:00'
cursor.execute(
"INSERT INTO messages (sender_id, receiver_id, content, created_at) VALUES (?,?,?,?)",
(users[sen], users[rec], content, ts))
# 群聊消息
group_msgs = [
('admin', '欢迎大家加入技术交流群!'),
('user1', '大家好,我是 user1很高兴加入'),
('user2', '大家好!我是 user2'),
('user3', '各位好,我是 user3'),
('admin', '大家可以在群里交流技术问题'),
('user1', 'Python 的异步编程大家有研究吗?'),
('user2', '我最近在看 asyncio感觉挺有意思的'),
('admin', 'asyncio 确实强大,适合高并发场景'),
('user3', '我主要做前端PyQt5 的界面开发也可以交流'),
]
for i, (sen, content) in enumerate(group_msgs):
if sen in users:
ts = f'2026-05-14 {9 + i // 4:02d}:{40 + i % 4 * 8:02d}:00'
cursor.execute(
"INSERT INTO messages (sender_id, group_id, content, created_at) VALUES (?,?,?,?)",
(users[sen], gid, content, ts))
conn.commit()
# ── 用户 ──────────────────────────────────────────────
def user_register(self, username, password, nickname=None):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("SELECT id FROM users WHERE username=?", (username,))
if cursor.fetchone():
return False, "用户名已存在"
pw = hashlib.md5(password.encode()).hexdigest()
cursor.execute(
"INSERT INTO users (username, password, nickname) VALUES (?,?,?)",
(username, pw, nickname or username))
conn.commit()
return True, "注册成功"
except Exception as e:
return False, f"注册失败: {e}"
finally:
conn.close()
def user_login(self, username, password):
conn = self.get_connection()
cursor = conn.cursor()
try:
pw = hashlib.md5(password.encode()).hexdigest()
cursor.execute(
"SELECT id, username, nickname, avatar FROM users WHERE username=? AND password=?",
(username, pw))
user = cursor.fetchone()
if user:
return True, "登录成功", dict(user)
return False, "用户名或密码错误", None
except Exception as e:
return False, f"登录失败: {e}", None
finally:
conn.close()
def change_username(self, user_id, new_username):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("SELECT id FROM users WHERE username=? AND id!=?", (new_username, user_id))
if cursor.fetchone():
return False, "用户名已被占用"
cursor.execute("UPDATE users SET username=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
(new_username, user_id))
conn.commit()
return True, "用户名修改成功"
except Exception as e:
return False, f"修改失败: {e}"
finally:
conn.close()
def change_nickname(self, user_id, new_nickname):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("UPDATE users SET nickname=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
(new_nickname, user_id))
conn.commit()
return True, "昵称修改成功"
except Exception as e:
return False, f"修改失败: {e}"
finally:
conn.close()
def change_password(self, user_id, old_password, new_password):
conn = self.get_connection()
cursor = conn.cursor()
try:
old_pw = hashlib.md5(old_password.encode()).hexdigest()
cursor.execute("SELECT id FROM users WHERE id=? AND password=?", (user_id, old_pw))
if not cursor.fetchone():
return False, "原密码错误"
new_pw = hashlib.md5(new_password.encode()).hexdigest()
cursor.execute("UPDATE users SET password=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
(new_pw, user_id))
conn.commit()
return True, "密码修改成功"
except Exception as e:
return False, f"修改失败: {e}"
finally:
conn.close()
def get_user_by_id(self, user_id):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute(
"SELECT id, username, nickname, avatar FROM users WHERE id=?", (user_id,))
user = cursor.fetchone()
conn.close()
return dict(user) if user else None
def get_user_by_username(self, username):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute(
"SELECT id, username, nickname, avatar FROM users WHERE username=?", (username,))
user = cursor.fetchone()
conn.close()
return dict(user) if user else None
def get_all_users(self, exclude_id=None):
conn = self.get_connection()
cursor = conn.cursor()
if exclude_id:
cursor.execute(
"SELECT id, username, nickname, avatar FROM users WHERE id!=? ORDER BY username",
(exclude_id,))
else:
cursor.execute(
"SELECT id, username, nickname, avatar FROM users ORDER BY username")
users = [dict(r) for r in cursor.fetchall()]
conn.close()
return users
def search_users(self, keyword, exclude_id=None):
"""搜索用户(用户名或昵称模糊匹配)"""
conn = self.get_connection()
cursor = conn.cursor()
like = f"%{keyword}%"
if exclude_id:
cursor.execute(
"SELECT id, username, nickname, avatar FROM users "
"WHERE id!=? AND (username LIKE ? OR nickname LIKE ?) ORDER BY username",
(exclude_id, like, like))
else:
cursor.execute(
"SELECT id, username, nickname, avatar FROM users "
"WHERE username LIKE ? OR nickname LIKE ? ORDER BY username",
(like, like))
users = [dict(r) for r in cursor.fetchall()]
conn.close()
return users
# ── 好友 ──────────────────────────────────────────────
def add_friend(self, user_id, friend_username):
conn = self.get_connection()
cursor = conn.cursor()
try:
friend = self.get_user_by_username(friend_username)
if not friend:
return False, "用户不存在"
friend_id = friend['id']
if user_id == friend_id:
return False, "不能添加自己为好友"
cursor.execute(
"SELECT id FROM friendships WHERE (user_id=? AND friend_id=?) OR (user_id=? AND friend_id=?)",
(user_id, friend_id, friend_id, user_id))
if cursor.fetchone():
return False, "已经是好友"
cursor.execute("INSERT INTO friendships (user_id, friend_id) VALUES (?,?)", (user_id, friend_id))
cursor.execute("INSERT INTO friendships (user_id, friend_id) VALUES (?,?)", (friend_id, user_id))
conn.commit()
return True, "添加好友成功"
except Exception as e:
return False, f"添加好友失败: {e}"
finally:
conn.close()
def get_friends(self, user_id):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
SELECT u.id, u.username, u.nickname, u.avatar
FROM users u JOIN friendships f ON u.id = f.friend_id
WHERE f.user_id=? ORDER BY u.username
''', (user_id,))
friends = [dict(r) for r in cursor.fetchall()]
conn.close()
return friends
def remove_friend(self, user_id, friend_id):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute(
"DELETE FROM friendships WHERE (user_id=? AND friend_id=?) OR (user_id=? AND friend_id=?)",
(user_id, friend_id, friend_id, user_id))
conn.commit()
return True, "已删除好友"
except Exception as e:
return False, str(e)
finally:
conn.close()
# ── 消息 ──────────────────────────────────────────────
def save_message(self, sender_id, receiver_id, content, msg_type='text', group_id=None):
conn = self.get_connection()
cursor = conn.cursor()
try:
now = beijing_now_str()
cursor.execute(
"INSERT INTO messages (sender_id, receiver_id, group_id, content, msg_type, created_at) VALUES (?,?,?,?,?,?)",
(sender_id, receiver_id, group_id, content, msg_type, now))
conn.commit()
return True, cursor.lastrowid
except Exception as e:
return False, str(e)
finally:
conn.close()
def get_chat_history(self, user1_id, user2_id, limit=50):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
SELECT m.id, m.content, m.created_at, u.username, u.nickname
FROM messages m JOIN users u ON m.sender_id = u.id
WHERE ((m.sender_id=? AND m.receiver_id=?) OR (m.sender_id=? AND m.receiver_id=?))
AND m.group_id IS NULL
ORDER BY m.created_at ASC LIMIT ?
''', (user1_id, user2_id, user2_id, user1_id, limit))
messages = [dict(r) for r in cursor.fetchall()]
conn.close()
return messages
def get_group_history(self, group_id, limit=50):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
SELECT m.id, m.content, m.created_at, u.username, u.nickname
FROM messages m JOIN users u ON m.sender_id = u.id
WHERE m.group_id=?
ORDER BY m.created_at ASC LIMIT ?
''', (group_id, limit))
messages = [dict(r) for r in cursor.fetchall()]
conn.close()
return messages
# ── 群组 ──────────────────────────────────────────────
def create_group(self, group_name, creator_id):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("INSERT INTO groups (name, creator_id) VALUES (?,?)", (group_name, creator_id))
group_id = cursor.lastrowid
cursor.execute("INSERT INTO group_members (group_id, user_id) VALUES (?,?)", (group_id, creator_id))
conn.commit()
return True, group_id
except Exception as e:
return False, str(e)
finally:
conn.close()
def get_user_groups(self, user_id):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
SELECT g.id, g.name, g.creator_id
FROM groups g JOIN group_members gm ON g.id = gm.group_id
WHERE gm.user_id=? ORDER BY g.name
''', (user_id,))
groups = [dict(r) for r in cursor.fetchall()]
conn.close()
return groups
def get_group_members(self, group_id):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
SELECT u.id, u.username, u.nickname, u.avatar
FROM users u JOIN group_members gm ON u.id = gm.user_id
WHERE gm.group_id=? ORDER BY u.username
''', (group_id,))
members = [dict(r) for r in cursor.fetchall()]
conn.close()
return members
def join_group(self, group_id, user_id):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("SELECT id FROM groups WHERE id=?", (group_id,))
if not cursor.fetchone():
return False, "群组不存在"
cursor.execute(
"INSERT OR IGNORE INTO group_members (group_id, user_id) VALUES (?,?)",
(group_id, user_id))
conn.commit()
return True, "加入成功"
except Exception as e:
return False, str(e)
finally:
conn.close()
def get_all_groups(self):
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute("SELECT id, name, creator_id FROM groups ORDER BY name")
groups = [dict(r) for r in cursor.fetchall()]
conn.close()
return groups
def leave_group(self, group_id, user_id):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("SELECT id FROM groups WHERE id=?", (group_id,))
if not cursor.fetchone():
return False, "群组不存在"
cursor.execute(
"DELETE FROM group_members WHERE group_id=? AND user_id=?",
(group_id, user_id))
if cursor.rowcount == 0:
return False, "你不在该群组中"
conn.commit()
return True, "已退出群组"
except Exception as e:
return False, str(e)
finally:
conn.close()
def invite_to_group(self, group_id, username):
conn = self.get_connection()
cursor = conn.cursor()
try:
cursor.execute("SELECT id FROM groups WHERE id=?", (group_id,))
if not cursor.fetchone():
return False, "群组不存在"
user = self.get_user_by_username(username)
if not user:
return False, "用户不存在"
cursor.execute(
"SELECT id FROM group_members WHERE group_id=? AND user_id=?",
(group_id, user['id']))
if cursor.fetchone():
return False, "用户已在群组中"
cursor.execute(
"INSERT INTO group_members (group_id, user_id) VALUES (?,?)",
(group_id, user['id']))
conn.commit()
return True, f"已邀请 {username} 加入群组"
except Exception as e:
return False, str(e)
finally:
conn.close()
def get_all_user_history(self, user_id, limit=200):
"""获取用户的所有聊天记录(私聊+群聊),用于全局历史浏览"""
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute('''
SELECT m.id, m.content, m.created_at, m.msg_type,
u.username as sender_name, u.nickname as sender_nick,
CASE WHEN m.group_id IS NOT NULL THEN g.name ELSE ru.username END as target_name,
CASE WHEN m.group_id IS NOT NULL THEN 'group' ELSE 'private' END as chat_type,
m.group_id, m.receiver_id
FROM messages m
JOIN users u ON m.sender_id = u.id
LEFT JOIN groups g ON m.group_id = g.id
LEFT JOIN users ru ON m.receiver_id = ru.id
WHERE (m.sender_id = ? OR m.receiver_id = ?
OR m.group_id IN (SELECT group_id FROM group_members WHERE user_id = ?))
ORDER BY m.created_at DESC LIMIT ?
''', (user_id, user_id, user_id, limit))
messages = [dict(r) for r in cursor.fetchall()]
conn.close()
return messages
def get_recent_history(self, user_id, chat_type='private', target_id=None, limit=50):
"""获取与特定对象或群组的聊天记录"""
conn = self.get_connection()
cursor = conn.cursor()
if chat_type == 'private':
friend = self.get_user_by_id(target_id) if target_id else None
if not friend:
return []
cursor.execute('''
SELECT m.id, m.content, m.created_at, u.username, u.nickname
FROM messages m JOIN users u ON m.sender_id = u.id
WHERE m.group_id IS NULL
AND ((m.sender_id=? AND m.receiver_id=?) OR (m.sender_id=? AND m.receiver_id=?))
ORDER BY m.created_at DESC LIMIT ?
''', (user_id, friend['id'], friend['id'], user_id, limit))
else:
cursor.execute('''
SELECT m.id, m.content, m.created_at, u.username, u.nickname
FROM messages m JOIN users u ON m.sender_id = u.id
WHERE m.group_id=?
ORDER BY m.created_at DESC LIMIT ?
''', (target_id, limit))
messages = [dict(r) for r in cursor.fetchall()]
conn.close()
return messages
def search_messages(self, user_id, keyword, chat_type='private', target_id=None, limit=100):
conn = self.get_connection()
cursor = conn.cursor()
like = f"%{keyword}%"
if chat_type == 'private':
# target_id may be a username string (from client) or user_id int
friend = None
if isinstance(target_id, str) and not target_id.isdigit():
friend = self.get_user_by_username(target_id)
elif target_id is not None:
friend = self.get_user_by_id(int(target_id))
if not friend:
return []
friend_id = friend['id']
cursor.execute('''
SELECT m.id, m.content, m.created_at, m.msg_type, u.username, u.nickname
FROM messages m JOIN users u ON m.sender_id = u.id
WHERE m.group_id IS NULL
AND ((m.sender_id=? AND m.receiver_id=?) OR (m.sender_id=? AND m.receiver_id=?))
AND m.content LIKE ?
ORDER BY m.created_at DESC LIMIT ?
''', (user_id, friend_id, friend_id, user_id, like, limit))
else:
cursor.execute('''
SELECT m.id, m.content, m.created_at, m.msg_type, u.username, u.nickname
FROM messages m JOIN users u ON m.sender_id = u.id
WHERE m.group_id=? AND m.content LIKE ?
ORDER BY m.created_at DESC LIMIT ?
''', (target_id, like, limit))
messages = [dict(r) for r in cursor.fetchall()]
conn.close()
return messages