forked from vmware-archive/pg2mysql
/
migrator.go
210 lines (176 loc) · 4.88 KB
/
migrator.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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
package pg2mysql
import (
"database/sql"
"errors"
"fmt"
"os"
"strings"
)
type Migrator interface {
Migrate() error
}
func NewMigrator(src, dst DB, truncateFirst bool, watcher MigratorWatcher) Migrator {
return &migrator{
src: src,
dst: dst,
truncateFirst: truncateFirst,
watcher: watcher,
}
}
type migrator struct {
src, dst DB
truncateFirst bool
watcher MigratorWatcher
}
func (m *migrator) Migrate() error {
srcSchema, err := BuildSchema(m.src)
if err != nil {
return fmt.Errorf("failed to build source schema: %s", err)
}
m.watcher.WillDisableConstraints()
err = m.dst.DisableConstraints()
if err != nil {
return fmt.Errorf("failed to disable constraints: %s", err)
}
m.watcher.DidDisableConstraints()
defer func() {
m.watcher.WillEnableConstraints()
err = m.dst.EnableConstraints()
if err != nil {
m.watcher.EnableConstraintsDidFailWithError(err)
} else {
m.watcher.EnableConstraintsDidFinish()
}
}()
for _, table := range srcSchema.Tables {
if m.truncateFirst {
m.watcher.WillTruncateTable(table.Name)
_, err := m.dst.DB().Exec(fmt.Sprintf("TRUNCATE TABLE %s", table.Name))
if err != nil {
return fmt.Errorf("failed truncating: %s", err)
}
m.watcher.TruncateTableDidFinish(table.Name)
}
columnNamesForInsert := make([]string, len(table.Columns))
placeholders := make([]string, len(table.Columns))
for i := range table.Columns {
columnNamesForInsert[i] = fmt.Sprintf("`%s`", table.Columns[i].Name)
placeholders[i] = "?"
}
preparedStmt, err := m.dst.DB().Prepare(fmt.Sprintf(
"INSERT INTO %s (%s) VALUES (%s)",
table.Name,
strings.Join(columnNamesForInsert, ","),
strings.Join(placeholders, ","),
))
if err != nil {
return fmt.Errorf("failed creating prepared statement: %s", err)
}
var recordsInserted int64
m.watcher.TableMigrationDidStart(table.Name)
if table.HasColumn("id") {
err = migrateWithIDs(m.watcher, m.src, m.dst, table, &recordsInserted, preparedStmt)
if err != nil {
return fmt.Errorf("failed migrating table with ids: %s", err)
}
} else {
err = EachMissingRow(m.src, m.dst, table, func(scanArgs []interface{}) {
err = insert(preparedStmt, scanArgs)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to insert into %s: %s\n", table.Name, err)
return
}
recordsInserted++
})
if err != nil {
return fmt.Errorf("failed migrating table without ids: %s", err)
}
}
m.watcher.TableMigrationDidFinish(table.Name, recordsInserted)
}
return nil
}
func migrateWithIDs(
watcher MigratorWatcher,
src DB,
dst DB,
table *Table,
recordsInserted *int64,
preparedStmt *sql.Stmt,
) error {
columnNamesForSelect := make([]string, len(table.Columns))
values := make([]interface{}, len(table.Columns))
scanArgs := make([]interface{}, len(table.Columns))
for i := range table.Columns {
columnNamesForSelect[i] = table.Columns[i].Name
scanArgs[i] = &values[i]
}
// find ids already in dst
rows, err := dst.DB().Query(fmt.Sprintf("SELECT id FROM %s", table.Name))
if err != nil {
return fmt.Errorf("failed to select id from rows: %s", err)
}
var dstIDs []interface{}
for rows.Next() {
var id interface{}
if err = rows.Scan(&id); err != nil {
return fmt.Errorf("failed to scan id from row: %s", err)
}
dstIDs = append(dstIDs, id)
}
if err = rows.Err(); err != nil {
return fmt.Errorf("failed iterating through rows: %s", err)
}
if err = rows.Close(); err != nil {
return fmt.Errorf("failed closing rows: %s", err)
}
// select data for ids to migrate from src
stmt := fmt.Sprintf(
"SELECT %s FROM %s",
strings.Join(columnNamesForSelect, ","),
table.Name,
)
if len(dstIDs) > 0 {
placeholders := make([]string, len(dstIDs))
for i := range dstIDs {
placeholders[i] = fmt.Sprintf("$%d", i+1)
}
stmt = fmt.Sprintf("%s WHERE id NOT IN (%s)", stmt, strings.Join(placeholders, ","))
}
rows, err = src.DB().Query(stmt, dstIDs...)
if err != nil {
return fmt.Errorf("failed to select rows: %s", err)
}
for rows.Next() {
if err = rows.Scan(scanArgs...); err != nil {
return fmt.Errorf("failed to scan row: %s", err)
}
err = insert(preparedStmt, scanArgs)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to insert into %s: %s\n", table.Name, err)
continue
}
*recordsInserted++
}
if err = rows.Err(); err != nil {
return fmt.Errorf("failed iterating through rows: %s", err)
}
if err = rows.Close(); err != nil {
return fmt.Errorf("failed closing rows: %s", err)
}
return nil
}
func insert(stmt *sql.Stmt, values []interface{}) error {
result, err := stmt.Exec(values...)
if err != nil {
return fmt.Errorf("failed to exec stmt: %s", err)
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return fmt.Errorf("failed getting rows affected by insert: %s", err)
}
if rowsAffected == 0 {
return errors.New("no rows affected by insert")
}
return nil
}