migrations_runner.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601
  1. package repository
  2. import (
  3. "context"
  4. "crypto/sha256"
  5. "database/sql"
  6. "encoding/hex"
  7. "errors"
  8. "fmt"
  9. "io/fs"
  10. "sort"
  11. "strings"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/migrations"
  14. )
  15. // schemaMigrationsTableDDL 定义迁移记录表的 DDL。
  16. // 该表用于跟踪已应用的迁移文件及其校验和。
  17. // - filename: 迁移文件名,作为主键唯一标识每个迁移
  18. // - checksum: 文件内容的 SHA256 哈希值,用于检测迁移文件是否被篡改
  19. // - applied_at: 迁移应用时间戳
  20. const schemaMigrationsTableDDL = `
  21. CREATE TABLE IF NOT EXISTS schema_migrations (
  22. filename TEXT PRIMARY KEY,
  23. checksum TEXT NOT NULL,
  24. applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
  25. );
  26. `
  27. const atlasSchemaRevisionsTableDDL = `
  28. CREATE TABLE IF NOT EXISTS atlas_schema_revisions (
  29. version TEXT PRIMARY KEY,
  30. description TEXT NOT NULL,
  31. type INTEGER NOT NULL,
  32. applied INTEGER NOT NULL DEFAULT 0,
  33. total INTEGER NOT NULL DEFAULT 0,
  34. executed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
  35. execution_time BIGINT NOT NULL DEFAULT 0,
  36. error TEXT NULL,
  37. error_stmt TEXT NULL,
  38. hash TEXT NOT NULL DEFAULT '',
  39. partial_hashes TEXT[] NULL,
  40. operator_version TEXT NULL
  41. );
  42. `
  43. // migrationsAdvisoryLockID 是用于序列化迁移操作的 PostgreSQL Advisory Lock ID。
  44. // 在多实例部署场景下,该锁确保同一时间只有一个实例执行迁移。
  45. // 任何稳定的 int64 值都可以,只要不与同一数据库中的其他锁冲突即可。
  46. const migrationsAdvisoryLockID int64 = 694208311321144027
  47. const migrationsLockRetryInterval = 500 * time.Millisecond
  48. const nonTransactionalMigrationSuffix = "_notx.sql"
  49. const paymentOrdersOutTradeNoUniqueMigration = "120_enforce_payment_orders_out_trade_no_unique_notx.sql"
  50. const paymentOrdersOutTradeNoUniqueIndex = "paymentorder_out_trade_no_unique"
  51. const schedulerOutboxPendingDedupKeyMigration = "153_scheduler_outbox_pending_dedup_key_index_notx.sql"
  52. const schedulerOutboxPendingDedupKeyIndex = "idx_scheduler_outbox_pending_dedup_key"
  53. const latestAPIKeyIPIndexMigration = "174_add_usage_logs_api_key_latest_ip_index_notx.sql"
  54. const latestAPIKeyIPIndex = "idx_usage_logs_api_key_latest_ip"
  55. const usageLogsUpstreamModelMismatchIndexMigration = "195_add_usage_log_upstream_model_mismatch_index_notx.sql"
  56. const usageLogsUpstreamModelMismatchIndex = "idx_usage_logs_upstream_model_mismatch_created_at"
  57. const usageLogsEffectiveModelIndexesMigration = "226_add_usage_log_effective_model_indexes_notx.sql"
  58. const usageLogsEffectiveRequestedModelIndex = "idx_usage_logs_effective_requested_model_created"
  59. const usageLogsEffectiveUpstreamModelIndex = "idx_usage_logs_effective_upstream_model_created"
  60. const usageLogsUpstreamRequestIDIndexMigration = "233_add_usage_log_upstream_request_id_index_notx.sql"
  61. const usageLogsUpstreamRequestIDIndex = "idx_usage_logs_upstream_request_id"
  62. type migrationChecksumCompatibilityRule struct {
  63. fileChecksum string
  64. acceptedDBChecksum map[string]struct{}
  65. acceptedChecksums map[string]struct{}
  66. }
  67. // migrationChecksumCompatibilityRules 仅用于兼容历史上误修改过的迁移文件 checksum。
  68. // 规则必须同时匹配「迁移名 + 数据库 checksum + 当前文件 checksum」且两者都落在该迁移的已知版本集合内才会放行,
  69. // 避免放宽全局校验,也允许将误改的历史 migration 回滚为已发布版本而不要求人工修 checksum。
  70. var migrationChecksumCompatibilityRules = map[string]migrationChecksumCompatibilityRule{
  71. "054_drop_legacy_cache_columns.sql": newMigrationChecksumCompatibilityRule("82de761156e03876653e7a6a4eee883cd927847036f779b0b9f34c42a8af7a7d", "182c193f3359946cf094090cd9e57d5c3fd9abaffbc1e8fc378646b8a6fa12b4"),
  72. "061_add_usage_log_request_type.sql": newMigrationChecksumCompatibilityRule("66207e7aa5dd0429c2e2c0fabdaf79783ff157fa0af2e81adff2ee03790ec65c", "08a248652cbab7cfde147fc6ef8cda464f2477674e20b718312faa252e0481c0", "222b4a09c797c22e5922b6b172327c824f5463aaa8760e4f621bc5c22e2be0f3"),
  73. "109_auth_identity_compat_backfill.sql": newMigrationChecksumCompatibilityRule("0580b4602d85435edf9aca1633db580bb3932f26517f75134106f80275ec2ace", "551e498aa5616d2d91096e9d72cf9fb36e418ee22eacc557f8811cadbc9e20ee"),
  74. "110_pending_auth_and_provider_default_grants.sql": newMigrationChecksumCompatibilityRule("32cf87ee787b1bb36b5c691367c96eee37518fa3eed6f3322cf68795e3745279", "e3d1f433be2b564cfbdc549adf98fce13c5c7b363ebc20fd05b765d0563b0925"),
  75. "112_add_payment_order_provider_key_snapshot.sql": newMigrationChecksumCompatibilityRule("b75f8f56d39455682787696a3d92ad25b055444ca328fb7fca9a460a15d68d99", "ffd3e8a2c9295fa9cbefefd629a78268877e5b51bc970a82d9b3f46ec4ebd15e"),
  76. "115_auth_identity_legacy_external_backfill.sql": newMigrationChecksumCompatibilityRule("022aadd97bb53e755f0cf7a3a957e0cb1a1353b0c39ec4de3234acd2871fd04f", "4cf39e508be9fd1a5aa41610cbbebeb80385c9adda45bf78a706de9db4f1385f"),
  77. "116_auth_identity_legacy_external_safety_reports.sql": newMigrationChecksumCompatibilityRule("07edb09fa8d04ffb172b0621e3c22f4d1757d20a24ae267b3b36b087ab72d488", "f7757bd929ac67ffb08ce69fa4cf20fad39dbff9d5a5085fb2adabb7607e5877"),
  78. "118_wechat_dual_mode_and_auth_source_defaults.sql": newMigrationChecksumCompatibilityRule("b54194d7a3e4fbf710e0a3590d22a2fe7966804c487052a356e0b55f53ef96b0", "e0cdf835d6c688d64100f483d31bc02ac9ebad414bf1837af239a84bf75b8227", "a38243ca0a72c3a01c0a92b7986423054d6133c0399441f853b99802852720fb"),
  79. "119_enforce_payment_orders_out_trade_no_unique.sql": newMigrationChecksumCompatibilityRule("0bbe809ae48a9d811dabda1ba1c74955bd71c4a9cc610f9128816818dfa6c11e", "ebd2c67cce0116393fb4f1b5d5116a67c6aceb73820dfb5133d1ff6f36d72d34"),
  80. "120_enforce_payment_orders_out_trade_no_unique_notx.sql": newMigrationChecksumCompatibilityRule("34aadc0db59a4e390f92a12b73bd74642d9724f33124f73638ae00089ea5e074", "e77921f79d539bc24575cb9c16cbe566d2b23ce816190343d0a7568f6a3fcf61", "707431450603e70a43ce9fbd61e0c12fa67da4875158ccefabacea069587ab22", "04b082b5a239c525154fe9185d324ee2b05ff90da9297e10dba19f9be79aa59a"),
  81. "123_fix_legacy_auth_source_grant_on_signup_defaults.sql": newMigrationChecksumCompatibilityRule("2ce43c2cd89e9f9e1febd34a407ed9e84d177386c5544b6f02c1f58a21129f57", "6cd33422f215dcd1f486ab6f35c0ea5805d9ca69bb25906d94bc649156657145"),
  82. "159_batch_image_foundation.sql": newMigrationChecksumCompatibilityRule("d902b70982025ec519749faf058aab7631e82c3f48167b9a4ae4db718eb72cce", "82da85b5d98e67a0507647b873a40373e84538e4adafdeed6767c0ac8b6570b2"),
  83. "161_batch_image_pricing_snapshot.sql": newMigrationChecksumCompatibilityRule("4012af3e43636cb6af22e0176d59d1fcc70615c0f310194329461ae462c4fbd6", "96d915c9b7a6941ae99039e0ff3f1a61481eb9bddd933d11c6fadb2274554e87"),
  84. // 195 originally seeded mode=v2; flipped to v1 (safe default / opt-in v2). Existing DBs
  85. // that already applied the v2 seed keep their row and the historical checksum.
  86. "195_channel_monitor_mode.sql": newMigrationChecksumCompatibilityRule("13f3792f3e3e53ee96e26415c884cf8062c77172824b54fcc9a8c0c2b1f185ec", "4c74fe33ef2274cc72e1bb49671e651274532c034b29f5b2982c2a4c88d101a6"),
  87. // 220 originally cleared video prices for all non-grok platforms (including composite);
  88. // composite is now preserved because it may route to Grok accounts.
  89. "220_clear_non_grok_video_generation_config.sql": newMigrationChecksumCompatibilityRule("85e320b9ec64f2d3fcd8cf705b2b4e76a7b49f7a57140c14bff97f32691c818b", "3da48c8fdffe6390325f43d08b8e353e0a365df43d44a78dbbe655d0deb18402"),
  90. "219_group_search_price_per_1k.sql": newMigrationChecksumCompatibilityRule("e86786ebcc3b14206fd2d321380a4e50e80cdadbfcf4962c639255e6a14008db", "df6ffd71b97e30ec2c8fe7b95e15783042dea58c553e32701ee7c42a5619af80"),
  91. "218_group_audio_voice_pricing.sql": newMigrationChecksumCompatibilityRule("40ee9f3a2af0e0a5e99dabc878fd0fe98be1011f26bcfcefcac7197f7081f0e7", "c2a5e5b4ffd6968ad1c10593289fbc11192cdea19fec3ed9bce3a84eff9a8351"),
  92. }
  93. // ApplyMigrations 将嵌入的 SQL 迁移文件应用到指定的数据库。
  94. //
  95. // 该函数可以在每次应用启动时安全调用:
  96. // - 已应用的迁移会被自动跳过(通过校验 filename 判断)
  97. // - 如果迁移文件内容被修改(checksum 不匹配),会返回错误
  98. // - 使用 PostgreSQL Advisory Lock 确保多实例并发安全
  99. //
  100. // 参数:
  101. // - ctx: 上下文,用于超时控制和取消
  102. // - db: 数据库连接
  103. //
  104. // 返回:
  105. // - error: 迁移过程中的任何错误
  106. func ApplyMigrations(ctx context.Context, db *sql.DB) error {
  107. if db == nil {
  108. return errors.New("nil sql db")
  109. }
  110. return applyMigrationsFS(ctx, db, migrations.FS)
  111. }
  112. // applyMigrationsFS 是迁移执行的核心实现。
  113. // 它从指定的文件系统读取 SQL 迁移文件并按顺序应用。
  114. //
  115. // 迁移执行流程:
  116. // 1. 获取 PostgreSQL Advisory Lock,防止多实例并发迁移
  117. // 2. 确保 schema_migrations 表存在
  118. // 3. 按文件名排序读取所有 .sql 文件
  119. // 4. 对于每个迁移文件:
  120. // - 计算文件内容的 SHA256 校验和
  121. // - 检查该迁移是否已应用(通过 filename 查询)
  122. // - 如果已应用,验证校验和是否匹配
  123. // - 如果未应用,在事务中执行迁移并记录
  124. // 5. 释放 Advisory Lock
  125. //
  126. // 参数:
  127. // - ctx: 上下文
  128. // - db: 数据库连接
  129. // - fsys: 包含迁移文件的文件系统(通常是 embed.FS)
  130. func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error {
  131. if db == nil {
  132. return errors.New("nil sql db")
  133. }
  134. // 获取分布式锁,确保多实例部署时只有一个实例执行迁移。
  135. // 这是 PostgreSQL 特有的 Advisory Lock 机制。
  136. lockConn, err := db.Conn(ctx)
  137. if err != nil {
  138. return fmt.Errorf("acquire migrations lock connection: %w", err)
  139. }
  140. defer func() { _ = lockConn.Close() }()
  141. if err := pgAdvisoryLock(ctx, lockConn); err != nil {
  142. return err
  143. }
  144. defer func() {
  145. // 无论迁移是否成功,都要释放锁。
  146. // 独立超时确保原 ctx 取消后仍会尝试释放,但数据库链路异常不会
  147. // 无限阻塞进程退出。
  148. unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  149. defer cancel()
  150. _ = pgAdvisoryUnlock(unlockCtx, lockConn)
  151. }()
  152. // 创建迁移记录表(如果不存在)。
  153. // 该表记录所有已应用的迁移及其校验和。
  154. if _, err := lockConn.ExecContext(ctx, schemaMigrationsTableDDL); err != nil {
  155. return fmt.Errorf("create schema_migrations: %w", err)
  156. }
  157. // 自动对齐 Atlas 基线(如果检测到 legacy schema_migrations 且缺失 atlas_schema_revisions)。
  158. if err := ensureAtlasBaselineAligned(ctx, lockConn, fsys); err != nil {
  159. return err
  160. }
  161. // 获取所有 .sql 迁移文件并按文件名排序。
  162. // 命名规范:使用零填充数字前缀(如 001_init.sql, 002_add_users.sql)。
  163. files, err := fs.Glob(fsys, "*.sql")
  164. if err != nil {
  165. return fmt.Errorf("list migrations: %w", err)
  166. }
  167. sort.Strings(files) // 确保按文件名顺序执行迁移
  168. for _, name := range files {
  169. // 读取迁移文件内容
  170. contentBytes, err := fs.ReadFile(fsys, name)
  171. if err != nil {
  172. return fmt.Errorf("read migration %s: %w", name, err)
  173. }
  174. content := strings.TrimSpace(string(contentBytes))
  175. if content == "" {
  176. continue // 跳过空文件
  177. }
  178. // 计算文件内容的 SHA256 校验和,用于检测文件是否被修改。
  179. // 这是一种防篡改机制:如果有人修改了已应用的迁移文件,系统会拒绝启动。
  180. sum := sha256.Sum256([]byte(content))
  181. checksum := hex.EncodeToString(sum[:])
  182. // 检查该迁移是否已经应用
  183. var existing string
  184. rowErr := lockConn.QueryRowContext(ctx, "SELECT checksum FROM schema_migrations WHERE filename = $1", name).Scan(&existing)
  185. if rowErr == nil {
  186. // 迁移已应用,验证校验和是否匹配
  187. if existing != checksum {
  188. // 兼容特定历史误改场景(仅白名单规则),其余仍保持严格不可变约束。
  189. if isMigrationChecksumCompatible(name, existing, checksum) {
  190. continue
  191. }
  192. // 校验和不匹配意味着迁移文件在应用后被修改,这是危险的。
  193. // 正确的做法是创建新的迁移文件来进行变更。
  194. return fmt.Errorf(
  195. "migration %s checksum mismatch (db=%s file=%s)\n"+
  196. "This means the migration file was modified after being applied to the database.\n"+
  197. "Solutions:\n"+
  198. " 1. Revert to original: git log --oneline -- migrations/%s && git checkout <commit> -- migrations/%s\n"+
  199. " 2. For new changes, create a new migration file instead of modifying existing ones\n"+
  200. "Note: Modifying applied migrations breaks the immutability principle and can cause inconsistencies across environments",
  201. name, existing, checksum, name, name,
  202. )
  203. }
  204. continue // 迁移已应用且校验和匹配,跳过
  205. }
  206. if !errors.Is(rowErr, sql.ErrNoRows) {
  207. return fmt.Errorf("check migration %s: %w", name, rowErr)
  208. }
  209. nonTx, err := validateMigrationExecutionMode(name, content)
  210. if err != nil {
  211. return fmt.Errorf("validate migration %s: %w", name, err)
  212. }
  213. if nonTx {
  214. if err := prepareNonTransactionalMigration(ctx, lockConn, name); err != nil {
  215. return fmt.Errorf("prepare migration %s: %w", name, err)
  216. }
  217. // *_notx.sql:用于 CREATE/DROP INDEX CONCURRENTLY 场景,必须非事务执行。
  218. // 逐条语句执行,避免将多条 CONCURRENTLY 语句放入同一个隐式事务块。
  219. statements := splitSQLStatements(content)
  220. for i, stmt := range statements {
  221. trimmed := strings.TrimSpace(stmt)
  222. if trimmed == "" {
  223. continue
  224. }
  225. if stripSQLLineComment(trimmed) == "" {
  226. continue
  227. }
  228. if _, err := lockConn.ExecContext(ctx, trimmed); err != nil {
  229. return fmt.Errorf("apply migration %s (non-tx statement %d): %w", name, i+1, err)
  230. }
  231. }
  232. if _, err := lockConn.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil {
  233. return fmt.Errorf("record migration %s (non-tx): %w", name, err)
  234. }
  235. continue
  236. }
  237. // 默认迁移在事务中执行,确保原子性:要么完全成功,要么完全回滚。
  238. tx, err := lockConn.BeginTx(ctx, nil)
  239. if err != nil {
  240. return fmt.Errorf("begin migration %s: %w", name, err)
  241. }
  242. // 执行迁移 SQL
  243. if _, err := tx.ExecContext(ctx, content); err != nil {
  244. _ = tx.Rollback()
  245. return fmt.Errorf("apply migration %s: %w", name, err)
  246. }
  247. // 记录迁移已完成,保存文件名和校验和
  248. if _, err := tx.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil {
  249. _ = tx.Rollback()
  250. return fmt.Errorf("record migration %s: %w", name, err)
  251. }
  252. // 提交事务
  253. if err := tx.Commit(); err != nil {
  254. _ = tx.Rollback()
  255. return fmt.Errorf("commit migration %s: %w", name, err)
  256. }
  257. }
  258. return nil
  259. }
  260. type migrationConnection interface {
  261. ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
  262. QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
  263. QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
  264. BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
  265. }
  266. func prepareNonTransactionalMigration(ctx context.Context, db migrationConnection, name string) error {
  267. switch name {
  268. case paymentOrdersOutTradeNoUniqueMigration:
  269. return preparePaymentOrdersOutTradeNoUniqueMigration(ctx, db)
  270. case schedulerOutboxPendingDedupKeyMigration:
  271. return dropInvalidIndexIfPresent(ctx, db, schedulerOutboxPendingDedupKeyIndex)
  272. case latestAPIKeyIPIndexMigration:
  273. return dropInvalidIndexIfPresent(ctx, db, latestAPIKeyIPIndex)
  274. case usageLogsUpstreamModelMismatchIndexMigration:
  275. return dropInvalidIndexIfPresent(ctx, db, usageLogsUpstreamModelMismatchIndex)
  276. case usageLogsEffectiveModelIndexesMigration:
  277. for _, indexName := range []string{usageLogsEffectiveRequestedModelIndex, usageLogsEffectiveUpstreamModelIndex} {
  278. if err := dropInvalidIndexIfPresent(ctx, db, indexName); err != nil {
  279. return err
  280. }
  281. }
  282. return nil
  283. case usageLogsUpstreamRequestIDIndexMigration:
  284. return dropInvalidIndexIfPresent(ctx, db, usageLogsUpstreamRequestIDIndex)
  285. default:
  286. return nil
  287. }
  288. }
  289. func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db migrationConnection) error {
  290. duplicates, err := findDuplicatePaymentOrderOutTradeNos(ctx, db)
  291. if err != nil {
  292. return fmt.Errorf("precheck duplicate out_trade_no: %w", err)
  293. }
  294. if len(duplicates) > 0 {
  295. return fmt.Errorf(
  296. "duplicate out_trade_no values block %s; remediate duplicates before retrying: %s",
  297. paymentOrdersOutTradeNoUniqueMigration,
  298. strings.Join(duplicates, ", "),
  299. )
  300. }
  301. return dropInvalidIndexIfPresent(ctx, db, paymentOrdersOutTradeNoUniqueIndex)
  302. }
  303. func dropInvalidIndexIfPresent(ctx context.Context, db migrationConnection, indexName string) error {
  304. invalid, err := indexIsInvalid(ctx, db, indexName)
  305. if err != nil {
  306. return fmt.Errorf("check invalid index %s: %w", indexName, err)
  307. }
  308. if !invalid {
  309. return nil
  310. }
  311. if _, err := db.ExecContext(ctx, fmt.Sprintf("DROP INDEX CONCURRENTLY IF EXISTS %s", indexName)); err != nil {
  312. return fmt.Errorf("drop invalid index %s: %w", indexName, err)
  313. }
  314. return nil
  315. }
  316. func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db migrationConnection) ([]string, error) {
  317. rows, err := db.QueryContext(ctx, `
  318. SELECT out_trade_no, COUNT(*) AS duplicate_count
  319. FROM payment_orders
  320. WHERE out_trade_no <> ''
  321. GROUP BY out_trade_no
  322. HAVING COUNT(*) > 1
  323. ORDER BY duplicate_count DESC, out_trade_no
  324. LIMIT 5
  325. `)
  326. if err != nil {
  327. return nil, err
  328. }
  329. defer func() {
  330. _ = rows.Close()
  331. }()
  332. duplicates := make([]string, 0, 5)
  333. for rows.Next() {
  334. var outTradeNo string
  335. var duplicateCount int
  336. if err := rows.Scan(&outTradeNo, &duplicateCount); err != nil {
  337. return nil, err
  338. }
  339. duplicates = append(duplicates, fmt.Sprintf("%s (count=%d)", outTradeNo, duplicateCount))
  340. }
  341. if err := rows.Err(); err != nil {
  342. return nil, err
  343. }
  344. return duplicates, nil
  345. }
  346. func indexIsInvalid(ctx context.Context, db migrationConnection, indexName string) (bool, error) {
  347. var invalid bool
  348. err := db.QueryRowContext(ctx, `
  349. SELECT EXISTS (
  350. SELECT 1
  351. FROM pg_class idx
  352. JOIN pg_namespace ns ON ns.oid = idx.relnamespace
  353. JOIN pg_index i ON i.indexrelid = idx.oid
  354. WHERE ns.nspname = 'public'
  355. AND idx.relname = $1
  356. AND NOT i.indisvalid
  357. )
  358. `, indexName).Scan(&invalid)
  359. return invalid, err
  360. }
  361. func ensureAtlasBaselineAligned(ctx context.Context, db migrationConnection, fsys fs.FS) error {
  362. hasLegacy, err := tableExists(ctx, db, "schema_migrations")
  363. if err != nil {
  364. return fmt.Errorf("check schema_migrations: %w", err)
  365. }
  366. if !hasLegacy {
  367. return nil
  368. }
  369. hasAtlas, err := tableExists(ctx, db, "atlas_schema_revisions")
  370. if err != nil {
  371. return fmt.Errorf("check atlas_schema_revisions: %w", err)
  372. }
  373. if !hasAtlas {
  374. if _, err := db.ExecContext(ctx, atlasSchemaRevisionsTableDDL); err != nil {
  375. return fmt.Errorf("create atlas_schema_revisions: %w", err)
  376. }
  377. }
  378. var count int
  379. if err := db.QueryRowContext(ctx, "SELECT COUNT(*) FROM atlas_schema_revisions").Scan(&count); err != nil {
  380. return fmt.Errorf("count atlas_schema_revisions: %w", err)
  381. }
  382. if count > 0 {
  383. return nil
  384. }
  385. version, description, hash, err := latestMigrationBaseline(fsys)
  386. if err != nil {
  387. return fmt.Errorf("atlas baseline version: %w", err)
  388. }
  389. if _, err := db.ExecContext(ctx, `
  390. INSERT INTO atlas_schema_revisions (version, description, type, applied, total, executed_at, execution_time, hash)
  391. VALUES ($1, $2, $3, 0, 0, NOW(), 0, $4)
  392. `, version, description, 1, hash); err != nil {
  393. return fmt.Errorf("insert atlas baseline: %w", err)
  394. }
  395. return nil
  396. }
  397. func tableExists(ctx context.Context, db migrationConnection, tableName string) (bool, error) {
  398. var exists bool
  399. err := db.QueryRowContext(ctx, `
  400. SELECT EXISTS (
  401. SELECT 1
  402. FROM information_schema.tables
  403. WHERE table_schema = 'public' AND table_name = $1
  404. )
  405. `, tableName).Scan(&exists)
  406. return exists, err
  407. }
  408. func latestMigrationBaseline(fsys fs.FS) (string, string, string, error) {
  409. files, err := fs.Glob(fsys, "*.sql")
  410. if err != nil {
  411. return "", "", "", err
  412. }
  413. if len(files) == 0 {
  414. return "baseline", "baseline", "", nil
  415. }
  416. sort.Strings(files)
  417. name := files[len(files)-1]
  418. contentBytes, err := fs.ReadFile(fsys, name)
  419. if err != nil {
  420. return "", "", "", err
  421. }
  422. content := strings.TrimSpace(string(contentBytes))
  423. sum := sha256.Sum256([]byte(content))
  424. hash := hex.EncodeToString(sum[:])
  425. version := strings.TrimSuffix(name, ".sql")
  426. return version, version, hash, nil
  427. }
  428. func checksumSet(values ...string) map[string]struct{} {
  429. out := make(map[string]struct{}, len(values))
  430. for _, value := range values {
  431. out[value] = struct{}{}
  432. }
  433. return out
  434. }
  435. func newMigrationChecksumCompatibilityRule(fileChecksum string, acceptedDBChecksums ...string) migrationChecksumCompatibilityRule {
  436. return migrationChecksumCompatibilityRule{
  437. fileChecksum: fileChecksum,
  438. acceptedDBChecksum: checksumSet(acceptedDBChecksums...),
  439. acceptedChecksums: checksumSet(append([]string{fileChecksum}, acceptedDBChecksums...)...),
  440. }
  441. }
  442. func isMigrationChecksumCompatible(name, dbChecksum, fileChecksum string) bool {
  443. rule, ok := migrationChecksumCompatibilityRules[name]
  444. if !ok {
  445. return false
  446. }
  447. _, dbOK := rule.acceptedChecksums[dbChecksum]
  448. if !dbOK {
  449. return false
  450. }
  451. _, fileOK := rule.acceptedChecksums[fileChecksum]
  452. return fileOK
  453. }
  454. func validateMigrationExecutionMode(name, content string) (bool, error) {
  455. normalizedName := strings.ToLower(strings.TrimSpace(name))
  456. upperContent := strings.ToUpper(content)
  457. nonTx := strings.HasSuffix(normalizedName, nonTransactionalMigrationSuffix)
  458. if !nonTx {
  459. if strings.Contains(upperContent, "CONCURRENTLY") {
  460. return false, errors.New("CONCURRENTLY statements must be placed in *_notx.sql migrations")
  461. }
  462. return false, nil
  463. }
  464. if strings.Contains(upperContent, "BEGIN") || strings.Contains(upperContent, "COMMIT") || strings.Contains(upperContent, "ROLLBACK") {
  465. return false, errors.New("*_notx.sql must not contain transaction control statements (BEGIN/COMMIT/ROLLBACK)")
  466. }
  467. statements := splitSQLStatements(content)
  468. for _, stmt := range statements {
  469. normalizedStmt := strings.ToUpper(stripSQLLineComment(strings.TrimSpace(stmt)))
  470. if normalizedStmt == "" {
  471. continue
  472. }
  473. if strings.Contains(normalizedStmt, "CONCURRENTLY") {
  474. isCreateIndex := strings.Contains(normalizedStmt, "CREATE") && strings.Contains(normalizedStmt, "INDEX")
  475. isDropIndex := strings.Contains(normalizedStmt, "DROP") && strings.Contains(normalizedStmt, "INDEX")
  476. if !isCreateIndex && !isDropIndex {
  477. return false, errors.New("*_notx.sql currently only supports CREATE/DROP INDEX CONCURRENTLY statements")
  478. }
  479. if isCreateIndex && !strings.Contains(normalizedStmt, "IF NOT EXISTS") {
  480. return false, errors.New("CREATE INDEX CONCURRENTLY in *_notx.sql must include IF NOT EXISTS for idempotency")
  481. }
  482. if isDropIndex && !strings.Contains(normalizedStmt, "IF EXISTS") {
  483. return false, errors.New("DROP INDEX CONCURRENTLY in *_notx.sql must include IF EXISTS for idempotency")
  484. }
  485. continue
  486. }
  487. return false, errors.New("*_notx.sql must not mix non-CONCURRENTLY SQL statements")
  488. }
  489. return true, nil
  490. }
  491. func splitSQLStatements(content string) []string {
  492. parts := strings.Split(content, ";")
  493. out := make([]string, 0, len(parts))
  494. for _, part := range parts {
  495. if strings.TrimSpace(part) == "" {
  496. continue
  497. }
  498. out = append(out, part)
  499. }
  500. return out
  501. }
  502. func stripSQLLineComment(s string) string {
  503. lines := strings.Split(s, "\n")
  504. for i, line := range lines {
  505. if idx := strings.Index(line, "--"); idx >= 0 {
  506. lines[i] = line[:idx]
  507. }
  508. }
  509. return strings.TrimSpace(strings.Join(lines, "\n"))
  510. }
  511. // pgAdvisoryLock 获取 PostgreSQL Advisory Lock。
  512. // Advisory Lock 是一种轻量级的锁机制,不与任何特定的数据库对象关联。
  513. // 它非常适合用于应用层面的分布式锁场景,如迁移序列化。
  514. type advisoryLockConnection interface {
  515. QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
  516. ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
  517. }
  518. func pgAdvisoryLock(ctx context.Context, db advisoryLockConnection) error {
  519. ticker := time.NewTicker(migrationsLockRetryInterval)
  520. defer ticker.Stop()
  521. for {
  522. var locked bool
  523. if err := db.QueryRowContext(ctx, "SELECT pg_try_advisory_lock($1)", migrationsAdvisoryLockID).Scan(&locked); err != nil {
  524. return fmt.Errorf("acquire migrations lock: %w", err)
  525. }
  526. if locked {
  527. return nil
  528. }
  529. select {
  530. case <-ctx.Done():
  531. return fmt.Errorf("acquire migrations lock: %w", ctx.Err())
  532. case <-ticker.C:
  533. }
  534. }
  535. }
  536. // pgAdvisoryUnlock 释放 PostgreSQL Advisory Lock。
  537. // 必须在获取锁后确保释放,否则会阻塞其他实例的迁移操作。
  538. func pgAdvisoryUnlock(ctx context.Context, db advisoryLockConnection) error {
  539. _, err := db.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", migrationsAdvisoryLockID)
  540. if err != nil {
  541. return fmt.Errorf("release migrations lock: %w", err)
  542. }
  543. return nil
  544. }