package store import ( "database/sql" "fmt" "git.yeij.top/AskaEth/Cyrene/pkg/logger" "time" _ "github.com/lib/pq" ) // Session 会话模型 type Session struct { ID string `json:"id"` UserID string `json:"user_id"` Title string `json:"title"` IsMain bool `json:"is_main"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` } // Message 消息模型 type Message struct { ID int `json:"id"` SessionID string `json:"session_id"` Role string `json:"role"` MsgType string `json:"msg_type"` Content string `json:"content"` ClientID string `json:"client_id"` CreatedAt time.Time `json:"created_at"` } // ClientRecord 客户端记录 (持久化) type ClientRecord struct { ClientID string `json:"client_id"` UserID string `json:"user_id"` DeviceName string `json:"device_name"` UserAgent string `json:"user_agent"` Note string `json:"note"` FirstSeenAt time.Time `json:"first_seen_at"` LastSeenAt time.Time `json:"last_seen_at"` } // SessionStore 会话持久化存储 type SessionStore struct { db *sql.DB } // NewSessionStore 初始化数据库连接并自动建表 // 如果连接失败,返回 nil 和错误(调用方可以选择降级为仅内存模式) func NewSessionStore(databaseURL string) (*SessionStore, error) { db, err := sql.Open("postgres", databaseURL) if err != nil { return nil, fmt.Errorf("无法打开数据库连接: %w", err) } // 配置连接池 db.SetMaxOpenConns(25) db.SetMaxIdleConns(5) db.SetConnMaxLifetime(5 * time.Minute) // 验证连接 if err := db.Ping(); err != nil { db.Close() return nil, fmt.Errorf("数据库连接验证失败: %w", err) } store := &SessionStore{db: db} // 自动建表 if err := store.migrate(); err != nil { db.Close() return nil, fmt.Errorf("数据库迁移失败: %w", err) } logger.Println("[SessionStore] PostgreSQL 持久化存储已初始化") return store, nil } // migrate 自动创建表结构 func (s *SessionStore) migrate() error { queries := []string{ `CREATE TABLE IF NOT EXISTS sessions ( id VARCHAR(64) PRIMARY KEY, user_id VARCHAR(128) NOT NULL, title VARCHAR(256) DEFAULT '', is_main BOOLEAN DEFAULT FALSE, created_at TIMESTAMP DEFAULT NOW(), updated_at TIMESTAMP DEFAULT NOW() )`, `CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id)`, `CREATE INDEX IF NOT EXISTS idx_sessions_updated_at ON sessions(updated_at DESC)`, `CREATE TABLE IF NOT EXISTS messages ( id SERIAL PRIMARY KEY, session_id VARCHAR(64) REFERENCES sessions(id) ON DELETE CASCADE, role VARCHAR(16) NOT NULL, msg_type VARCHAR(16) DEFAULT 'chat', content TEXT NOT NULL, created_at TIMESTAMP DEFAULT NOW() )`, `CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id)`, `CREATE INDEX IF NOT EXISTS idx_messages_created_at ON messages(session_id, created_at)`, // 为已存在的数据库添加 msg_type 列 (Phase 0.1) `ALTER TABLE messages ADD COLUMN IF NOT EXISTS msg_type VARCHAR(16) DEFAULT 'chat'`, // 为已存在的数据库添加 client_id 列 (Phase 5: 多端客户端追踪) `ALTER TABLE messages ADD COLUMN IF NOT EXISTS client_id VARCHAR(128) DEFAULT ''`, `CREATE TABLE IF NOT EXISTS clients ( client_id VARCHAR(128) PRIMARY KEY, user_id VARCHAR(128) NOT NULL, device_name VARCHAR(256) DEFAULT '', user_agent VARCHAR(512) DEFAULT '', note VARCHAR(256) DEFAULT '', first_seen_at TIMESTAMP DEFAULT NOW(), last_seen_at TIMESTAMP DEFAULT NOW() )`, `CREATE INDEX IF NOT EXISTS idx_clients_user_id ON clients(user_id)`, } for _, q := range queries { if _, err := s.db.Exec(q); err != nil { return fmt.Errorf("迁移SQL执行失败: %w\nSQL: %s", err, q) } } return nil } // CreateSession 创建新会话 func (s *SessionStore) CreateSession(userID, sessionID, title string, isMain bool) error { now := time.Now() _, err := s.db.Exec( `INSERT INTO sessions (id, user_id, title, is_main, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $5) ON CONFLICT (id) DO UPDATE SET updated_at = $5`, sessionID, userID, title, isMain, now, ) if err != nil { return fmt.Errorf("创建会话失败: %w", err) } return nil } // GetUserSessions 获取用户的所有会话(按 updated_at DESC 排序) func (s *SessionStore) GetUserSessions(userID string) ([]Session, error) { rows, err := s.db.Query( `SELECT id, user_id, title, is_main, created_at, updated_at FROM sessions WHERE user_id = $1 ORDER BY updated_at DESC`, userID, ) if err != nil { return nil, fmt.Errorf("查询用户会话失败: %w", err) } defer rows.Close() var sessions []Session for rows.Next() { var sess Session if err := rows.Scan(&sess.ID, &sess.UserID, &sess.Title, &sess.IsMain, &sess.CreatedAt, &sess.UpdatedAt); err != nil { return nil, fmt.Errorf("扫描会话行失败: %w", err) } sessions = append(sessions, sess) } if sessions == nil { sessions = []Session{} } return sessions, rows.Err() } // GetSession 获取单个会话 func (s *SessionStore) GetSession(sessionID string) (*Session, error) { var sess Session err := s.db.QueryRow( `SELECT id, user_id, title, is_main, created_at, updated_at FROM sessions WHERE id = $1`, sessionID, ).Scan(&sess.ID, &sess.UserID, &sess.Title, &sess.IsMain, &sess.CreatedAt, &sess.UpdatedAt) if err != nil { if err == sql.ErrNoRows { return nil, nil } return nil, fmt.Errorf("查询会话失败: %w", err) } return &sess, nil } // UpdateSessionTitle 更新会话标题 func (s *SessionStore) UpdateSessionTitle(sessionID, title string) error { _, err := s.db.Exec( `UPDATE sessions SET title = $1, updated_at = NOW() WHERE id = $2`, title, sessionID, ) if err != nil { return fmt.Errorf("更新会话标题失败: %w", err) } return nil } // UpdateSessionTime 更新会话的 updated_at 时间戳 func (s *SessionStore) UpdateSessionTime(sessionID string) error { _, err := s.db.Exec( `UPDATE sessions SET updated_at = NOW() WHERE id = $1`, sessionID, ) if err != nil { return fmt.Errorf("更新会话时间失败: %w", err) } return nil } // DeleteSession 删除会话(级联删除消息,但不删除记忆) func (s *SessionStore) DeleteSession(sessionID string) error { _, err := s.db.Exec(`DELETE FROM sessions WHERE id = $1`, sessionID) if err != nil { return fmt.Errorf("删除会话失败: %w", err) } return nil } // DeleteAllUserSessions 删除用户的所有会话(但不删除记忆) func (s *SessionStore) DeleteAllUserSessions(userID string) error { _, err := s.db.Exec(`DELETE FROM sessions WHERE user_id = $1`, userID) if err != nil { return fmt.Errorf("删除用户所有会话失败: %w", err) } return nil } // AddMessage 添加一条消息到会话 func (s *SessionStore) AddMessage(sessionID, role, msgType, content, clientID string) error { _, err := s.db.Exec( `INSERT INTO messages (session_id, role, msg_type, content, client_id) VALUES ($1, $2, $3, $4, $5)`, sessionID, role, msgType, content, clientID, ) if err != nil { return fmt.Errorf("添加消息失败: %w", err) } return nil } // GetMessages 获取会话的消息列表(按时间正序,支持分页) func (s *SessionStore) GetMessages(sessionID string, limit, offset int) ([]Message, error) { if limit <= 0 { limit = 50 } if offset < 0 { offset = 0 } rows, err := s.db.Query( `SELECT id, session_id, role, COALESCE(msg_type, 'chat'), content, COALESCE(client_id, ''), created_at FROM messages WHERE session_id = $1 ORDER BY created_at ASC LIMIT $2 OFFSET $3`, sessionID, limit, offset, ) if err != nil { return nil, fmt.Errorf("查询消息失败: %w", err) } defer rows.Close() var messages []Message for rows.Next() { var msg Message if err := rows.Scan(&msg.ID, &msg.SessionID, &msg.Role, &msg.MsgType, &msg.Content, &msg.ClientID, &msg.CreatedAt); err != nil { return nil, fmt.Errorf("扫描消息行失败: %w", err) } messages = append(messages, msg) } if messages == nil { messages = []Message{} } return messages, rows.Err() } // ClearSessionMessages 清空会话的所有消息但不删除会话本身 func (s *SessionStore) ClearSessionMessages(sessionID string) error { _, err := s.db.Exec(`DELETE FROM messages WHERE session_id = $1`, sessionID) if err != nil { return fmt.Errorf("清空会话消息失败: %w", err) } return nil } // SearchResult 搜索结果 type SearchResult struct { MessageID int `json:"message_id"` SessionID string `json:"session_id"` SessionTitle string `json:"session_title"` Role string `json:"role"` MsgType string `json:"msg_type"` Content string `json:"content"` CreatedAt time.Time `json:"created_at"` } // SearchMessages 全文搜索消息 (使用 ILIKE 进行模糊匹配) // 返回搜索结果列表、总数和可能的错误 func (s *SessionStore) SearchMessages(userID, query string, limit, offset int) ([]SearchResult, int, error) { if limit <= 0 { limit = 50 } if offset < 0 { offset = 0 } // 获取匹配总数 var total int countSQL := `SELECT COUNT(*) FROM messages m JOIN sessions s ON m.session_id = s.id WHERE s.user_id = $1 AND m.content ILIKE '%' || $2 || '%'` if err := s.db.QueryRow(countSQL, userID, query).Scan(&total); err != nil { return nil, 0, fmt.Errorf("搜索计数失败: %w", err) } // 分页查询,关联 sessions 获取会话标题 rows, err := s.db.Query( `SELECT m.id, m.session_id, COALESCE(s.title, '') AS session_title, m.role, COALESCE(m.msg_type, 'chat'), m.content, m.created_at FROM messages m JOIN sessions s ON m.session_id = s.id WHERE s.user_id = $1 AND m.content ILIKE '%' || $2 || '%' ORDER BY m.created_at DESC LIMIT $3 OFFSET $4`, userID, query, limit, offset, ) if err != nil { return nil, 0, fmt.Errorf("搜索消息失败: %w", err) } defer rows.Close() var results []SearchResult for rows.Next() { var r SearchResult if err := rows.Scan(&r.MessageID, &r.SessionID, &r.SessionTitle, &r.Role, &r.MsgType, &r.Content, &r.CreatedAt); err != nil { return nil, 0, fmt.Errorf("扫描搜索结果行失败: %w", err) } results = append(results, r) } if results == nil { results = []SearchResult{} } return results, total, rows.Err() } // DB 返回底层数据库连接,供其他 store 复用 func (s *SessionStore) DB() *sql.DB { return s.db } // Close 关闭数据库连接 func (s *SessionStore) Close() error { if s.db != nil { return s.db.Close() } return nil } // IsAvailable 检查存储是否可用(数据库连接正常) func (s *SessionStore) IsAvailable() bool { if s.db == nil { return false } return s.db.Ping() == nil } // ========== 多端客户端追踪 ========== // UpsertClient inserts or updates a client record. func (s *SessionStore) UpsertClient(clientID, userID, deviceName, userAgent string) error { _, err := s.db.Exec( `INSERT INTO clients (client_id, user_id, device_name, user_agent, first_seen_at, last_seen_at) VALUES ($1, $2, $3, $4, NOW(), NOW()) ON CONFLICT (client_id) DO UPDATE SET device_name = EXCLUDED.device_name, user_agent = EXCLUDED.user_agent, last_seen_at = NOW()`, clientID, userID, deviceName, userAgent, ) if err != nil { return fmt.Errorf("upsert client failed: %w", err) } return nil } // GetClients returns all known clients for a user. func (s *SessionStore) GetClients(userID string) ([]ClientRecord, error) { rows, err := s.db.Query( `SELECT client_id, user_id, device_name, user_agent, note, first_seen_at, last_seen_at FROM clients WHERE user_id = $1 ORDER BY last_seen_at DESC`, userID, ) if err != nil { return nil, fmt.Errorf("query clients failed: %w", err) } defer rows.Close() var result []ClientRecord for rows.Next() { var cr ClientRecord if err := rows.Scan(&cr.ClientID, &cr.UserID, &cr.DeviceName, &cr.UserAgent, &cr.Note, &cr.FirstSeenAt, &cr.LastSeenAt); err != nil { return nil, fmt.Errorf("scan client row failed: %w", err) } result = append(result, cr) } if result == nil { result = []ClientRecord{} } return result, rows.Err() } // UpdateClientNote sets the user-defined note for a client. func (s *SessionStore) UpdateClientNote(clientID, note string) error { _, err := s.db.Exec(`UPDATE clients SET note = $1 WHERE client_id = $2`, note, clientID) if err != nil { return fmt.Errorf("update client note failed: %w", err) } return nil }