Skip to content

Commit bb20573

Browse files
nickvanwnickvanw
authored
Merge pull request #1273 from planetscale/improve-dumper-escaping
Improve escaping
2 parents 70c8ccc + 2c3ae87 commit bb20573

7 files changed

Lines changed: 307 additions & 24 deletions

File tree

‎internal/cmd/database/dump.go‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -254,13 +254,13 @@ func dump(ch *cmdutil.Helper, cmd *cobra.Command, flags *dumpFlags, args []strin
254254

255255
if flags.shard != "" {
256256
if flags.replica {
257-
useCmd := fmt.Sprintf("USE `%s/%s@replica`;", dbName, flags.shard)
257+
useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
258258
cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
259259
} else if flags.rdonly {
260-
useCmd := fmt.Sprintf("USE `%s/%s@rdonly`;", dbName, flags.shard)
260+
useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
261261
cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
262262
} else {
263-
useCmd := fmt.Sprintf("USE `%s/%s`;", dbName, flags.shard)
263+
useCmd := shardUseCommand(dbName, flags.shard, flags.replica, flags.rdonly)
264264
cfg.SessionVars = append([]string{useCmd}, cfg.SessionVars...)
265265
}
266266
}
@@ -392,3 +392,17 @@ func parseColumnIncludes(columns []string) (map[string]map[string]bool, error) {
392392

393393
return result, nil
394394
}
395+
396+
func shardUseCommand(dbName string, shard string, replica bool, rdonly bool) string {
397+
target := fmt.Sprintf("%s/%s", dbName, shard)
398+
if replica {
399+
target += "@replica"
400+
} else if rdonly {
401+
target += "@rdonly"
402+
}
403+
return fmt.Sprintf("USE %s;", quoteIdentifier(target))
404+
}
405+
406+
func quoteIdentifier(identifier string) string {
407+
return "`" + strings.ReplaceAll(identifier, "`", "``") + "`"
408+
}

‎internal/cmd/database/dump_test.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,3 +100,12 @@ func TestParseColumnIncludes(t *testing.T) {
100100
})
101101
}
102102
}
103+
104+
func TestShardUseCommand(t *testing.T) {
105+
c := qt.New(t)
106+
107+
c.Assert(shardUseCommand("commerce", "-80", false, false), qt.Equals, "USE `commerce/-80`;")
108+
c.Assert(shardUseCommand("commerce", "-80", true, false), qt.Equals, "USE `commerce/-80@replica`;")
109+
c.Assert(shardUseCommand("commerce", "-80", false, true), qt.Equals, "USE `commerce/-80@rdonly`;")
110+
c.Assert(shardUseCommand("key`space", "sh`ard", false, false), qt.Equals, "USE `key``space/sh``ard`;")
111+
}

‎internal/dumper/dumper.go‎

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,7 @@ func writeMetaData(outdir string) error {
249249
}
250250

