1287 lines
42 KiB
Go
1287 lines
42 KiB
Go
package sqlite
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jmoiron/sqlx"
|
|
"github.com/kandev/kandev/internal/db/dialect"
|
|
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
|
|
"github.com/kandev/kandev/internal/task/models"
|
|
)
|
|
|
|
const attachmentSelectColumns = `id, owner_id, workspace_id, task_id, session_id,
|
|
message_id, queue_id, name, mime_type, kind, delivery_mode, size_bytes,
|
|
storage_key, state, expires_at, created_at, updated_at`
|
|
|
|
func (r *Repository) CreateMessageAttachment(ctx context.Context, attachment *models.TaskMessageAttachment) error {
|
|
if attachment.ID == "" {
|
|
attachment.ID = uuid.NewString()
|
|
}
|
|
if attachment.StorageKey == "" {
|
|
attachment.StorageKey = attachment.ID
|
|
}
|
|
now := time.Now().UTC()
|
|
if attachment.CreatedAt.IsZero() {
|
|
attachment.CreatedAt = now
|
|
}
|
|
attachment.UpdatedAt = now
|
|
if attachment.State == "" {
|
|
attachment.State = models.AttachmentStateStaged
|
|
}
|
|
_, err := r.db.ExecContext(ctx, r.db.Rebind(`
|
|
INSERT INTO task_message_attachments
|
|
(id, owner_id, workspace_id, task_id, session_id, message_id, queue_id,
|
|
name, mime_type, kind, delivery_mode, size_bytes, storage_key, state,
|
|
expires_at, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`), attachment.ID, attachment.OwnerID, attachment.WorkspaceID, attachment.TaskID,
|
|
attachment.SessionID, attachment.MessageID, attachment.QueueID, attachment.Name,
|
|
attachment.MimeType, attachment.Kind, attachment.DeliveryMode, attachment.SizeBytes,
|
|
attachment.StorageKey, attachment.State, attachment.ExpiresAt, attachment.CreatedAt,
|
|
attachment.UpdatedAt)
|
|
if err != nil {
|
|
return fmt.Errorf("create message attachment: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) GetMessageAttachment(ctx context.Context, id string) (*models.TaskMessageAttachment, error) {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
err := r.ro.QueryRowxContext(ctx, r.ro.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments WHERE id = ?
|
|
`), id).StructScan(attachment)
|
|
if err == sql.ErrNoRows {
|
|
return nil, models.ErrAttachmentNotFound
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("get message attachment: %w", err)
|
|
}
|
|
return attachment, nil
|
|
}
|
|
|
|
func (r *Repository) ListMessageAttachments(ctx context.Context, ids []string) ([]*models.TaskMessageAttachment, error) {
|
|
if len(ids) == 0 {
|
|
return nil, nil
|
|
}
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(ids)), ",")
|
|
args := make([]interface{}, len(ids))
|
|
for i := range ids {
|
|
args[i] = ids[i]
|
|
}
|
|
rows, err := r.ro.QueryxContext(ctx, r.ro.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments
|
|
WHERE id IN (`+placeholders+`)
|
|
`), args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list message attachments: %w", err)
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
var out []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
return nil, fmt.Errorf("scan message attachment: %w", err)
|
|
}
|
|
out = append(out, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("iterate message attachments: %w", err)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (r *Repository) ClaimMessageAttachments(ctx context.Context, ids []string, ownerID, workspaceID, taskID, sessionID string) error {
|
|
if len(ids) == 0 {
|
|
return nil
|
|
}
|
|
if len(ids) > models.MaxMessageAttachmentCount {
|
|
return models.ErrTooManyAttachments
|
|
}
|
|
queueLockPresent, err := r.queueSessionLockTablePresent(ctx)
|
|
if err != nil {
|
|
return fmt.Errorf("check queue session lock schema: %w", err)
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("begin attachment claim: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if queueLockPresent {
|
|
if err := messagequeue.LockSessionInTransaction(ctx, tx, r.db, sessionID); err != nil {
|
|
return fmt.Errorf("lock attachment claim session: %w", err)
|
|
}
|
|
}
|
|
if taskID != "" {
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE tasks SET updated_at = updated_at WHERE id = ?
|
|
`), taskID)
|
|
if err != nil {
|
|
return fmt.Errorf("lock attachment task %q: %w", taskID, err)
|
|
}
|
|
rows, _ := result.RowsAffected()
|
|
if rows == 0 {
|
|
return fmt.Errorf("%w: %s", ErrTaskNotFound, taskID)
|
|
}
|
|
}
|
|
if sessionID != "" {
|
|
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), sessionID); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
now := time.Now().UTC()
|
|
selection, err := r.selectAttachmentsForClaim(ctx, tx, ids, ownerID, workspaceID, taskID, sessionID, now)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if selection.selectedSize > models.MaxMessageAttachmentBytes {
|
|
return models.ErrAttachmentTotalTooLarge
|
|
}
|
|
if err := r.markAttachmentsClaimed(ctx, tx, selection.claimIDs, ownerID, workspaceID, taskID, sessionID); err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("commit attachment claim: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ClaimQueuedMessageAttachments marks only newly staged rows with queueID.
|
|
// Existing claims for the same task/session remain usable but are not adopted,
|
|
// so a failed admission cannot roll back another accepted message's claim.
|
|
func (r *Repository) ClaimQueuedMessageAttachments(
|
|
ctx context.Context,
|
|
ids []string,
|
|
ownerID, workspaceID, taskID, sessionID, queueID string,
|
|
) error {
|
|
if len(ids) == 0 {
|
|
return nil
|
|
}
|
|
if queueID == "" {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
if len(ids) > models.MaxMessageAttachmentCount {
|
|
return models.ErrTooManyAttachments
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("begin queued attachment claim: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
|
|
now := time.Now().UTC()
|
|
selection, err := r.selectQueuedAttachmentsForClaim(
|
|
ctx, tx, ids, ownerID, workspaceID, taskID, sessionID, queueID, now,
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if selection.selectedSize > models.MaxMessageAttachmentBytes {
|
|
return models.ErrAttachmentTotalTooLarge
|
|
}
|
|
if err := r.markQueuedAttachmentsClaimed(
|
|
ctx, tx, selection.claimIDs, ownerID, workspaceID, taskID, sessionID, queueID, now,
|
|
); err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("commit queued attachment claim: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) selectQueuedAttachmentsForClaim(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
ids []string,
|
|
ownerID, workspaceID, taskID, sessionID, queueID string,
|
|
now time.Time,
|
|
) (attachmentClaimSelection, error) {
|
|
selection := attachmentClaimSelection{}
|
|
seen := make(map[string]struct{}, len(ids))
|
|
for _, id := range ids {
|
|
if !recordUniqueAttachmentID(seen, id) {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentClaimConflict
|
|
}
|
|
attachment, err := loadAttachmentForClaim(ctx, tx, id)
|
|
if err != nil {
|
|
return attachmentClaimSelection{}, err
|
|
}
|
|
if attachment.OwnerID != ownerID || attachment.WorkspaceID != workspaceID {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentClaimConflict
|
|
}
|
|
if err := addQueuedAttachmentToClaim(&selection, attachment, id, taskID, sessionID, queueID, now); err != nil {
|
|
return attachmentClaimSelection{}, err
|
|
}
|
|
}
|
|
return selection, nil
|
|
}
|
|
|
|
func addQueuedAttachmentToClaim(
|
|
selection *attachmentClaimSelection,
|
|
attachment *models.TaskMessageAttachment,
|
|
id, taskID, sessionID, queueID string,
|
|
now time.Time,
|
|
) error {
|
|
if attachment.State == models.AttachmentStateStaged {
|
|
return addStagedAttachmentToClaim(selection, attachment, id, now)
|
|
}
|
|
if attachment.State != models.AttachmentStateClaimed ||
|
|
attachment.TaskID != taskID ||
|
|
(attachment.SessionID != "" && attachment.SessionID != sessionID) ||
|
|
attachment.MessageID != "" ||
|
|
(attachment.QueueID != "" && attachment.QueueID != queueID) {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) markQueuedAttachmentsClaimed(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
ids []string,
|
|
ownerID, workspaceID, taskID, sessionID, queueID string,
|
|
now time.Time,
|
|
) error {
|
|
for _, id := range ids {
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments
|
|
SET task_id = ?, session_id = ?, queue_id = ?, state = ?, updated_at = ?
|
|
WHERE id = ? AND owner_id = ? AND workspace_id = ? AND state = ?
|
|
`), taskID, sessionID, queueID, models.AttachmentStateClaimed, now,
|
|
id, ownerID, workspaceID, models.AttachmentStateStaged)
|
|
if err != nil {
|
|
return fmt.Errorf("claim queued attachment: %w", err)
|
|
}
|
|
if changed, _ := result.RowsAffected(); changed != 1 {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// RestoreQueuedMessageAttachments returns only rows provisionally claimed by
|
|
// this queue admission to staged state. A pre-existing claim has no queueID and
|
|
// is deliberately untouched.
|
|
func (r *Repository) RestoreQueuedMessageAttachments(
|
|
ctx context.Context,
|
|
ids []string,
|
|
ownerID, taskID, sessionID, queueID string,
|
|
) error {
|
|
if len(ids) == 0 {
|
|
return nil
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("begin queued attachment restore: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
now := time.Now().UTC()
|
|
for _, id := range ids {
|
|
if _, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments
|
|
SET task_id = '', session_id = '', message_id = '', queue_id = '', state = ?, updated_at = ?
|
|
WHERE id = ? AND owner_id = ? AND task_id = ? AND session_id = ?
|
|
AND queue_id = ? AND state = ?
|
|
`), models.AttachmentStateStaged, now, id, ownerID, taskID, sessionID,
|
|
queueID, models.AttachmentStateClaimed); err != nil {
|
|
return fmt.Errorf("restore queued attachment: %w", err)
|
|
}
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("commit queued attachment restore: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) claimDirectMessageAttachmentsTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
claim *messagequeue.QueueAttachmentClaim,
|
|
taskID, sessionID, messageID string,
|
|
) error {
|
|
if claim == nil || len(claim.IDs) == 0 {
|
|
return nil
|
|
}
|
|
if claim.OwnerID == "" || claim.WorkspaceID == "" || messageID == "" || len(claim.IDs) > models.MaxMessageAttachmentCount {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
now := time.Now().UTC()
|
|
selection, err := r.selectDirectAttachmentsForClaim(
|
|
ctx, tx, claim.IDs, claim.OwnerID, claim.WorkspaceID, taskID, sessionID, messageID, now,
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if selection.selectedSize > models.MaxMessageAttachmentBytes {
|
|
return models.ErrAttachmentTotalTooLarge
|
|
}
|
|
if err := r.markDirectAttachmentsClaimed(
|
|
ctx, tx, selection.claimIDs, claim.OwnerID, claim.WorkspaceID, taskID, sessionID, messageID, now,
|
|
); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) selectDirectAttachmentsForClaim(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
ids []string,
|
|
ownerID, workspaceID, taskID, sessionID, messageID string,
|
|
now time.Time,
|
|
) (attachmentClaimSelection, error) {
|
|
selection := attachmentClaimSelection{}
|
|
seen := make(map[string]struct{}, len(ids))
|
|
for _, id := range ids {
|
|
if !recordUniqueAttachmentID(seen, id) {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentClaimConflict
|
|
}
|
|
attachment, err := loadAttachmentForClaim(ctx, tx, id)
|
|
if err != nil {
|
|
return attachmentClaimSelection{}, err
|
|
}
|
|
if attachment.OwnerID != ownerID || attachment.WorkspaceID != workspaceID {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentClaimConflict
|
|
}
|
|
if err := addDirectAttachmentToClaim(&selection, attachment, id, taskID, sessionID, messageID, now); err != nil {
|
|
return attachmentClaimSelection{}, err
|
|
}
|
|
}
|
|
return selection, nil
|
|
}
|
|
|
|
func addDirectAttachmentToClaim(
|
|
selection *attachmentClaimSelection,
|
|
attachment *models.TaskMessageAttachment,
|
|
id, taskID, sessionID, messageID string,
|
|
now time.Time,
|
|
) error {
|
|
if attachment.State == models.AttachmentStateStaged {
|
|
return addStagedAttachmentToClaim(selection, attachment, id, now)
|
|
}
|
|
if attachment.State != models.AttachmentStateClaimed ||
|
|
attachment.TaskID != taskID ||
|
|
(attachment.SessionID != "" && attachment.SessionID != sessionID) ||
|
|
attachment.QueueID != "" ||
|
|
(attachment.MessageID != "" && attachment.MessageID != messageID) {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) markDirectAttachmentsClaimed(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
ids []string,
|
|
ownerID, workspaceID, taskID, sessionID, messageID string,
|
|
now time.Time,
|
|
) error {
|
|
for _, id := range ids {
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments
|
|
SET task_id = ?, session_id = ?, message_id = ?, state = ?, updated_at = ?
|
|
WHERE id = ? AND owner_id = ? AND workspace_id = ? AND state = ?
|
|
`), taskID, sessionID, messageID, models.AttachmentStateClaimed, now,
|
|
id, ownerID, workspaceID, models.AttachmentStateStaged)
|
|
if err != nil {
|
|
return fmt.Errorf("claim direct attachment: %w", err)
|
|
}
|
|
if changed, _ := result.RowsAffected(); changed != 1 {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type attachmentClaimSelection struct {
|
|
claimIDs []string
|
|
selectedSize int64
|
|
}
|
|
|
|
func recordUniqueAttachmentID(seen map[string]struct{}, id string) bool {
|
|
if _, exists := seen[id]; exists {
|
|
return false
|
|
}
|
|
seen[id] = struct{}{}
|
|
return true
|
|
}
|
|
|
|
func loadAttachmentForClaim(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
id string,
|
|
) (*models.TaskMessageAttachment, error) {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
err := tx.GetContext(ctx, attachment, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments WHERE id = ?
|
|
`), id)
|
|
if errorsIsNoRows(err) {
|
|
return nil, models.ErrAttachmentNotFound
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("load attachment for claim: %w", err)
|
|
}
|
|
return attachment, nil
|
|
}
|
|
|
|
func addStagedAttachmentToClaim(
|
|
selection *attachmentClaimSelection,
|
|
attachment *models.TaskMessageAttachment,
|
|
id string,
|
|
now time.Time,
|
|
) error {
|
|
if !attachment.ExpiresAt.IsZero() && !attachment.ExpiresAt.After(now) {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
if attachment.SizeBytes < 0 || attachment.SizeBytes > models.MaxMessageAttachmentBytes {
|
|
return models.ErrAttachmentTooLarge
|
|
}
|
|
selection.selectedSize += attachment.SizeBytes
|
|
selection.claimIDs = append(selection.claimIDs, id)
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) selectAttachmentsForClaim(ctx context.Context, tx *sqlx.Tx, ids []string, ownerID, workspaceID, taskID, sessionID string, now time.Time) (attachmentClaimSelection, error) {
|
|
selection := attachmentClaimSelection{}
|
|
seen := make(map[string]struct{}, len(ids))
|
|
for _, id := range ids {
|
|
if _, ok := seen[id]; ok {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentClaimConflict
|
|
}
|
|
seen[id] = struct{}{}
|
|
var attachment models.TaskMessageAttachment
|
|
err := tx.GetContext(ctx, &attachment, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments WHERE id = ?
|
|
`), id)
|
|
if errorsIsNoRows(err) {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentNotFound
|
|
}
|
|
if err != nil {
|
|
return attachmentClaimSelection{}, fmt.Errorf("load attachment for claim: %w", err)
|
|
}
|
|
if attachment.OwnerID != ownerID || attachment.WorkspaceID != workspaceID {
|
|
return attachmentClaimSelection{}, models.ErrAttachmentClaimConflict
|
|
}
|
|
if err := addAttachmentToClaim(&selection, &attachment, id, taskID, sessionID, now); err != nil {
|
|
return attachmentClaimSelection{}, err
|
|
}
|
|
}
|
|
return selection, nil
|
|
}
|
|
|
|
func addAttachmentToClaim(selection *attachmentClaimSelection, attachment *models.TaskMessageAttachment, id, taskID, sessionID string, now time.Time) error {
|
|
switch attachment.State {
|
|
case models.AttachmentStateStaged:
|
|
if !attachment.ExpiresAt.IsZero() && !attachment.ExpiresAt.After(now) {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
if attachment.SizeBytes < 0 || attachment.SizeBytes > models.MaxMessageAttachmentBytes {
|
|
return models.ErrAttachmentTooLarge
|
|
}
|
|
selection.selectedSize += attachment.SizeBytes
|
|
selection.claimIDs = append(selection.claimIDs, id)
|
|
return nil
|
|
case models.AttachmentStateClaimed:
|
|
if attachment.TaskID != taskID ||
|
|
(attachment.SessionID != "" && attachment.SessionID != sessionID) {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
return nil
|
|
default:
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
}
|
|
|
|
func (r *Repository) markAttachmentsClaimed(ctx context.Context, tx *sqlx.Tx, ids []string, ownerID, workspaceID, taskID, sessionID string) error {
|
|
now := time.Now().UTC()
|
|
for _, id := range ids {
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments
|
|
SET task_id = ?, session_id = ?, state = ?, updated_at = ?
|
|
WHERE id = ? AND owner_id = ? AND workspace_id = ? AND state = ?
|
|
`), taskID, sessionID, models.AttachmentStateClaimed, now, id, ownerID, workspaceID, models.AttachmentStateStaged)
|
|
if err != nil {
|
|
return fmt.Errorf("claim attachment: %w", err)
|
|
}
|
|
affected, err := result.RowsAffected()
|
|
if err != nil {
|
|
return fmt.Errorf("claim attachment rows affected: %w", err)
|
|
}
|
|
if affected != 1 {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) DeleteMessageAttachment(ctx context.Context, id, ownerID string) error {
|
|
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
|
|
DELETE FROM task_message_attachments WHERE id = ? AND owner_id = ?
|
|
`), id, ownerID)
|
|
if err != nil {
|
|
return fmt.Errorf("delete message attachment: %w", err)
|
|
}
|
|
rows, _ := result.RowsAffected()
|
|
if rows == 0 {
|
|
return models.ErrAttachmentNotFound
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) DeleteClaimedMessageAttachments(
|
|
ctx context.Context, ids []string, ownerID, taskID, sessionID string,
|
|
) ([]*models.TaskMessageAttachment, error) {
|
|
return r.deleteClaimedMessageAttachments(ctx, ids, ownerID, taskID, sessionID, true)
|
|
}
|
|
|
|
func (r *Repository) DeleteClaimedMessageAttachmentsByTaskSession(
|
|
ctx context.Context, ids []string, taskID, sessionID string,
|
|
) ([]*models.TaskMessageAttachment, error) {
|
|
return r.deleteClaimedMessageAttachments(ctx, ids, "", taskID, sessionID, false)
|
|
}
|
|
|
|
//nolint:cyclop,funlen // Attachment release coordinates ownership, queue locks, and transactional cleanup.
|
|
func (r *Repository) deleteClaimedMessageAttachments(
|
|
ctx context.Context,
|
|
ids []string,
|
|
ownerID, taskID, sessionID string,
|
|
ownerRequired bool,
|
|
) ([]*models.TaskMessageAttachment, error) {
|
|
if len(ids) == 0 {
|
|
return nil, nil
|
|
}
|
|
queueTablePresent, err := r.tableExistsContext(ctx, "queued_messages")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("probe queue attachment references: %w", err)
|
|
}
|
|
queueLockPresent := false
|
|
if sessionID != "" {
|
|
queueLockPresent, err = r.queueSessionLockTablePresent(ctx)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("check queue session lock schema: %w", err)
|
|
}
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("begin claimed attachment release: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if err := r.lockClaimedAttachmentReleaseTx(ctx, tx, sessionID, queueLockPresent); err != nil {
|
|
return nil, err
|
|
}
|
|
if sessionID != "" {
|
|
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), sessionID); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if queueTablePresent {
|
|
if err := lockAttachmentQueueSession(ctx, tx, r.db.DriverName(), sessionID); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
referenced, err := referencedAttachmentIDsTx(ctx, r, tx, taskID, sessionID, queueTablePresent)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
releasable := make([]string, 0, len(ids))
|
|
seen := make(map[string]struct{}, len(ids))
|
|
for _, id := range ids {
|
|
if id == "" {
|
|
continue
|
|
}
|
|
if _, ok := referenced[id]; ok {
|
|
continue
|
|
}
|
|
if _, ok := seen[id]; ok {
|
|
continue
|
|
}
|
|
seen[id] = struct{}{}
|
|
releasable = append(releasable, id)
|
|
}
|
|
if len(releasable) == 0 {
|
|
return nil, nil
|
|
}
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(releasable)), ",")
|
|
ownerClause := ""
|
|
args := []interface{}{taskID, sessionID, models.AttachmentStateClaimed}
|
|
if ownerRequired {
|
|
ownerClause = "owner_id = ? AND "
|
|
args = append([]interface{}{ownerID}, args...)
|
|
}
|
|
args = append(args, idsToInterfaces(releasable)...)
|
|
released, err := listClaimedAttachmentsForReleaseTx(
|
|
ctx, tx, ownerClause, attachmentSessionScopeClause(sessionID), placeholders, args,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err := deleteClaimedAttachmentsTx(
|
|
ctx, tx, released, ownerClause, attachmentSessionScopeClause(sessionID),
|
|
ownerID, ownerRequired, taskID, sessionID,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return nil, fmt.Errorf("commit claimed attachment release: %w", err)
|
|
}
|
|
return released, nil
|
|
}
|
|
|
|
func (r *Repository) lockClaimedAttachmentReleaseTx(
|
|
ctx context.Context, tx *sqlx.Tx, sessionID string, queueLockPresent bool,
|
|
) error {
|
|
if !queueLockPresent {
|
|
return nil
|
|
}
|
|
if err := messagequeue.LockSessionInTransaction(ctx, tx, r.db, sessionID); err != nil {
|
|
return fmt.Errorf("lock attachment release session %s: %w", sessionID, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func lockAttachmentQueueSession(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
driverName, sessionID string,
|
|
) error {
|
|
if sessionID == "" || !dialect.IsPostgres(driverName) {
|
|
return nil
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `
|
|
INSERT INTO queue_session_locks (session_id) VALUES ($1)
|
|
ON CONFLICT(session_id) DO NOTHING
|
|
`, sessionID); err != nil {
|
|
return fmt.Errorf("ensure attachment queue session lock %q: %w", sessionID, err)
|
|
}
|
|
var lockedSessionID string
|
|
if err := tx.GetContext(
|
|
ctx,
|
|
&lockedSessionID,
|
|
`SELECT session_id FROM queue_session_locks WHERE session_id = $1 FOR UPDATE`,
|
|
sessionID,
|
|
); err != nil {
|
|
return fmt.Errorf("lock attachment queue session %q: %w", sessionID, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func listClaimedAttachmentsForReleaseTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
ownerClause, sessionScopeClause, placeholders string,
|
|
args []interface{},
|
|
) ([]*models.TaskMessageAttachment, error) {
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments
|
|
WHERE `+ownerClause+`task_id = ? AND `+sessionScopeClause+` AND state = ?
|
|
AND id IN (`+placeholders+`)
|
|
`), args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list claimed attachments for release: %w", err)
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
var released []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
return nil, fmt.Errorf("scan claimed attachment for release: %w", err)
|
|
}
|
|
released = append(released, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("iterate claimed attachments for release: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close claimed attachments for release: %w", err)
|
|
}
|
|
return released, nil
|
|
}
|
|
|
|
func deleteClaimedAttachmentsTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
released []*models.TaskMessageAttachment,
|
|
ownerClause, sessionScopeClause, ownerID string,
|
|
ownerRequired bool,
|
|
taskID, sessionID string,
|
|
) error {
|
|
for _, attachment := range released {
|
|
deleteArgs := []interface{}{attachment.ID}
|
|
if ownerRequired {
|
|
deleteArgs = append(deleteArgs, ownerID)
|
|
}
|
|
deleteArgs = append(deleteArgs, taskID, sessionID, models.AttachmentStateClaimed)
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
DELETE FROM task_message_attachments
|
|
WHERE id = ? AND `+ownerClause+`task_id = ? AND `+sessionScopeClause+` AND state = ?
|
|
`), deleteArgs...)
|
|
if err != nil {
|
|
return fmt.Errorf("release claimed attachment: %w", err)
|
|
}
|
|
affected, err := result.RowsAffected()
|
|
if err != nil {
|
|
return fmt.Errorf("release claimed attachment rows affected: %w", err)
|
|
}
|
|
if affected != 1 {
|
|
return models.ErrAttachmentClaimConflict
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func attachmentSessionScopeClause(sessionID string) string {
|
|
if sessionID == "" {
|
|
return "session_id = ?"
|
|
}
|
|
return "(session_id = ? OR session_id = '')"
|
|
}
|
|
|
|
type attachmentReference struct {
|
|
AttachmentID string `json:"attachment_id"`
|
|
}
|
|
|
|
func referencedAttachmentIDsTx(
|
|
ctx context.Context,
|
|
r *Repository,
|
|
tx *sqlx.Tx,
|
|
taskID, sessionID string,
|
|
queueTablePresent bool,
|
|
) (map[string]struct{}, error) {
|
|
referenced := make(map[string]struct{})
|
|
//nolint:nestif // durable queue and transcript references use separate optional tables.
|
|
if queueTablePresent {
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT attachments_json FROM queued_messages WHERE task_id = ? AND session_id = ?
|
|
`), taskID, sessionID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list queued attachment references: %w", err)
|
|
}
|
|
for rows.Next() {
|
|
var encoded string
|
|
if err := rows.Scan(&encoded); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("scan queued attachment references: %w", err)
|
|
}
|
|
var attachments []attachmentReference
|
|
if err := json.Unmarshal([]byte(encoded), &attachments); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("decode queued attachment references: %w", err)
|
|
}
|
|
addAttachmentReferences(referenced, attachments)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("iterate queued attachment references: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close queued attachment references: %w", err)
|
|
}
|
|
}
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT metadata FROM task_session_messages WHERE task_id = ? AND task_session_id = ?
|
|
`), taskID, sessionID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list transcript attachment references: %w", err)
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
for rows.Next() {
|
|
var encoded string
|
|
if err := rows.Scan(&encoded); err != nil {
|
|
return nil, fmt.Errorf("scan transcript attachment references: %w", err)
|
|
}
|
|
var metadata struct {
|
|
Attachments []attachmentReference `json:"attachments"`
|
|
}
|
|
if err := json.Unmarshal([]byte(encoded), &metadata); err != nil {
|
|
return nil, fmt.Errorf("decode transcript attachment references: %w", err)
|
|
}
|
|
addAttachmentReferences(referenced, metadata.Attachments)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("iterate transcript attachment references: %w", err)
|
|
}
|
|
return referenced, nil
|
|
}
|
|
|
|
func (r *Repository) queuedAttachmentIDsForTaskTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
taskID string,
|
|
) (map[string]struct{}, error) {
|
|
ids := make(map[string]struct{})
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT attachments_json FROM queued_messages WHERE task_id = ?
|
|
`), taskID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list queued attachment references for archive: %w", err)
|
|
}
|
|
for rows.Next() {
|
|
var encoded string
|
|
if err := rows.Scan(&encoded); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("scan queued attachment references for archive: %w", err)
|
|
}
|
|
var attachments []attachmentReference
|
|
if err := json.Unmarshal([]byte(encoded), &attachments); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("decode queued attachment references for archive: %w", err)
|
|
}
|
|
addAttachmentReferences(ids, attachments)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("iterate queued attachment references for archive: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close queued attachment references for archive: %w", err)
|
|
}
|
|
return ids, nil
|
|
}
|
|
|
|
// releaseUnreferencedTaskAttachmentClaimsTx stages queue-owned claims that no
|
|
// transcript message references. Only IDs present in queuedIDs are eligible:
|
|
// a direct prompt Claim commits before its transcript row exists, so staging
|
|
// every unreferenced claimed row would release an in-flight direct claim whose
|
|
// message inserts after the archive commits. Queue admission claims atomically
|
|
// with its queue row, so the pre-purge queued set is the exact ownership
|
|
// boundary. Rows are collected before any UPDATE so no cursor stays open
|
|
// across statements on a single-connection Postgres transaction.
|
|
func (r *Repository) releaseUnreferencedTaskAttachmentClaimsTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
taskID string,
|
|
queuedIDs map[string]struct{},
|
|
) error {
|
|
if len(queuedIDs) == 0 {
|
|
return nil
|
|
}
|
|
referenced := make(map[string]struct{})
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT metadata FROM task_session_messages WHERE task_id = ?
|
|
`), taskID)
|
|
if err != nil {
|
|
return fmt.Errorf("list task attachment references for archive: %w", err)
|
|
}
|
|
for rows.Next() {
|
|
var encoded string
|
|
if err := rows.Scan(&encoded); err != nil {
|
|
_ = rows.Close()
|
|
return fmt.Errorf("scan task attachment references for archive: %w", err)
|
|
}
|
|
var metadata struct {
|
|
Attachments []attachmentReference `json:"attachments"`
|
|
}
|
|
if err := json.Unmarshal([]byte(encoded), &metadata); err != nil {
|
|
_ = rows.Close()
|
|
return fmt.Errorf("decode task attachment references for archive: %w", err)
|
|
}
|
|
addAttachmentReferences(referenced, metadata.Attachments)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return fmt.Errorf("iterate task attachment references for archive: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return fmt.Errorf("close task attachment references for archive: %w", err)
|
|
}
|
|
var unreferenced []string
|
|
for id := range queuedIDs {
|
|
if _, ok := referenced[id]; !ok {
|
|
unreferenced = append(unreferenced, id)
|
|
}
|
|
}
|
|
now := time.Now().UTC()
|
|
for _, attachmentID := range unreferenced {
|
|
if _, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments
|
|
SET state = ?, expires_at = ?, updated_at = ?
|
|
WHERE id = ? AND task_id = ? AND state = ?
|
|
`), models.AttachmentStateStaged, now, now, attachmentID, taskID, models.AttachmentStateClaimed); err != nil {
|
|
return fmt.Errorf("release task attachment claim for archive: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) deleteUnreferencedSessionAttachmentClaimsTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
taskID, sessionID string,
|
|
) ([]*models.TaskMessageAttachment, error) {
|
|
// The session, its transcript, and its queue rows are deleted by this
|
|
// transaction. Return descriptors before deleting their registry rows so
|
|
// the caller can remove private bytes after commit.
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+`
|
|
FROM task_message_attachments
|
|
WHERE task_id = ? AND session_id = ? AND state = ?
|
|
`), taskID, sessionID, models.AttachmentStateClaimed)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list deleted-session attachment claims: %w", err)
|
|
}
|
|
var attachments []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("scan deleted-session attachment claim: %w", err)
|
|
}
|
|
attachments = append(attachments, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("iterate deleted-session attachment claims: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close deleted-session attachment claims: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
DELETE FROM task_message_attachments
|
|
WHERE task_id = ? AND session_id = ? AND state = ?
|
|
`), taskID, sessionID, models.AttachmentStateClaimed); err != nil {
|
|
return nil, fmt.Errorf("delete deleted-session attachment claims: %w", err)
|
|
}
|
|
return attachments, nil
|
|
}
|
|
|
|
func addAttachmentReferences(target map[string]struct{}, attachments []attachmentReference) {
|
|
for _, attachment := range attachments {
|
|
if attachment.AttachmentID != "" {
|
|
target[attachment.AttachmentID] = struct{}{}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (r *Repository) DeleteMessageAttachmentsByTask(ctx context.Context, taskID string) ([]*models.TaskMessageAttachment, error) {
|
|
if taskID == "" {
|
|
return nil, nil
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("begin task attachment cleanup: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments WHERE task_id = ?
|
|
`), taskID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list task attachments for cleanup: %w", err)
|
|
}
|
|
var attachments []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("scan task attachment for cleanup: %w", err)
|
|
}
|
|
attachments = append(attachments, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("iterate task attachments for cleanup: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close task attachments for cleanup: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, tx.Rebind(`DELETE FROM task_message_attachments WHERE task_id = ?`), taskID); err != nil {
|
|
return nil, fmt.Errorf("delete task attachments: %w", err)
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return nil, fmt.Errorf("commit task attachment cleanup: %w", err)
|
|
}
|
|
return attachments, nil
|
|
}
|
|
|
|
// DeleteMessageAttachmentsBySession removes claimed attachment descriptors
|
|
// owned by a deleted task session. Staged uploads have no session binding and
|
|
// remain governed by their expiry cleanup.
|
|
func (r *Repository) DeleteMessageAttachmentsBySession(ctx context.Context, taskID, sessionID string) ([]*models.TaskMessageAttachment, error) {
|
|
if taskID == "" || sessionID == "" {
|
|
return nil, nil
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("begin session attachment cleanup: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments
|
|
WHERE task_id = ? AND session_id = ?
|
|
`), taskID, sessionID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list session attachments for cleanup: %w", err)
|
|
}
|
|
var attachments []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("scan session attachment for cleanup: %w", err)
|
|
}
|
|
attachments = append(attachments, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("iterate session attachments for cleanup: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close session attachments for cleanup: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
DELETE FROM task_message_attachments WHERE task_id = ? AND session_id = ?
|
|
`), taskID, sessionID); err != nil {
|
|
return nil, fmt.Errorf("delete session attachments: %w", err)
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return nil, fmt.Errorf("commit session attachment cleanup: %w", err)
|
|
}
|
|
return attachments, nil
|
|
}
|
|
|
|
func (r *Repository) DeleteMessageAttachmentsByWorkspaceTx(
|
|
ctx context.Context,
|
|
tx *sqlx.Tx,
|
|
workspaceID string,
|
|
) ([]*models.TaskMessageAttachment, error) {
|
|
if workspaceID == "" {
|
|
return nil, nil
|
|
}
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments WHERE workspace_id = ?
|
|
`), workspaceID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list workspace attachments for cleanup: %w", err)
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
var attachments []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
return nil, fmt.Errorf("scan workspace attachment for cleanup: %w", err)
|
|
}
|
|
attachments = append(attachments, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("iterate workspace attachments for cleanup: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, tx.Rebind(`DELETE FROM task_message_attachments WHERE workspace_id = ?`), workspaceID); err != nil {
|
|
return nil, fmt.Errorf("delete workspace attachments: %w", err)
|
|
}
|
|
return attachments, nil
|
|
}
|
|
|
|
func (r *Repository) TransferMessageAttachments(
|
|
ctx context.Context,
|
|
taskID, oldSessionID, newSessionID string,
|
|
attachmentIDs []string,
|
|
) error {
|
|
if taskID == "" || oldSessionID == "" || newSessionID == "" ||
|
|
oldSessionID == newSessionID || len(attachmentIDs) == 0 {
|
|
return nil
|
|
}
|
|
uniqueAttachmentIDs := uniqueAttachmentIDs(attachmentIDs)
|
|
queueLockPresent, err := r.queueSessionLockTablePresent(ctx)
|
|
if err != nil {
|
|
return fmt.Errorf("check queue session lock schema: %w", err)
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("begin transfer session attachments: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if queueLockPresent {
|
|
if err := messagequeue.LockSessionPairInTransaction(
|
|
ctx, tx, r.db, oldSessionID, newSessionID,
|
|
); err != nil {
|
|
return fmt.Errorf("lock transfer session pair: %w", err)
|
|
}
|
|
}
|
|
locations, err := transferAttachmentLocations(ctx, tx, taskID, uniqueAttachmentIDs)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
sourceAttachmentIDs, err := validateTransferAttachmentLocations(
|
|
taskID, oldSessionID, newSessionID, uniqueAttachmentIDs, locations,
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := updateTransferSessionAttachments(
|
|
ctx, tx, taskID, oldSessionID, newSessionID, sourceAttachmentIDs,
|
|
); err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("commit transfer session attachments: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func uniqueAttachmentIDs(attachmentIDs []string) []string {
|
|
unique := make([]string, 0, len(attachmentIDs))
|
|
seen := make(map[string]struct{}, len(attachmentIDs))
|
|
for _, attachmentID := range attachmentIDs {
|
|
if _, ok := seen[attachmentID]; ok {
|
|
continue
|
|
}
|
|
seen[attachmentID] = struct{}{}
|
|
unique = append(unique, attachmentID)
|
|
}
|
|
return unique
|
|
}
|
|
|
|
func transferAttachmentLocations(
|
|
ctx context.Context, tx *sqlx.Tx, taskID string, attachmentIDs []string,
|
|
) (map[string]string, error) {
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(attachmentIDs)), ",")
|
|
args := make([]interface{}, 0, 2+len(attachmentIDs))
|
|
args = append(args, taskID, models.AttachmentStateClaimed)
|
|
for _, attachmentID := range attachmentIDs {
|
|
args = append(args, attachmentID)
|
|
}
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT id, session_id
|
|
FROM task_message_attachments
|
|
WHERE task_id = ? AND state = ? AND id IN (`+placeholders+`)
|
|
`), args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("inspect transfer session attachments: %w", err)
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
locations := make(map[string]string, len(attachmentIDs))
|
|
for rows.Next() {
|
|
var attachmentID, sessionID string
|
|
if err := rows.Scan(&attachmentID, &sessionID); err != nil {
|
|
return nil, fmt.Errorf("scan transfer session attachment: %w", err)
|
|
}
|
|
locations[attachmentID] = sessionID
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("iterate transfer session attachments: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close transfer session attachments: %w", err)
|
|
}
|
|
return locations, nil
|
|
}
|
|
|
|
func validateTransferAttachmentLocations(
|
|
taskID, oldSessionID, newSessionID string, attachmentIDs []string, locations map[string]string,
|
|
) ([]string, error) {
|
|
sourceAttachmentIDs := make([]string, 0, len(attachmentIDs))
|
|
for _, attachmentID := range attachmentIDs {
|
|
sessionID, ok := locations[attachmentID]
|
|
if !ok || (sessionID != "" && sessionID != oldSessionID && sessionID != newSessionID) {
|
|
return nil, fmt.Errorf(
|
|
"transfer session attachments: attachment %q is not a claimed attachment for task %q in either transfer session",
|
|
attachmentID, taskID,
|
|
)
|
|
}
|
|
if sessionID == oldSessionID {
|
|
sourceAttachmentIDs = append(sourceAttachmentIDs, attachmentID)
|
|
}
|
|
}
|
|
return sourceAttachmentIDs, nil
|
|
}
|
|
|
|
func updateTransferSessionAttachments(
|
|
ctx context.Context, tx *sqlx.Tx, taskID, oldSessionID, newSessionID string, attachmentIDs []string,
|
|
) error {
|
|
if len(attachmentIDs) == 0 {
|
|
return nil
|
|
}
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(attachmentIDs)), ",")
|
|
args := make([]interface{}, 0, 5+len(attachmentIDs))
|
|
args = append(args, newSessionID, time.Now().UTC(), taskID, oldSessionID, models.AttachmentStateClaimed)
|
|
for _, attachmentID := range attachmentIDs {
|
|
args = append(args, attachmentID)
|
|
}
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments
|
|
SET session_id = ?, updated_at = ?
|
|
WHERE task_id = ? AND session_id = ? AND state = ?
|
|
AND id IN (`+placeholders+`)
|
|
`), args...)
|
|
if err != nil {
|
|
return fmt.Errorf("transfer session attachments: %w", err)
|
|
}
|
|
affected, err := result.RowsAffected()
|
|
if err != nil {
|
|
return fmt.Errorf("count transferred session attachments: %w", err)
|
|
}
|
|
if affected != int64(len(attachmentIDs)) {
|
|
return fmt.Errorf("transfer session attachments: expected %d claimed source attachments, moved %d", len(attachmentIDs), affected)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) MarkExpiredMessageAttachments(ctx context.Context, now time.Time) ([]*models.TaskMessageAttachment, error) {
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("begin expired attachment cleanup: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
rows, err := tx.QueryxContext(ctx, tx.Rebind(`
|
|
SELECT `+attachmentSelectColumns+` FROM task_message_attachments
|
|
WHERE state = ? AND expires_at <= ?
|
|
`), models.AttachmentStateStaged, now)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list expired attachments: %w", err)
|
|
}
|
|
var expired []*models.TaskMessageAttachment
|
|
for rows.Next() {
|
|
attachment := &models.TaskMessageAttachment{}
|
|
if err := rows.StructScan(attachment); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("scan expired attachment: %w", err)
|
|
}
|
|
expired = append(expired, attachment)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return nil, fmt.Errorf("iterate expired attachments: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, fmt.Errorf("close expired attachments: %w", err)
|
|
}
|
|
transitioned := make([]*models.TaskMessageAttachment, 0, len(expired))
|
|
for _, attachment := range expired {
|
|
result, err := tx.ExecContext(ctx, tx.Rebind(`
|
|
UPDATE task_message_attachments SET state = ?, updated_at = ?
|
|
WHERE id = ? AND state = ? AND expires_at <= ?
|
|
`), models.AttachmentStateExpired, now, attachment.ID, models.AttachmentStateStaged, now)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mark expired attachment: %w", err)
|
|
}
|
|
affected, err := result.RowsAffected()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mark expired attachment rows affected: %w", err)
|
|
}
|
|
if affected == 1 {
|
|
attachment.State = models.AttachmentStateExpired
|
|
attachment.UpdatedAt = now
|
|
transitioned = append(transitioned, attachment)
|
|
}
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return nil, fmt.Errorf("commit expired attachment cleanup: %w", err)
|
|
}
|
|
return transitioned, nil
|
|
}
|
|
|
|
func errorsIsNoRows(err error) bool { return err == sql.ErrNoRows }
|
|
|
|
func idsToInterfaces(ids []string) []interface{} {
|
|
args := make([]interface{}, len(ids))
|
|
for i, id := range ids {
|
|
args[i] = id
|
|
}
|
|
return args
|
|
}
|