-
-
Notifications
You must be signed in to change notification settings - Fork 121
/
migrate.go
154 lines (128 loc) · 4.42 KB
/
migrate.go
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
package dao
import (
"os"
"github.com/ArtalkJS/Artalk/internal/entity"
"github.com/ArtalkJS/Artalk/internal/log"
)
func (dao *Dao) MigrateModels() {
// Upgrade the database
if dao.DB().Migrator().HasTable(&entity.Comment{}) &&
(!dao.DB().Migrator().HasColumn(&entity.Comment{}, "root_id") ||
os.Getenv("ATK_DB_MIGRATOR_FUNC_MIGRATE_ROOT_ID") == "1") {
dao.MigrateRootID()
}
// Migrate the schema
dao.DB().AutoMigrate(&entity.Site{}, &entity.Page{}, &entity.User{},
&entity.AuthIdentity{}, &entity.UserEmailVerify{},
&entity.Comment{}, &entity.Notify{}, &entity.Vote{})
// Delete all foreign key constraints
// Leave relationship maintenance to the program and reduce the difficulty of database management.
// because there are many different DBs and the implementation of foreign keys may be different,
// and the DB may not support foreign keys, so don't rely on the foreign key function of the DB system.
dao.DropConstraintsIfExist()
// Merge pages
if os.Getenv("ATK_DB_MIGRATOR_FUNC_MERGE_PAGES") == "1" {
dao.MergePages()
}
}
// Remove all constraints
func (dao *Dao) DropConstraintsIfExist() {
if dao.DB().Dialector.Name() == "sqlite" {
return // sqlite dose not support constraints by default
}
TAG := "[DB Migrator] "
list := []struct {
model any
constraint string
}{
{&entity.Comment{}, "fk_comments_page"},
{&entity.Comment{}, "fk_comments_user"},
{&entity.Page{}, "fk_pages_site"},
{&entity.User{}, "fk_comments_user"},
}
for _, item := range list {
if dao.DB().Migrator().HasConstraint(item.model, item.constraint) {
log.Info(TAG, "Dropping constraint: ", item.constraint)
err := dao.DB().Migrator().DropConstraint(item.model, item.constraint)
if err != nil {
log.Fatal(TAG, "Failed to drop constraint: ", item.constraint)
}
}
}
}
func (dao *Dao) MigrateRootID() {
const TAG = "[DB Migrator] "
log.Info(TAG, "Generating Root IDs...")
if !dao.DB().Migrator().HasColumn(&entity.Comment{}, "root_id") {
dao.DB().Migrator().AddColumn(&entity.Comment{}, "root_id")
}
if err := dao.DB().Raw(`WITH RECURSIVE CommentHierarchy AS (
SELECT id, id AS root_id, rid
FROM comments
WHERE rid = 0
UNION ALL
SELECT c.id, ch.root_id, c.rid
FROM comments c
INNER JOIN CommentHierarchy ch ON c.rid = ch.id
)
UPDATE comments SET root_id = (
SELECT root_id
FROM CommentHierarchy
WHERE comments.id = CommentHierarchy.id
);
`).Scan(&struct{}{}).Error; err == nil {
// no error, then do some patch
dao.DB().Table("comments").Where("id = root_id").Update("root_id", 0)
} else {
// try backup plan (if recursive CTE is not supported)
log.Info(TAG, "Recursive CTE is not supported, trying backup plan... Please wait a moment. This may take a long time if there are many comments.")
comments := []entity.Comment{}
if err := dao.DB().Find(&comments).Error; err != nil {
log.Fatal(TAG, "Failed to load comments. ", err.Error)
}
// update root_id
for _, comment := range comments {
if err := dao.DB().Model(&comment).Update("root_id", dao.FindCommentRootID(comment.ID)).Error; err != nil {
log.Error(TAG, "Failed to update root ID. ", err.Error, " ID=", comment.ID)
}
}
}
log.Info(TAG, "Root IDs generated successfully.")
}
func (dao *Dao) MergePages() {
// merge pages with same key and site_name, sum pv
pages := []*entity.Page{}
// load all pages
if err := dao.DB().Order("id ASC").Find(&pages).Error; err != nil {
log.Fatal("Failed to load pages. ", err.Error)
}
beforeLen := len(pages)
// merge pages
mergedPages := map[string]*entity.Page{}
for _, page := range pages {
key := page.SiteName + page.Key
if _, ok := mergedPages[key]; !ok {
mergedPages[key] = page
} else {
mergedPages[key].PV += page.PV
mergedPages[key].VoteUp += page.VoteUp
mergedPages[key].VoteDown += page.VoteDown
}
}
// delete all pages
dao.DB().Exec("DELETE FROM pages")
// insert merged pages
pages = []*entity.Page{}
for _, page := range mergedPages {
pages = append(pages, page)
}
if err := dao.DB().CreateInBatches(pages, 1000); err.Error != nil {
log.Fatal("Failed to insert merged pages. ", err.Error)
}
// drop page AccessibleURL column
if dao.DB().Migrator().HasColumn(&entity.Page{}, "accessible_url") {
dao.DB().Migrator().DropColumn(&entity.Page{}, "accessible_url")
}
log.Info("Pages merged successfully. Before pages: ", beforeLen, ", After pages: ", len(mergedPages), ", Deleted pages: ", beforeLen-len(mergedPages))
os.Exit(0)
}