133 lines
4.2 KiB
Go
133 lines
4.2 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/pkg/errors"
|
|
"gorm.io/gorm"
|
|
|
|
"github.com/tamakyi/TamaBox/internal/dbutil"
|
|
)
|
|
|
|
var Messages MessageStore
|
|
|
|
type MessageStore interface {
|
|
Create(ctx context.Context, opts CreateMessageOptions) (*Message, error)
|
|
GetByConversationID(ctx context.Context, conversationID uint, cursor *dbutil.Cursor) ([]*Message, error)
|
|
// GetFirstMessagesByConversationIDs 批量获取多个会话的第一条消息（解决 N+1 查询）
|
|
GetFirstMessagesByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint]string, error)
|
|
// GetAllByConversationIDs 批量获取多个会话的全部消息（导出用，不分页），按 conversation_id 分组。
|
|
// 将每个会话原本的 1 次查询合并为 1 次 IN 查询，避免 N+1。
|
|
GetAllByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint][]*Message, error)
|
|
}
|
|
|
|
func NewMessageStore(db *gorm.DB) MessageStore {
|
|
return &messages{db}
|
|
}
|
|
|
|
type messages struct {
|
|
*gorm.DB
|
|
}
|
|
|
|
type Message struct {
|
|
dbutil.Model
|
|
ConversationID uint `gorm:"index:idx_message_conversation" json:"conversation_id"`
|
|
SenderID uint `json:"sender_id"`
|
|
Content string `gorm:"type:text" json:"content"`
|
|
Read bool `gorm:"default:false" json:"read"`
|
|
}
|
|
|
|
type CreateMessageOptions struct {
|
|
ConversationID uint
|
|
SenderID uint
|
|
Content string
|
|
}
|
|
|
|
func (db *messages) Create(ctx context.Context, opts CreateMessageOptions) (*Message, error) {
|
|
msg := &Message{
|
|
ConversationID: opts.ConversationID,
|
|
SenderID: opts.SenderID,
|
|
Content: opts.Content,
|
|
}
|
|
if err := db.WithContext(ctx).Create(msg).Error; err != nil {
|
|
return nil, errors.Wrap(err, "create message")
|
|
}
|
|
return msg, nil
|
|
}
|
|
|
|
func (db *messages) GetByConversationID(ctx context.Context, conversationID uint, cursor *dbutil.Cursor) ([]*Message, error) {
|
|
var messages []*Message
|
|
q := db.WithContext(ctx).Where("conversation_id = ?", conversationID)
|
|
|
|
limit := 50
|
|
if cursor != nil {
|
|
if cursor.Value != nil && fmt.Sprintf("%v", cursor.Value) != "" {
|
|
q = q.Where("id < ?", cursor.Value)
|
|
}
|
|
limit = cursor.Limit()
|
|
}
|
|
|
|
q = q.Order("created_at ASC")
|
|
q = q.Limit(limit)
|
|
|
|
if err := q.Find(&messages).Error; err != nil {
|
|
return nil, errors.Wrap(err, "get messages")
|
|
}
|
|
return messages, nil
|
|
}
|
|
|
|
// GetFirstMessagesByConversationIDs 批量获取多个会话的第一条消息内容。
|
|
// 使用窗口函数 ROW_NUMBER() 选出每个会话最早的消息，避免 N+1 查询。
|
|
func (db *messages) GetFirstMessagesByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint]string, error) {
|
|
if len(conversationIDs) == 0 {
|
|
return make(map[uint]string), nil
|
|
}
|
|
|
|
// 使用子查询 + ROW_NUMBER 窗口函数获取每个 conversation 的第一条消息
|
|
// 兼容 MySQL 8+ / PostgreSQL / SQLite 3.25+
|
|
var results []struct {
|
|
ConversationID uint `gorm:"column:conversation_id"`
|
|
Content string `gorm:"column:content"`
|
|
}
|
|
|
|
err := db.WithContext(ctx).Raw(`
|
|
SELECT conversation_id, content FROM (
|
|
SELECT conversation_id, content,
|
|
ROW_NUMBER() OVER (PARTITION BY conversation_id ORDER BY created_at ASC) AS rn
|
|
FROM messages
|
|
WHERE conversation_id IN ?
|
|
) sub WHERE rn = 1
|
|
`, conversationIDs).Scan(&results).Error
|
|
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "get first messages batch")
|
|
}
|
|
|
|
resultMap := make(map[uint]string, len(results))
|
|
for _, r := range results {
|
|
resultMap[r.ConversationID] = r.Content
|
|
}
|
|
return resultMap, nil
|
|
}
|
|
|
|
// GetAllByConversationIDs 批量获取多个会话的全部消息（导出用，不分页），按 conversation_id 分组。
|
|
// 用于账号导出：将每个会话原本的 1 次查询合并为 1 次 IN 查询，避免 N+1。
|
|
func (db *messages) GetAllByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint][]*Message, error) {
|
|
result := make(map[uint][]*Message)
|
|
if len(conversationIDs) == 0 {
|
|
return result, nil
|
|
}
|
|
var items []*Message
|
|
if err := db.WithContext(ctx).
|
|
Where("conversation_id IN ?", conversationIDs).
|
|
Order("conversation_id ASC, created_at ASC").
|
|
Find(&items).Error; err != nil {
|
|
return nil, errors.Wrap(err, "get messages by conversation ids")
|
|
}
|
|
for _, m := range items {
|
|
result[m.ConversationID] = append(result[m.ConversationID], m)
|
|
}
|
|
return result, nil
|
|
}
|