outbox/message_tracker_sqlx.go at master · nilorg/outbox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package outbox

import (
	"context"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/nilorg/sqlxplus"
)

// NewSqlxMessageTracker 创建基于 Sqlx 的消息追踪
func NewSqlxMessageTracker(db *sqlx.DB) MessageTracker {
	// 创建表(如果不存在)
	createTableSQL := `
		CREATE TABLE IF NOT EXISTS ` + ProcessedMessageTableName + ` (
			id BIGINT AUTO_INCREMENT PRIMARY KEY,
			message_id VARCHAR(255) NOT NULL UNIQUE,
			processed_at DATETIME NOT NULL,
			INDEX idx_message_id (message_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
	`
	db.Exec(createTableSQL)

	return &sqlxMessageTracker{
		db: db,
	}
}

type sqlxMessageTracker struct {
	db *sqlx.DB
}

func (t *sqlxMessageTracker) HasProcessed(ctx context.Context, msgID string) (exist bool, err error) {
	var count int
	query := `SELECT COUNT(*) FROM ` + ProcessedMessageTableName + ` WHERE message_id = ?`
	err = t.db.GetContext(ctx, &count, query, msgID)
	exist = count > 0
	return
}

func (t *sqlxMessageTracker) MarkAsProcessed(ctx context.Context, msgID string) (err error) {
	query := `INSERT INTO ` + ProcessedMessageTableName + ` (message_id, processed_at) VALUES (?, ?)`
	_, err = t.db.ExecContext(ctx, query, msgID, time.Now())
	return
}

// NewSqlxPlusMessageTracker 创建基于 SqlxPlus 的消息追踪
func NewSqlxPlusMessageTracker(db *sqlx.DB, log sqlxplus.Logger) MessageTracker {
	// 创建表(如果不存在)
	createTableSQL := `
		CREATE TABLE IF NOT EXISTS ` + ProcessedMessageTableName + ` (
			id BIGINT AUTO_INCREMENT PRIMARY KEY,
			message_id VARCHAR(255) NOT NULL UNIQUE,
			processed_at DATETIME NOT NULL,
			INDEX idx_message_id (message_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
	`
	db.Exec(createTableSQL)

	return &sqlxPlusMessageTracker{
		db:  db,
		log: log,
	}
}

type sqlxPlusMessageTracker struct {
	db  *sqlx.DB
	log sqlxplus.Logger
}

func (t *sqlxPlusMessageTracker) HasProcessed(ctx context.Context, msgID string) (exist bool, err error) {
	var count int
	query := `SELECT COUNT(*) FROM ` + ProcessedMessageTableName + ` WHERE message_id = ?`

	// 尝试从 context 获取 DbConnect（如果在事务中调用）
	var conn *sqlxplus.DBConnect
	conn, err = sqlxplus.FromSqlxContext(ctx)
	if err != nil {
		// 降级方案：不在事务中，直接使用传入的 db 执行查询
		err = t.db.GetContext(ctx, &count, query, msgID)
	} else {
		// 在事务中，使用 DBConnect.Adapter() 执行查询
		err = sqlx.GetContext(ctx, conn.Adapter(), &count, query, msgID)
	}
	exist = count > 0
	return
}

func (t *sqlxPlusMessageTracker) MarkAsProcessed(ctx context.Context, msgID string) (err error) {
	query := `INSERT INTO ` + ProcessedMessageTableName + ` (message_id, processed_at) VALUES (?, ?)`

	// 尝试从 context 获取 DbConnect（如果在事务中调用）
	var conn *sqlxplus.DBConnect
	conn, err = sqlxplus.FromSqlxContext(ctx)
	if err != nil {
		// 降级方案：不在事务中，直接使用传入的 db 执行
		_, err = t.db.ExecContext(ctx, query, msgID, time.Now())
	} else {
		// 在事务中，使用 sqlxplus.ExecContext 执行（带日志）
		_, err = sqlxplus.ExecContext(ctx, conn.Adapter(), query, msgID, time.Now())
	}
	return
}