251251
func (d *Dumper) dumpTableSchema(conn *Connection, database string, table string, views map[string]bool) error {
252-
qr, err := conn.Fetch(fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s`", database, table))
252+
qr, err := conn.Fetch(fmt.Sprintf("SHOW CREATE TABLE %s.%s", quoteIdentifier(database), quoteIdentifier(table)))
253253
if err != nil {
254254
return err
255255
}
@@ -293,7 +293,7 @@ func (d *Dumper) dumpTable(ctx context.Context, conn *Connection, database strin
293293
return err
294294
}
295295

296-
cursor, err := conn.StreamFetch(fmt.Sprintf("SELECT %s FROM `%s`.`%s` %s", strings.Join(dumpCtx.selfields, ", "), database, table, dumpCtx.where))
296+
cursor, err := conn.StreamFetch(fmt.Sprintf("SELECT %s FROM %s.%s %s", strings.Join(dumpCtx.selfields, ", "), quoteIdentifier(database), quoteIdentifier(table), dumpCtx.where))
297297
if err != nil {
298298
return err
299299
}
@@ -387,9 +387,9 @@ func (d *Dumper) tableDumpContext(conn *Connection, table string) (*dumpContext,
387387
ctx.fieldNames = append(ctx.fieldNames, name)
388388
replacement, ok := d.cfg.Selects[table][name]
389389
if ok {
390-
ctx.selfields = append(ctx.selfields, fmt.Sprintf("%s AS `%s`", replacement, name))
390+
ctx.selfields = append(ctx.selfields, fmt.Sprintf("%s AS %s", replacement, quoteIdentifier(name)))
391391
} else {
392-
ctx.selfields = append(ctx.selfields, fmt.Sprintf("`%s`", name))
392+
ctx.selfields = append(ctx.selfields, quoteIdentifier(name))
393393
}
394394
}
395395

@@ -406,7 +406,7 @@ func (d *Dumper) tableDumpContext(conn *Connection, table string) (*dumpContext,
406406
}
407407

408408
func (d *Dumper) allTables(conn *Connection, database string) ([]string, error) {
409-
qr, err := conn.Fetch(fmt.Sprintf("SHOW TABLES FROM `%s`", database))
409+
qr, err := conn.Fetch(fmt.Sprintf("SHOW TABLES FROM %s", quoteIdentifier(database)))
410410
if err != nil {
411411
return nil, err
412412
}
@@ -419,12 +419,12 @@ func (d *Dumper) allTables(conn *Connection, database string) ([]string, error)
419419
}
420420

421421
func (d *Dumper) allViews(conn *Connection, database string) (map[string]bool, error) {
422-
query := `SELECT TABLE_NAME
423-
FROM information_schema.TABLES
424-
WHERE TABLE_SCHEMA LIKE '%s'
425-
AND TABLE_TYPE = 'VIEW'
426-
`
427-
qr, err := conn.Fetch(fmt.Sprintf(query, database))
422+
query := "SELECT TABLE_NAME \n" +
423+
"\t\t\t FROM information_schema.TABLES \n" +
424+
"\t\t\t WHERE TABLE_SCHEMA LIKE %s \n" +
425+
"\t\t\t AND TABLE_TYPE = 'VIEW'\n" +
426+
"\t\t\t"
427+
qr, err := conn.Fetch(fmt.Sprintf(query, quoteStringLiteral(database)))
428428
if err != nil {
429429
return nil, err
430430
}
@@ -466,7 +466,7 @@ func (d *Dumper) filterDatabases(conn *Connection, filter *regexp.Regexp, invert
466466

467467
// dumpableFieldNames returns a slice that contains valid field names for the dump.
468468
func (d *Dumper) dumpableFieldNames(conn *Connection, table string) ([]string, error) {
469-
qr, err := conn.Fetch(fmt.Sprintf("SHOW FIELDS FROM `%s`", table))
469+
qr, err := conn.Fetch(fmt.Sprintf("SHOW FIELDS FROM %s", quoteIdentifier(table)))
470470
if err != nil {
471471
return nil, err
472472
}
@@ -515,6 +515,14 @@ func writeFile(file string, data string) error {
515515
return nil
516516
}
517517

518+
func quoteIdentifier(identifier string) string {
519+
return "`" + strings.ReplaceAll(identifier, "`", "``") + "`"
520+
}
521+
522+
func quoteStringLiteral(s string) string {
523+
return "'" + strings.ReplaceAll(s, "'", "''") + "'"
524+
}
525+
518526
// escapeBytes used to escape the literal byte.
519527
// See https://dev.mysql.com/doc/refman/5.7/en/string-literals.html
520528
// for more information on how to escape string literals in MySQL.

‎internal/dumper/dumper_test.go‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1574,6 +1574,124 @@ func TestEscapeBytes(t *testing.T) {
15741574
}
15751575
}
15761576

