本站源代码
Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.

135 lines
2.5KB

  1. package testfixtures
  2. import (
  3. "database/sql"
  4. "fmt"
  5. )
  6. // MySQL is the MySQL helper for this package
  7. type MySQL struct {
  8. baseHelper
  9. tables []string
  10. tablesChecksum map[string]int64
  11. }
  12. func (h *MySQL) init(db *sql.DB) error {
  13. var err error
  14. h.tables, err = h.tableNames(db)
  15. if err != nil {
  16. return err
  17. }
  18. return nil
  19. }
  20. func (*MySQL) paramType() int {
  21. return paramTypeQuestion
  22. }
  23. func (*MySQL) quoteKeyword(str string) string {
  24. return fmt.Sprintf("`%s`", str)
  25. }
  26. func (*MySQL) databaseName(q queryable) (string, error) {
  27. var dbName string
  28. err := q.QueryRow("SELECT DATABASE()").Scan(&dbName)
  29. return dbName, err
  30. }
  31. func (h *MySQL) tableNames(q queryable) ([]string, error) {
  32. query := `
  33. SELECT table_name
  34. FROM information_schema.tables
  35. WHERE table_schema = ?
  36. AND table_type = 'BASE TABLE';
  37. `
  38. dbName, err := h.databaseName(q)
  39. if err != nil {
  40. return nil, err
  41. }
  42. rows, err := q.Query(query, dbName)
  43. if err != nil {
  44. return nil, err
  45. }
  46. defer rows.Close()
  47. var tables []string
  48. for rows.Next() {
  49. var table string
  50. if err = rows.Scan(&table); err != nil {
  51. return nil, err
  52. }
  53. tables = append(tables, table)
  54. }
  55. if err = rows.Err(); err != nil {
  56. return nil, err
  57. }
  58. return tables, nil
  59. }
  60. func (h *MySQL) disableReferentialIntegrity(db *sql.DB, loadFn loadFunction) (err error) {
  61. // re-enable after load
  62. defer func() {
  63. if _, err2 := db.Exec("SET FOREIGN_KEY_CHECKS = 1"); err2 != nil && err == nil {
  64. err = err2
  65. }
  66. }()
  67. tx, err := db.Begin()
  68. if err != nil {
  69. return err
  70. }
  71. defer tx.Rollback()
  72. if _, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 0"); err != nil {
  73. return err
  74. }
  75. if err = loadFn(tx); err != nil {
  76. return err
  77. }
  78. return tx.Commit()
  79. }
  80. func (h *MySQL) isTableModified(q queryable, tableName string) (bool, error) {
  81. checksum, err := h.getChecksum(q, tableName)
  82. if err != nil {
  83. return true, err
  84. }
  85. oldChecksum := h.tablesChecksum[tableName]
  86. return oldChecksum == 0 || checksum != oldChecksum, nil
  87. }
  88. func (h *MySQL) afterLoad(q queryable) error {
  89. if h.tablesChecksum != nil {
  90. return nil
  91. }
  92. h.tablesChecksum = make(map[string]int64, len(h.tables))
  93. for _, t := range h.tables {
  94. checksum, err := h.getChecksum(q, t)
  95. if err != nil {
  96. return err
  97. }
  98. h.tablesChecksum[t] = checksum
  99. }
  100. return nil
  101. }
  102. func (h *MySQL) getChecksum(q queryable, tableName string) (int64, error) {
  103. sql := fmt.Sprintf("CHECKSUM TABLE %s", h.quoteKeyword(tableName))
  104. var (
  105. table string
  106. checksum int64
  107. )
  108. if err := q.QueryRow(sql).Scan(&table, &checksum); err != nil {
  109. return 0, err
  110. }
  111. return checksum, nil
  112. }
上海开阖软件有限公司 沪ICP备12045867号-1