Skip to content

Commit d66f8c4

Browse files
BrandonRoehlBrandonRoehl
authored
Merge pull request #3 from BrandonRoehl/master
Now stream INSERT statements.
2 parents 298b70a + e58e36d commit d66f8c4

File tree

3 files changed

+178
-28
lines changed

3 files changed

+178
-28
lines changed

‎dump.go

Lines changed: 62 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
package mysqldump
22

33
import (
4+
"bytes"
45
"database/sql"
56
"errors"
67
"fmt"
78
"io"
89
"reflect"
9-
"strings"
1010
"text/template"
1111
"time"
1212
)
@@ -19,9 +19,10 @@ Data struct to configure dump behavior
1919
IgnoreTables: Mark sensitive tables to ignore
2020
*/
2121
type Data struct {
22-
Out io.Writer
23-
Connection *sql.DB
24-
IgnoreTables []string
22+
Out io.Writer
23+
Connection *sql.DB
24+
IgnoreTables []string
25+
MaxAllowedPacket int
2526

2627
headerTmpl *template.Template
2728
tableTmpl *template.Template
@@ -45,7 +46,12 @@ type metaData struct {
4546
CompleteTime string
4647
}
4748

48-
const version = "0.4.0"
49+
const (
50+
// Version of this plugin for easy reference
51+
Version = "0.4.1"
52+
53+
defaultMaxAllowedPacket = 4194304
54+
)
4955

5056
// takes a *metaData
5157
const headerTmpl = `-- Go SQL Dump {{ .DumpVersion }}
@@ -97,10 +103,9 @@ DROP TABLE IF EXISTS {{ .NameEsc }};
97103
98104
LOCK TABLES {{ .NameEsc }} WRITE;
99105
/*!40000 ALTER TABLE {{ .NameEsc }} DISABLE KEYS */;
100-
{{- if .Next }}
101-
INSERT INTO {{ .NameEsc }} VALUES {{ .RowValues }}
102-
{{- range $value := .Stream }},{{ $value }}{{ end -}};
103-
{{- end }}
106+
{{ range $value := .Stream }}
107+
{{- $value }}
108+
{{ end -}}
104109
/*!40000 ALTER TABLE {{ .NameEsc }} ENABLE KEYS */;
105110
UNLOCK TABLES;
106111
`
@@ -110,7 +115,11 @@ const nullType = "NULL"
110115
// Dump data using struct
111116
func (data *Data) Dump() error {
112117
meta := metaData{
113-
DumpVersion: version,
118+
DumpVersion: Version,
119+
}
120+
121+
if data.MaxAllowedPacket == 0 {
122+
data.MaxAllowedPacket = defaultMaxAllowedPacket
114123
}
115124

116125
if err := meta.updateServerVersion(data.Connection); err != nil {
@@ -219,7 +228,7 @@ func (data *Data) isIgnoredTable(name string) bool {
219228

220229
func (data *metaData) updateServerVersion(db *sql.DB) (err error) {
221230
var serverVersion sql.NullString
222-
err = db.QueryRow("SELECT version()").Scan(&serverVersion)
231+
err = db.QueryRow("SELECT version();").Scan(&serverVersion)
223232
data.ServerVersion = serverVersion.String
224233
return
225234
}
@@ -323,45 +332,73 @@ func (table *table) Next() bool {
323332
}
324333

325334
func (table *table) RowValues() string {
326-
dataStrings := make([]string, len(table.values))
335+
return table.RowBuffer().String()
336+
}
337+
338+
func (table *table) RowBuffer() *bytes.Buffer {
339+
var b bytes.Buffer
340+
b.WriteString("(")
327341

328342
for key, value := range table.values {
343+
if key != 0 {
344+
b.WriteString(",")
345+
}
329346
switch s := value.(type) {
330347
case nil:
331-
dataStrings[key] = nullType
348+
b.WriteString(nullType)
332349
case *sql.NullString:
333350
if s.Valid {
334-
dataStrings[key] = "'" + sanitize(s.String) + "'"
351+
fmt.Fprintf(&b, "'%s'", sanitize(s.String))
335352
} else {
336-
dataStrings[key] = nullType
353+
b.WriteString(nullType)
337354
}
338355
case *sql.NullInt64:
339356
if s.Valid {
340-
dataStrings[key] = fmt.Sprintf("%d", s.Int64)
357+
fmt.Fprintf(&b, "%d", s.Int64)
341358
} else {
342-
dataStrings[key] = nullType
359+
b.WriteString(nullType)
343360
}
344361
case *sql.RawBytes:
345362
if len(*s) == 0 {
346-
dataStrings[key] = nullType
363+
b.WriteString(nullType)
347364
} else {
348-
dataStrings[key] = "_binary '" + sanitize(string(*s)) + "'"
365+
fmt.Fprintf(&b, "_binary '%s'", sanitize(string(*s)))
349366
}
350367
default:
351-
dataStrings[key] = fmt.Sprint("'", value, "'")
368+
fmt.Fprintf(&b, "'%s'", value)
352369
}
353370
}
371+
b.WriteString(")")
354372

355-
return "(" + strings.Join(dataStrings, ",") + ")"
373+
return &b
356374
}
357375

358376
func (table *table) Stream() <-chan string {
359377
valueOut := make(chan string, 1)
360-
go func(out chan string) {
361-
defer close(out)
378+
go func() {
379+
defer close(valueOut)
380+
var insert bytes.Buffer
381+
362382
for table.Next() {
363-
out <- table.RowValues()
383+
b := table.RowBuffer()
384+
// Truncate our insert if it won't fit
385+
if insert.Len() != 0 && insert.Len()+b.Len() > table.data.MaxAllowedPacket-1 {
386+
insert.WriteString(";")
387+
valueOut <- insert.String()
388+
insert.Reset()
389+
}
390+
391+
if insert.Len() == 0 {
392+
fmt.Fprintf(&insert, "INSERT INTO %s VALUES ", table.NameEsc())
393+
} else {
394+
insert.WriteString(",")
395+
}
396+
b.WriteTo(&insert)
397+
}
398+
if insert.Len() != 0 {
399+
insert.WriteString(";")
400+
valueOut <- insert.String()
364401
}
365-
}(valueOut)
402+
}()
366403
return valueOut
367404
}

‎dump_test.go

Lines changed: 115 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,59 @@ func TestCreateTableRowValues(t *testing.T) {
165165
assert.EqualValues(t, "('1','[email protected]','Test Name 1')", result)
166166
}
167167

168+
func TestCreateTableValuesSteam(t *testing.T) {
169+
db, mock, err := sqlmock.New()
170+
assert.NoError(t, err, "an error was not expected when opening a stub database connection")
171+
defer db.Close()
172+
173+
rows := sqlmock.NewRows([]string{"id", "email", "name"}).
174+
AddRow(1, "[email protected]", "Test Name 1").
175+
AddRow(2, "[email protected]", "Test Name 2")
176+
177+
mock.ExpectQuery("^SELECT (.+) FROM `test`$").WillReturnRows(rows)
178+
179+
data := Data{
180+
Connection: db,
181+
MaxAllowedPacket: 4096,
182+
}
183+
184+
table, err := data.createTable("test")
185+
assert.NoError(t, err)
186+
187+
s := table.Stream()
188+
assert.EqualValues(t, "INSERT INTO `test` VALUES ('1','[email protected]','Test Name 1'),('2','[email protected]','Test Name 2');", <-s)
189+
190+
// we make sure that all expectations were met
191+
assert.NoError(t, mock.ExpectationsWereMet(), "there were unfulfilled expections")
192+
}
193+
194+
func TestCreateTableValuesSteamSmallPackets(t *testing.T) {
195+
db, mock, err := sqlmock.New()
196+
assert.NoError(t, err, "an error was not expected when opening a stub database connection")
197+
defer db.Close()
198+
199+
rows := sqlmock.NewRows([]string{"id", "email", "name"}).
200+
AddRow(1, "[email protected]", "Test Name 1").
201+
AddRow(2, "[email protected]", "Test Name 2")
202+
203+
mock.ExpectQuery("^SELECT (.+) FROM `test`$").WillReturnRows(rows)
204+
205+
data := Data{
206+
Connection: db,
207+
MaxAllowedPacket: 64,
208+
}
209+
210+
table, err := data.createTable("test")
211+
assert.NoError(t, err)
212+
213+
s := table.Stream()
214+
assert.EqualValues(t, "INSERT INTO `test` VALUES ('1','[email protected]','Test Name 1');", <-s)
215+
assert.EqualValues(t, "INSERT INTO `test` VALUES ('2','[email protected]','Test Name 2');", <-s)
216+
217+
// we make sure that all expectations were met
218+
assert.NoError(t, mock.ExpectationsWereMet(), "there were unfulfilled expections")
219+
}
220+
168221
func TestCreateTableAllValuesWithNil(t *testing.T) {
169222
db, mock, err := sqlmock.New()
170223
assert.NoError(t, err, "an error was not expected when opening a stub database connection")
@@ -217,8 +270,9 @@ func TestCreateTableOk(t *testing.T) {
217270

218271
var buf bytes.Buffer
219272
data := Data{
220-
Connection: db,
221-
Out: &buf,
273+
Connection: db,
274+
Out: &buf,
275+
MaxAllowedPacket: 4096,
222276
}
223277

224278
assert.NoError(t, data.getTemplates())
@@ -255,3 +309,62 @@ UNLOCK TABLES;
255309
result := strings.Replace(buf.String(), "`", "~", -1)
256310
assert.Equal(t, expectedResult, result)
257311
}
312+
313+
func TestCreateTableOkSmallPackets(t *testing.T) {
314+
db, mock, err := sqlmock.New()
315+
assert.NoError(t, err, "an error was not expected when opening a stub database connection")
316+
317+
defer db.Close()
318+
319+
createTableRows := sqlmock.NewRows([]string{"Table", "Create Table"}).
320+
AddRow("Test_Table", "CREATE TABLE 'Test_Table' (`id` int(11) NOT NULL AUTO_INCREMENT,`s` char(60) DEFAULT NULL, PRIMARY KEY (`id`))ENGINE=InnoDB DEFAULT CHARSET=latin1")
321+
322+
createTableValueRows := sqlmock.NewRows([]string{"id", "email", "name"}).
323+
AddRow(1, nil, "Test Name 1").
324+
AddRow(2, "[email protected]", "Test Name 2")
325+
326+
mock.ExpectQuery("^SHOW CREATE TABLE `Test_Table`$").WillReturnRows(createTableRows)
327+
mock.ExpectQuery("^SELECT (.+) FROM `Test_Table`$").WillReturnRows(createTableValueRows)
328+
329+
var buf bytes.Buffer
330+
data := Data{
331+
Connection: db,
332+
Out: &buf,
333+
MaxAllowedPacket: 64,
334+
}
335+
336+
assert.NoError(t, data.getTemplates())
337+
338+
table, err := data.createTable("Test_Table")
339+
assert.NoError(t, err)
340+
341+
data.writeTable(table)
342+
343+
// we make sure that all expectations were met
344+
assert.NoError(t, mock.ExpectationsWereMet(), "there were unfulfilled expections")
345+
346+
expectedResult := `
347+
--
348+
-- Table structure for table ~Test_Table~
349+
--
350+
351+
DROP TABLE IF EXISTS ~Test_Table~;
352+
/*!40101 SET @saved_cs_client = @@character_set_client */;
353+
SET character_set_client = utf8mb4 ;
354+
CREATE TABLE 'Test_Table' (~id~ int(11) NOT NULL AUTO_INCREMENT,~s~ char(60) DEFAULT NULL, PRIMARY KEY (~id~))ENGINE=InnoDB DEFAULT CHARSET=latin1;
355+
/*!40101 SET character_set_client = @saved_cs_client */;
356+
357+
--
358+
-- Dumping data for table ~Test_Table~
359+
--
360+
361+
LOCK TABLES ~Test_Table~ WRITE;
362+
/*!40000 ALTER TABLE ~Test_Table~ DISABLE KEYS */;
363+
INSERT INTO ~Test_Table~ VALUES ('1',NULL,'Test Name 1');
364+
INSERT INTO ~Test_Table~ VALUES ('2','[email protected]','Test Name 2');
365+
/*!40000 ALTER TABLE ~Test_Table~ ENABLE KEYS */;
366+
UNLOCK TABLES;
367+
`
368+
result := strings.Replace(buf.String(), "`", "~", -1)
369+
assert.Equal(t, expectedResult, result)
370+
}

‎mysqldump_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
"github.com/stretchr/testify/assert"
1010
)
1111

12-
const expected = `-- Go SQL Dump ` + version + `
12+
const expected = `-- Go SQL Dump ` + Version + `
1313
--
1414
-- ------------------------------------------------------
1515
-- Server version test_version

0 commit comments

Comments
 (0)