feat: support shared tenant database sharding

server/internal/logic/sys/im_tenant_tables.go

@@ -0,0 +1,131 @@
+package sys
+
+import (
+	"context"
+	"fmt"
+	"hotgo/internal/library/imtenant"
+	"time"
+
+	"github.com/gogf/gf/v2/database/gdb"
+	"github.com/gogf/gf/v2/errors/gerror"
+)
+
+const imTenantIDSequenceTable = "im_tenant_id_sequence"
+
+func sharedDatabaseConfig(ctx context.Context) (dbName, dbDSN string, err error) {
+	row, err := imtenant.RegistryDB().GetOne(ctx, "SELECT db_name,db_dsn FROM im_tenant WHERE code='default' LIMIT 1")
+	if err != nil {
+		return "", "", gerror.Wrap(err, "获取共享业务库配置失败")
+	}
+	if row.IsEmpty() {
+		return "", "", gerror.New("默认商户不存在，无法确定共享业务库")
+	}
+	return row["db_name"].String(), row["db_dsn"].String(), nil
+}
+
+func ensureTenantIDSequence(ctx context.Context) error {
+	db := imtenant.RegistryDB()
+	if _, err := db.Exec(ctx, `CREATE TABLE IF NOT EXISTS im_tenant_id_sequence (
+		sequence_name VARCHAR(32) NOT NULL,
+		next_id BIGINT NOT NULL,
+		updated_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
+		PRIMARY KEY(sequence_name)
+	) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci`); err != nil {
+		return gerror.Wrap(err, "初始化商户编号序列失败")
+	}
+	_, err := db.Exec(ctx, `INSERT INTO im_tenant_id_sequence(sequence_name,next_id)
+		SELECT 'tenant',COALESCE(MAX(id),0)+1 FROM im_tenant
+		ON DUPLICATE KEY UPDATE next_id=GREATEST(next_id,VALUES(next_id))`)
+	return gerror.Wrap(err, "校准商户编号序列失败")
+}
+
+func nextTenantID(tx gdb.TX) (int64, error) {
+	value, err := tx.GetValue("SELECT next_id FROM " + imTenantIDSequenceTable + " WHERE sequence_name='tenant' FOR UPDATE")
+	if err != nil {
+		return 0, gerror.Wrap(err, "锁定商户编号序列失败")
+	}
+	id := value.Int64()
+	if id <= 0 {
+		return 0, gerror.New("商户编号序列不正确")
+	}
+	if _, err = tx.Exec("UPDATE "+imTenantIDSequenceTable+" SET next_id=? WHERE sequence_name='tenant'", id+1); err != nil {
+		return 0, gerror.Wrap(err, "更新商户编号序列失败")
+	}
+	return id, nil
+}
+
+func tenantTablePrefix(tenantID int64) string {
+	return fmt.Sprintf("im_%03d_", tenantID)
+}
+
+func provisionTenantMessageTables(ctx context.Context, tenantID int64) error {
+	if tenantID <= 0 {
+		return gerror.New("商户ID不能为空")
+	}
+	db, err := imtenant.SharedDB(ctx)
+	if err != nil {
+		return err
+	}
+	prefix := tenantTablePrefix(tenantID)
+	baseTables := [][2]string{
+		{prefix + "friendships", "im_friendships"},
+		{prefix + "messages", "im_messages"},
+	}
+	for _, pair := range baseTables {
+		if _, err = db.Exec(ctx, fmt.Sprintf("CREATE TABLE IF NOT EXISTS `%s` LIKE `%s`", pair[0], pair[1])); err != nil {
+			return gerror.Wrapf(err, "创建商户业务表 %s 失败", pair[0])
+		}
+		if _, err = db.Exec(ctx, fmt.Sprintf("ALTER TABLE `%s` MODIFY tenant_id BIGINT NOT NULL DEFAULT %d COMMENT '商户ID'", pair[0], tenantID)); err != nil {
+			return gerror.Wrapf(err, "设置商户业务表 %s 默认商户失败", pair[0])
+		}
+	}
+
+	first := time.Date(time.Now().UTC().Year(), time.Now().UTC().Month(), 1, 0, 0, 0, 0, time.UTC)
+	for _, offset := range []int{-1, 0, 1, 2} {
+		start := first.AddDate(0, offset, 0)
+		month := start.Format("200601")
+		table := prefix + "group_messages_" + month
+		if err = createTenantGroupMessageTable(ctx, db, table, tenantID); err != nil {
+			return err
+		}
+		status := "prepared"
+		if offset == -1 {
+			status = "sealing"
+		} else if offset == 0 {
+			status = "writable"
+		}
+		if _, err = db.Exec(ctx, `INSERT INTO im_group_message_shards
+			(tenant_id,shard_month,table_name,shard_start,shard_end,status)
+			VALUES(?,?,?,?,?,?) ON DUPLICATE KEY UPDATE table_name=VALUES(table_name),
+			shard_start=VALUES(shard_start),shard_end=VALUES(shard_end),
+			status=IF(status IN ('sealed','archived'),status,VALUES(status))`,
+			tenantID, month, table, start, start.AddDate(0, 1, 0), status); err != nil {
+			return gerror.Wrapf(err, "登记商户群消息月表 %s 失败", table)
+		}
+	}
+	return nil
+}
+
+func createTenantGroupMessageTable(ctx context.Context, db gdb.DB, table string, tenantID int64) error {
+	ddl := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
+		tenant_id BIGINT NOT NULL DEFAULT %d, message_id BIGINT NOT NULL, topic VARCHAR(64) NOT NULL, group_id BIGINT NOT NULL,
+		seqid BIGINT NOT NULL, from_user BIGINT NOT NULL, client_msg_id VARCHAR(64),
+		msg_type VARCHAR(32) NOT NULL DEFAULT 'text', status VARCHAR(32) NOT NULL DEFAULT 'normal',
+		content JSON, content_nonce VARBINARY(12), key_version INT NOT NULL DEFAULT 0,
+		head JSON, mentioned_user_ids JSON, mention_all TINYINT(1) NOT NULL DEFAULT 0,
+		reply_seqid BIGINT, root_seqid BIGINT, file_count INT NOT NULL DEFAULT 0,
+		content_size INT NOT NULL DEFAULT 0, has_risk TINYINT(1) NOT NULL DEFAULT 0,
+		audit_status VARCHAR(32) NOT NULL DEFAULT 'pass', version INT NOT NULL DEFAULT 1,
+		recalled_at DATETIME(3), recalled_by BIGINT, deleted_at DATETIME(3), expires_at DATETIME(3),
+		created_at DATETIME(3) NOT NULL, updated_at DATETIME(3) NOT NULL,
+		PRIMARY KEY(message_id), UNIQUE INDEX uk_group_seq(group_id, seqid),
+		UNIQUE INDEX uk_topic_seq(topic, seqid), UNIQUE INDEX uk_sender_client(topic, from_user, client_msg_id),
+		INDEX idx_group_status_seq(group_id, status, seqid), INDEX idx_topic_created(topic, created_at),
+		INDEX idx_topic_from_seq(topic, from_user, seqid), INDEX idx_audit_created(audit_status, created_at),
+		INDEX idx_expires(expires_at)
+	) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci`, table, tenantID)
+	if _, err := db.Exec(ctx, ddl); err != nil {
+		return gerror.Wrapf(err, "创建商户群消息月表 %s 失败", table)
+	}
+	return nil
+}