package store import ( "database/sql" "fmt" "log" "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"` Content string `json:"content"` CreatedAt time.Time `json:"created_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) } log.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, 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)`, } 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, content string) error { _, err := s.db.Exec( `INSERT INTO messages (session_id, role, content) VALUES ($1, $2, $3)`, sessionID, role, content, ) 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, content, 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.Content, &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"` 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, 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.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 }