Skip to content

Commit 0504c16

Browse files
mscoutermarshmscoutermarsh
authored
Merge pull request #1155 from planetscale/issue-1462
2 parents 3287a73 + f459411 commit 0504c16

File tree

2 files changed

+240
-20
lines changed

2 files changed

+240
-20
lines changed

‎internal/dumper/loader.go‎

Lines changed: 56 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -304,32 +304,68 @@ func (l *Loader) restoreTableSchema(overwrite bool, tables []string, conn *Conne
304304
return err
305305
}
306306

307+
// Drop table once before processing queries (if overwrite is enabled)
308+
if overwrite {
309+
l.log.Info(
310+
"drop(overwrite.is.true)",
311+
zap.String("database", db),
312+
zap.String("table ", tbl),
313+
)
314+
315+
if l.cfg.ShowDetails {
316+
l.cfg.Printer.Println("Dropping Existing Table (if it exists): " + printer.BoldBlue(name))
317+
}
318+
dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s", name)
319+
err = conn.Execute(dropQuery)
320+
if err != nil {
321+
return err
322+
}
323+
}
324+
325+
// Execute each valid SQL statement (skip comments)
307326
for _, query := range queries {
308-
if !strings.HasPrefix(query, "/*") && query != "" {
309-
if overwrite {
310-
l.log.Info(
311-
"drop(overwrite.is.true)",
312-
zap.String("database", db),
313-
zap.String("table ", tbl),
314-
)
327+
// Skip empty queries and block comments
328+
trimmedQuery := strings.TrimSpace(query)
329+
if trimmedQuery == "" || strings.HasPrefix(trimmedQuery, "/*") {
330+
continue
331+
}
315332

316-
if l.cfg.ShowDetails {
317-
l.cfg.Printer.Println("Dropping Existing Table (if it exists): " + printer.BoldBlue(name))
318-
}
319-
dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s", name)
320-
err = conn.Execute(dropQuery)
321-
if err != nil {
322-
return err
323-
}
333+
// Filter out line comments (--) but keep the rest of the query
334+
var cleanedLines []string
335+
for _, line := range strings.Split(query, "\n") {
336+
trimmedLine := strings.TrimSpace(line)
337+
// Skip empty lines and comment-only lines
338+
if trimmedLine != "" && !strings.HasPrefix(trimmedLine, "--") {
339+
cleanedLines = append(cleanedLines, line)
324340
}
341+
}
325342

326-
if l.cfg.ShowDetails {
343+
// Skip if no non-comment content remains
344+
if len(cleanedLines) == 0 {
345+
continue
346+
}
347+
348+
// Reconstruct the query without comment lines
349+
cleanedQuery := strings.Join(cleanedLines, "\n")
350+
trimmedCleanedQuery := strings.TrimSpace(cleanedQuery)
351+
352+
if l.cfg.ShowDetails {
353+
// Detect query type and provide appropriate output
354+
upperQuery := strings.ToUpper(trimmedCleanedQuery)
355+
if strings.HasPrefix(upperQuery, "CREATE TABLE") {
327356
l.cfg.Printer.Printf("Creating Table: %s (Table %d of %d)\n", printer.BoldBlue(name), (idx + 1), numberOfTables)
357+
} else if strings.HasPrefix(upperQuery, "ALTER TABLE") {
358+
l.cfg.Printer.Printf("Altering Table: %s (Table %d of %d)\n", printer.BoldBlue(name), (idx + 1), numberOfTables)
359+
l.cfg.Printer.Printf("Query: %s\n", cleanedQuery)
360+
} else {
361+
// For any other query type, show what's being executed
362+
l.cfg.Printer.Printf("Executing Query for Table: %s (Table %d of %d)\n", printer.BoldBlue(name), (idx + 1), numberOfTables)
363+
l.cfg.Printer.Printf("Query: %s\n", cleanedQuery)
328364
}
329-
err = conn.Execute(query)
330-
if err != nil {
331-
return err
332-
}
365+
}
366+
err = conn.Execute(cleanedQuery)
367+
if err != nil {
368+
return err
333369
}
334370
}
335371
l.log.Info("restoring schema",

‎internal/dumper/loader_test.go‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package dumper
22

33
import (
44
"context"
5+
"os"
56
"testing"
67

78
qt "github.com/frankban/quicktest"
@@ -47,3 +48,186 @@ func TestLoader(t *testing.T) {
4748
err = loader.Run(context.Background())
4849
c.Assert(err, qt.IsNil)
4950
}
51+
52+
func TestRestoreTableSchema_WithComments(t *testing.T) {
53+
tests := []struct {
54+
name string
55+
schemaContent string
56+
setupQueries []string
57+
description string
58+
}{
59+
{
60+
name: "schema with line comments at end",
61+
schemaContent: `CREATE TABLE example_table (
62+
id INT PRIMARY KEY
63+
);
64+
-- This is a comment
65+
-- This is another comment`,
66+
setupQueries: []string{
67+
"CREATE TABLE example_table (\n id INT PRIMARY KEY\n)",
68+
},
69+
description: "Should execute CREATE TABLE and skip trailing comments",
70+
},
71+
{
72+
name: "schema with ALTER TABLE after comments",
73+
schemaContent: `CREATE TABLE example_table (
74+
id INT PRIMARY KEY,
75+
name VARCHAR(100)
76+
);
77+
-- This is a comment
78+
-- This is another comment
79+
ALTER TABLE example_table
80+
ADD INDEX idx_name (name);`,
81+
setupQueries: []string{
82+
"CREATE TABLE example_table (\n id INT PRIMARY KEY,\n name VARCHAR(100)\n)",
83+
},
84+
description: "Should execute CREATE TABLE and ALTER TABLE, skipping comments in between",
85+
},
86+
{
87+
name: "schema with block comments",
88+
schemaContent: `/* This is a block comment */
89+
CREATE TABLE example_table (
90+
id INT PRIMARY KEY
91+
);`,
92+
setupQueries: []string{
93+
"CREATE TABLE example_table (\n id INT PRIMARY KEY\n)",
94+
},
95+
description: "Should skip block comments and execute CREATE TABLE",
96+
},
97+
{
98+
name: "schema with interspersed comments",
99+
schemaContent: `CREATE TABLE example_table (
100+
id INT PRIMARY KEY
101+
);
102+
-- Comment between statements
103+
ALTER TABLE example_table ADD COLUMN name VARCHAR(100);
104+
-- Another comment
105+
ALTER TABLE example_table ADD INDEX idx_id (id);`,
106+
setupQueries: []string{
107+
"CREATE TABLE example_table (\n id INT PRIMARY KEY\n)",
108+
},
109+
description: "Should execute all SQL statements and skip all comment lines",
110+
},
111+
}
112+
113+
for _, tt := range tests {
114+
t.Run(tt.name, func(t *testing.T) {
115+
c := qt.New(t)
116+
117+
log := xlog.NewStdLog(xlog.Level(xlog.ERROR))
118+
fakedbs := driver.NewTestHandler(log)
119+
server, err := driver.MockMysqlServer(log, fakedbs)
120+
c.Assert(err, qt.IsNil)
121+
defer server.Close()
122+
123+
address := server.Addr()
124+
125+
// Set up mock expectations
126+
fakedbs.AddQueryPattern("use .*", &sqltypes.Result{})
127+
fakedbs.AddQueryPattern("set foreign_key_checks=.*", &sqltypes.Result{})
128+
fakedbs.AddQueryPattern("drop table .*", &sqltypes.Result{})
129+
fakedbs.AddQueryPattern("alter table .*", &sqltypes.Result{})
130+
131+
// Add expected queries - if these aren't executed, the test will fail
132+
for _, query := range tt.setupQueries {
133+
fakedbs.AddQuery(query, &sqltypes.Result{})
134+
}
135+
136+
// Create test schema file
137+
tempDir := c.TempDir()
138+
schemaFile := tempDir + "/testdb.test_table-schema.sql"
139+
err = os.WriteFile(schemaFile, []byte(tt.schemaContent), 0644)
140+
c.Assert(err, qt.IsNil)
141+
142+
// Create loader
143+
cfg := &Config{
144+
Database: "testdb",
145+
Outdir: tempDir,
146+
User: "mock",
147+
Password: "mock",
148+
Threads: 1,
149+
Address: address,
150+
IntervalMs: 500,
151+
OverwriteTables: true,
152+
ShowDetails: false,
153+
Debug: false,
154+
}
155+
loader, err := NewLoader(cfg)
156+
c.Assert(err, qt.IsNil)
157+
158+
// Create connection pool
159+
pool, err := NewPool(loader.log, cfg.Threads, cfg.Address, cfg.User, cfg.Password, cfg.SessionVars, "")
160+
c.Assert(err, qt.IsNil)
161+
defer pool.Close()
162+
163+
conn := pool.Get()
164+
defer pool.Put(conn)
165+
166+
// Execute restoreTableSchema - should not return error if all expected queries are executed
167+
err = loader.restoreTableSchema(cfg.OverwriteTables, []string{schemaFile}, conn)
168+
c.Assert(err, qt.IsNil, qt.Commentf("%s: failed to restore table schema", tt.description))
169+
})
170+
}
171+
}
172+
173+
func TestRestoreTableSchema_DropTableCalledOnce(t *testing.T) {
174+
c := qt.New(t)
175+
176+
log := xlog.NewStdLog(xlog.Level(xlog.ERROR))
177+
fakedbs := driver.NewTestHandler(log)
178+
server, err := driver.MockMysqlServer(log, fakedbs)
179+
c.Assert(err, qt.IsNil)
180+
defer server.Close()
181+
182+
address := server.Addr()
183+
184+
// Set up mock expectations
185+
fakedbs.AddQueryPattern("use .*", &sqltypes.Result{})
186+
fakedbs.AddQueryPattern("set foreign_key_checks=.*", &sqltypes.Result{})
187+
188+
// Add DROP TABLE query only once - if it's called twice, the second call will fail
189+
// because there's no matching handler for it
190+
fakedbs.AddQuery("DROP TABLE IF EXISTS `testdb`.`test_table`", &sqltypes.Result{})
191+
fakedbs.AddQuery("CREATE TABLE test_table (\n id INT PRIMARY KEY\n)", &sqltypes.Result{})
192+
193+
// Create test schema file with comments at the end (the original bug scenario)
194+
tempDir := c.TempDir()
195+
schemaFile := tempDir + "/testdb.test_table-schema.sql"
196+
schemaContent := `CREATE TABLE test_table (
197+
id INT PRIMARY KEY
198+
);
199+
-- This is a comment
200+
-- This is another comment`
201+
err = os.WriteFile(schemaFile, []byte(schemaContent), 0644)
202+
c.Assert(err, qt.IsNil)
203+
204+
// Create loader
205+
cfg := &Config{
206+
Database: "testdb",
207+
Outdir: tempDir,
208+
User: "mock",
209+
Password: "mock",
210+
Threads: 1,
211+
Address: address,
212+
IntervalMs: 500,
213+
OverwriteTables: true,
214+
ShowDetails: false,
215+
Debug: false,
216+
}
217+
loader, err := NewLoader(cfg)
218+
c.Assert(err, qt.IsNil)
219+
220+
// Create connection pool
221+
pool, err := NewPool(loader.log, cfg.Threads, cfg.Address, cfg.User, cfg.Password, cfg.SessionVars, "")
222+
c.Assert(err, qt.IsNil)
223+
defer pool.Close()
224+
225+
conn := pool.Get()
226+
defer pool.Put(conn)
227+
228+
// Execute restoreTableSchema
229+
// If DROP TABLE is called more than once, the test will fail because there's
230+
// only one handler registered for it
231+
err = loader.restoreTableSchema(cfg.OverwriteTables, []string{schemaFile}, conn)
232+
c.Assert(err, qt.IsNil, qt.Commentf("DROP TABLE should be called exactly once. If called multiple times, this test will fail."))
233+
}

0 commit comments

Comments
 (0)