1577+
func TestQuoteIdentifier(t *testing.T) {
1578+
c := qt.New(t)
1579+
1580+
c.Assert(quoteIdentifier("simple"), qt.Equals, "`simple`")
1581+
c.Assert(quoteIdentifier("customer_report`$probe"), qt.Equals, "`customer_report``$probe`")
1582+
c.Assert(quoteIdentifier("display`name"), qt.Equals, "`display``name`")
1583+
}
1584+
1585+
func TestQuoteStringLiteral(t *testing.T) {
1586+
c := qt.New(t)
1587+
1588+
c.Assert(quoteStringLiteral("simple"), qt.Equals, "'simple'")
1589+
c.Assert(quoteStringLiteral("test'db"), qt.Equals, "'test''db'")
1590+
}
1591+
1592+
func TestDumperEscapesDiscoveredIdentifiers(t *testing.T) {
1593+
c := qt.New(t)
1594+
1595+
log := xlog.NewStdLog(xlog.Level(xlog.INFO))
1596+
fakedbs := driver.NewTestHandler(log)
1597+
server, err := driver.MockMysqlServer(log, fakedbs)
1598+
c.Assert(err, qt.IsNil)
1599+
c.Cleanup(func() { server.Close() })
1600+
1601+
address := server.Addr()
1602+
database := "test`db"
1603+
table := "customer_report`$probe"
1604+
column := "display`name"
1605+
1606+
selectResult := &sqltypes.Result{
1607+
Fields: []*querypb.Field{
1608+
{Name: column, Type: querypb.Type_VARCHAR},
1609+
},
1610+
Rows: [][]sqltypes.Value{
1611+
{sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte("ok"))},
1612+
},
1613+
}
1614+
1615+
schemaResult := &sqltypes.Result{
1616+
Fields: []*querypb.Field{
1617+
{Name: "Table", Type: querypb.Type_VARCHAR},
1618+
{Name: "Create Table", Type: querypb.Type_VARCHAR},
1619+
},
1620+
Rows: [][]sqltypes.Value{
1621+
{
1622+
sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte(table)),
1623+
sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte("CREATE TABLE `customer_report``$probe` (`display``name` varchar(255)) ENGINE=InnoDB")),
1624+
},
1625+
},
1626+
}
1627+
1628+
tablesResult := &sqltypes.Result{
1629+
Fields: []*querypb.Field{
1630+
{Name: "Tables_in_test`db", Type: querypb.Type_VARCHAR},
1631+
},
1632+
Rows: [][]sqltypes.Value{
1633+
{sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte(table))},
1634+
},
1635+
}
1636+
1637+
viewsResult := &sqltypes.Result{
1638+
Fields: []*querypb.Field{
1639+
{Name: "TABLE_NAME", Type: querypb.Type_VARCHAR},
1640+
},
1641+
Rows: [][]sqltypes.Value{},
1642+
}
1643+
1644+
fieldsResult := &sqltypes.Result{
1645+
Fields: []*querypb.Field{
1646+
{Name: "Field", Type: querypb.Type_VARCHAR},
1647+
{Name: "Type", Type: querypb.Type_VARCHAR},
1648+
{Name: "Null", Type: querypb.Type_VARCHAR},
1649+
{Name: "Key", Type: querypb.Type_VARCHAR},
1650+
{Name: "Default", Type: querypb.Type_VARCHAR},
1651+
{Name: "Extra", Type: querypb.Type_VARCHAR},
1652+
},
1653+
Rows: [][]sqltypes.Value{
1654+
testRow(column, ""),
1655+
},
1656+
}
1657+
1658+
fakedbs.AddQueryPattern("use .*", &sqltypes.Result{})
1659+
fakedbs.AddQueryPattern("show create table .*", schemaResult)
1660+
fakedbs.AddQueryPattern("show tables from .*", tablesResult)
1661+
fakedbs.AddQueryPattern("select table_name .* from information_schema.tables .*", viewsResult)
1662+
fakedbs.AddQueryPattern("show fields from .*", fieldsResult)
1663+
fakedbs.AddQueryPattern("select .* from .*", selectResult)
1664+
fakedbs.AddQueryPattern("set .*", &sqltypes.Result{})
1665+
1666+
cfg := &Config{
1667+
Database: database,
1668+
Outdir: c.TempDir(),
1669+
User: "mock",
1670+
Password: "mock",
1671+
Address: address,
1672+
ChunksizeInMB: 1,
1673+
Threads: 16,
1674+
StmtSize: 10000,
1675+
IntervalMs: 500,
1676+
SessionVars: []string{"SET @@radon_streaming_fetch='ON'"},
1677+
}
1678+
1679+
d, err := NewDumper(cfg)
1680+
c.Assert(err, qt.IsNil)
1681+
1682+
err = d.Run(context.Background())
1683+
c.Assert(err, qt.IsNil)
1684+
1685+
c.Assert(fakedbs.GetQueryCalledNum("SHOW TABLES FROM `test``db`"), qt.Equals, 1)
1686+
c.Assert(fakedbs.GetQueryCalledNum("SHOW CREATE TABLE `test``db`.`customer_report``$probe`"), qt.Equals, 1)
1687+
c.Assert(fakedbs.GetQueryCalledNum("SHOW FIELDS FROM `customer_report``$probe`"), qt.Equals, 1)
1688+
c.Assert(fakedbs.GetQueryCalledNum("SELECT `display``name` FROM `test``db`.`customer_report``$probe` "), qt.Equals, 1)
1689+
1690+
dat, err := os.ReadFile(cfg.Outdir + "/" + database + "." + table + ".00001.sql")
1691+
c.Assert(err, qt.IsNil)
1692+
c.Assert(string(dat), qt.Contains, "INSERT INTO `customer_report``$probe`(`display``name`) VALUES")
1693+
}
1694+
15771695
func TestDumperColumnIncludes(t *testing.T) {
15781696
c := qt.New(t)
15791697

‎internal/dumper/loader.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -270,15 +270,15 @@ func (l *Loader) restoreTableSchema(overwrite bool, tables []string, conn *Conne
270270
name := strings.TrimSuffix(base, schemaSuffix)
271271
db := l.databaseNameFromFilename(name)
272272
tbl := strings.Split(name, ".")[1]
273-
name = fmt.Sprintf("`%v`.`%v`", db, tbl)
273+
name = fmt.Sprintf("%s.%s", quoteIdentifier(db), quoteIdentifier(tbl))
274274

275275
l.log.Info(
276276
"working table",
277277
zap.String("database", db),
278278
zap.String("table ", tbl),
279279
)
280280

281-
err := conn.Execute(fmt.Sprintf("USE `%s`", db))
281+
err := conn.Execute(fmt.Sprintf("USE %s", quoteIdentifier(db)))
282282
if err != nil {
283283
return err
284284
}
@@ -386,15 +386,15 @@ func (l *Loader) restoreViews(overwrite bool, views []string, conn *Connection)
386386
name := strings.TrimSuffix(base, viewSuffix)
387387
db := strings.Split(name, ".")[0]
388388
view := strings.Split(name, ".")[1]
389-
name = fmt.Sprintf("`%v`.`%v`", db, view)
389+
name = fmt.Sprintf("%s.%s", quoteIdentifier(db), quoteIdentifier(view))
390390

391391
l.log.Info(
392392
"working view",
393393
zap.String("database", db),
394394
zap.String("view ", view),
395395
)
396396

397-
err := conn.Execute(fmt.Sprintf("USE `%s`", db))
397+
err := conn.Execute(fmt.Sprintf("USE %s", quoteIdentifier(db)))
398398
if err != nil {
399399
return err
400400
}
@@ -483,7 +483,7 @@ func (l *Loader) restoreTable(ctx context.Context, table string, conn *Connectio
483483
zap.Int("thread_conn_id", conn.ID),
484484
)
485485

486-
err := conn.Execute(fmt.Sprintf("USE `%s`", db))
486+
err := conn.Execute(fmt.Sprintf("USE %s", quoteIdentifier(db)))
487487
if err != nil {
488488
return 0, err
489489
}

0 commit comments

Comments
 (0)