Repository: ariga/atlas
Pull request #1197 (2022-10-13 07:56:09)

sql/internal: check balanced quoting
Signed-off-by: Giau. Tran Minh <hello@giautm.dev>


Patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate.go b/cmd/atlas/internal/cmdapi/migrate.go
index 554aafd4..0c05ea8b 100644
--- a/cmd/atlas/internal/cmdapi/migrate.go
+++ b/cmd/atlas/internal/cmdapi/migrate.go
@@ -917,12 +917,12 @@ func migrateStatusRun(cmd *cobra.Command, _ []string, flags migrateStatusFlags)
 	}
 	var format = cmdmigrate.DefaultTemplate
 	if f := flags.logFormat; f != "" {
-		format, err = template.New("format").Funcs(cmdmigrate.TemplateFuncs).Parse(f)
+		format, err = template.New("format").Funcs(cmdmigrate.StatusTemplateFuncs).Parse(f)
 		if err != nil {
 			return fmt.Errorf("parse log format: %w", err)
 		}
 	}
-	return (&cmdmigrate.Reporter{
+	return (&cmdmigrate.StatusReporter{
 		Client:       client,
 		Dir:          dir,
 		ReportWriter: &cmdmigrate.TemplateWriter{T: format, W: cmd.OutOrStdout()},
diff --git a/cmd/atlas/internal/migrate/report.go b/cmd/atlas/internal/migrate/report.go
index c91f55eb..3cb4735a 100644
--- a/cmd/atlas/internal/migrate/report.go
+++ b/cmd/atlas/internal/migrate/report.go
@@ -22,94 +22,66 @@ import (
 	"github.com/olekukonko/tablewriter"
 )
 
-// Reporter is used to gather information about migration status.
-type Reporter struct {
-	// Client configures the connection to the database to file a StatusReport for.
-	Client *sqlclient.Client
-	// Dir is used for scanning and validating the migration directory.
-	Dir migrate.Dir
-	// ReportWriter writes the summary report.
-	ReportWriter ReportWriter
-	// Schema name the revision table resides in.
-	Schema string
-}
+var (
+	// ColorTemplateFuncs are globally available functions to color strings in a report template.
+	ColorTemplateFuncs = template.FuncMap{
+		"cyan":   color.CyanString,
+		"green":  color.HiGreenString,
+		"red":    color.HiRedString,
+		"yellow": color.YellowString,
+	}
+)
 
-// Status creates and writes a StatusReport.
-func (r *Reporter) Status(ctx context.Context) error {
-	rep, err := NewStatusReport(r.Client, r.Dir)
-	if err != nil {
-		return err
+type (
+	// A Report represents collected information about the execution of a 'atlas migrate' command execution.
+	Report interface {
+		report()
 	}
-	// Check if there already is a revision table in the defined schema.
-	// Inspect schema and check if the table does already exist.
-	sch, err := r.Client.InspectSchema(ctx, r.Schema, &schema.InspectOptions{Tables: []string{revision.Table}})
-	if err != nil && !schema.IsNotExistError(err) {
-		return err
+
+	// ReportWriter writes a Report.
+	ReportWriter interface {
+		WriteReport(Report) error
 	}
-	if schema.IsNotExistError(err) || func() bool { _, ok := sch.Table(revision.Table); return !ok }() {
-		// Either schema or table does not exist.
-		rep.Pending = rep.Available
-	} else {
-		// Both exist, fetch their data.
-		rrw, err := NewEntRevisions(ctx, r.Client, WithSchema(r.Schema))
-		if err != nil {
-			return err
-		}
-		if err := rrw.Migrate(ctx); err != nil {
-			return err
-		}
-		ex, err := migrate.NewExecutor(r.Client.Driver, r.Dir, rrw)
-		if err != nil {
-			return err
-		}
-		rep.Pending, err = ex.Pending(ctx)
-		if err != nil && !errors.Is(err, migrate.ErrNoPendingFiles) {
-			return err
-		}
-		rep.Applied, err = rrw.ReadRevisions(ctx)
-		if err != nil {
-			return err
-		}
+
+	// A TemplateWriter is ReportWriter that writes output according to a template.
+	TemplateWriter struct {
+		T *template.Template
+		W io.Writer
 	}
-	switch len(rep.Pending) {
-	case len(rep.Available):
-		rep.Current = "No migration applied yet"
-	default:
-		rep.Current = rep.Applied[len(rep.Applied)-1].Version
+
+	// Env holds the environment information.
+	Env struct {
+		Driver string         `json:"Driver,omitempty"` // Driver name.
+		URL    *sqlclient.URL `json:"URL,omitempty"`    // URL to dev database.
+		Dir    string         `json:"Dir,omitempty"`    // Path to migration directory.
 	}
-	if len(rep.Pending) == 0 {
-		rep.Status = "OK"
-		rep.Next = "Already at latest version"
-	} else {
-		rep.Status = "PENDING"
-		rep.Next = rep.Pending[0].Version()
+
+	// Files is a slice of migrate.File. Implements json.Marshaler.
+	Files []migrate.File
+)
+
+// WriteReport implements ReportWriter.
+func (w *TemplateWriter) WriteReport(r Report) error {
+	return w.T.Execute(w.W, r)
+}
+
+// MarshalJSON implements json.Marshaler.
+func (fs Files) MarshalJSON() ([]byte, error) {
+	type file struct {
+		Name        string `json:"Name,omitempty"`
+		Version     string `json:"Version,omitempty"`
+		Description string `json:"Description,omitempty"`
 	}
-	// If the last one is partially applied (and not manually resolved).
-	if len(rep.Applied) != 0 {
-		last := rep.Applied[len(rep.Applied)-1]
-		if !last.Type.Has(migrate.RevisionTypeResolved) && last.Applied < last.Total {
-			rep.SQL = strings.ReplaceAll(last.ErrorStmt, "\n", " ")
-			rep.Error = strings.ReplaceAll(last.Error, "\n", " ")
-			rep.Count = last.Applied
-			idx := migrate.FilesLastIndex(rep.Available, func(f migrate.File) bool {
-				return f.Version() == last.Version
-			})
-			if idx == -1 {
-				return fmt.Errorf("migration file with version %q not found", last.Version)
-			}
-			stmts, err := rep.Available[idx].Stmts()
-			if err != nil {
-				return err
-			}
-			rep.Total = len(stmts)
-		}
+	files := make([]file, len(fs))
+	for i, f := range fs {
+		files[i] = file{f.Name(), f.Version(), f.Desc()}
 	}
-	return r.ReportWriter.WriteReport(rep)
+	return json.Marshal(files)
 }
 
 var (
-	// TemplateFuncs are global functions available in templates.
-	TemplateFuncs = template.FuncMap{
+	// StatusTemplateFuncs are global functions available in templates.
+	StatusTemplateFuncs = merge(template.FuncMap{
 		"json": func(v any, args ...string) (string, error) {
 			var (
 				b   []byte
@@ -125,19 +97,10 @@ var (
 			}
 			return string(b), err
 		},
-		"table":  table,
-		"cyan":   color.CyanString,
-		"green":  color.HiGreenString,
-		"red":    color.HiRedString,
-		"yellow": color.YellowString,
+		"table": table,
 		"default": func(report *StatusReport) (string, error) {
 			var buf bytes.Buffer
-			t, err := template.New("report").Funcs(template.FuncMap{
-				"cyan":   color.CyanString,
-				"green":  color.HiGreenString,
-				"red":    color.HiRedString,
-				"yellow": color.YellowString,
-			}).Parse(`Migration Status:
+			t, err := template.New("report").Funcs(ColorTemplateFuncs).Parse(`Migration Status:
 {{- if eq .Status "OK"      }} {{ green .Status }}{{ end }}
 {{- if eq .Status "PENDING" }} {{ yellow .Status }}{{ end }}
   {{ yellow "--" }} Current Version: {{ cyan .Current }}
@@ -157,20 +120,26 @@ Last migration attempt had errors:
 			err = t.Execute(&buf, report)
 			return buf.String(), err
 		},
-	}
-	DefaultTemplate = template.Must(template.New("report").Funcs(TemplateFuncs).Parse("{{ default . }}"))
+	}, ColorTemplateFuncs)
+	DefaultTemplate = template.Must(template.New("report").Funcs(StatusTemplateFuncs).Parse("{{ default . }}"))
 )
 
 type (
-	// Env holds the environment information.
-	Env struct {
-		Driver string         `json:"Driver,omitempty"` // Driver name.
-		URL    *sqlclient.URL `json:"URL,omitempty"`    // URL to dev database.
-		Dir    string         `json:"Dir,omitempty"`    // Path to migration directory.
+	// StatusReporter is used to gather information about migration status.
+	StatusReporter struct {
+		// Client configures the connection to the database to file a StatusReport for.
+		Client *sqlclient.Client
+		// Dir is used for scanning and validating the migration directory.
+		Dir migrate.Dir
+		// ReportWriter writes the summary report.
+		ReportWriter ReportWriter
+		// Schema name the revision table resides in.
+		Schema string
 	}
 
 	// StatusReport contains a summary of the migration status of a database.
 	StatusReport struct {
+		Report
 		Env           `json:"Env"`
 		Available     Files               `json:"Available,omitempty"` // Available migration files.
 		Pending       Files               `json:"Pending,omitempty"`   // Pending migration files.
@@ -181,20 +150,6 @@ type (
 		Error         string              `json:"Error,omitempty"` // Last Error that occurred.
 		SQL           string              `json:"SQL,omitempty"`   // SQL that caused the last Error.
 	}
-
-	// ReportWriter writes a StatusReport.
-	ReportWriter interface {
-		WriteReport(*StatusReport) error
-	}
-
-	// A TemplateWriter is ReportWrite that writes output according to a template.
-	TemplateWriter struct {
-		T *template.Template
-		W io.Writer
-	}
-
-	// Files is a slice of migrate.File. Implements json.Marshaler.
-	Files []migrate.File
 )
 
 // NewStatusReport returns a new StatusReport.
@@ -216,23 +171,77 @@ func NewStatusReport(c *sqlclient.Client, dir migrate.Dir) (*StatusReport, error
 	return sum, nil
 }
 
-// WriteReport implements ReportWriter.
-func (w *TemplateWriter) WriteReport(r *StatusReport) error {
-	return w.T.Execute(w.W, r)
-}
-
-// MarshalJSON implements json.Marshaler.
-func (fs Files) MarshalJSON() ([]byte, error) {
-	type file struct {
-		Name        string `json:"Name,omitempty"`
-		Version     string `json:"Version,omitempty"`
-		Description string `json:"Description,omitempty"`
+// Status creates and writes a StatusReport.
+func (r *StatusReporter) Status(ctx context.Context) error {
+	rep, err := NewStatusReport(r.Client, r.Dir)
+	if err != nil {
+		return err
 	}
-	files := make([]file, len(fs))
-	for i, f := range fs {
-		files[i] = file{f.Name(), f.Version(), f.Desc()}
+	// Check if there already is a revision table in the defined schema.
+	// Inspect schema and check if the table does already exist.
+	sch, err := r.Client.InspectSchema(ctx, r.Schema, &schema.InspectOptions{Tables: []string{revision.Table}})
+	if err != nil && !schema.IsNotExistError(err) {
+		return err
 	}
-	return json.Marshal(files)
+	if schema.IsNotExistError(err) || func() bool { _, ok := sch.Table(revision.Table); return !ok }() {
+		// Either schema or table does not exist.
+		rep.Pending = rep.Available
+	} else {
+		// Both exist, fetch their data.
+		rrw, err := NewEntRevisions(ctx, r.Client, WithSchema(r.Schema))
+		if err != nil {
+			return err
+		}
+		if err := rrw.Migrate(ctx); err != nil {
+			return err
+		}
+		ex, err := migrate.NewExecutor(r.Client.Driver, r.Dir, rrw)
+		if err != nil {
+			return err
+		}
+		rep.Pending, err = ex.Pending(ctx)
+		if err != nil && !errors.Is(err, migrate.ErrNoPendingFiles) {
+			return err
+		}
+		rep.Applied, err = rrw.ReadRevisions(ctx)
+		if err != nil {
+			return err
+		}
+	}
+	switch len(rep.Pending) {
+	case len(rep.Available):
+		rep.Current = "No migration applied yet"
+	default:
+		rep.Current = rep.Applied[len(rep.Applied)-1].Version
+	}
+	if len(rep.Pending) == 0 {
+		rep.Status = "OK"
+		rep.Next = "Already at latest version"
+	} else {
+		rep.Status = "PENDING"
+		rep.Next = rep.Pending[0].Version()
+	}
+	// If the last one is partially applied (and not manually resolved).
+	if len(rep.Applied) != 0 {
+		last := rep.Applied[len(rep.Applied)-1]
+		if !last.Type.Has(migrate.RevisionTypeResolved) && last.Applied < last.Total {
+			rep.SQL = strings.ReplaceAll(last.ErrorStmt, "\n", " ")
+			rep.Error = strings.ReplaceAll(last.Error, "\n", " ")
+			rep.Count = last.Applied
+			idx := migrate.FilesLastIndex(rep.Available, func(f migrate.File) bool {
+				return f.Version() == last.Version
+			})
+			if idx == -1 {
+				return fmt.Errorf("migration file with version %q not found", last.Version)
+			}
+			stmts, err := rep.Available[idx].Stmts()
+			if err != nil {
+				return err
+			}
+			rep.Total = len(stmts)
+		}
+	}
+	return r.ReportWriter.WriteReport(rep)
 }
 
 // Left returns the amount of statements left to apply (if any).
@@ -287,3 +296,20 @@ func table(report *StatusReport) (string, error) {
 	tbl.Render()
 	return buf.String(), nil
 }
+
+func merge(maps ...template.FuncMap) template.FuncMap {
+	switch len(maps) {
+	case 0:
+		return nil
+	case 1:
+		return maps[0]
+	default:
+		m := maps[0]
+		for _, e := range maps[1:] {
+			for k, v := range e {
+				m[k] = v
+			}
+		}
+		return m
+	}
+}
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index 59f6d4b4..964b86f0 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -351,10 +351,25 @@ func (b *Builder) rewriteLastByte(c byte) {
 
 // IsQuoted reports if the given string is quoted with one of the given quotes (e.g. ', ", `).
 func IsQuoted(s string, q ...byte) bool {
-	for i := range q {
-		if l, r := strings.IndexByte(s, q[i]), strings.LastIndexByte(s, q[i]); l < r && l == 0 && r == len(s)-1 {
-			return true
+	last := len(s) - 1
+	if last < 1 {
+		return false
+	}
+Top:
+	for _, quote := range q {
+		if s[0] != quote || s[last] != quote {
+			continue
+		}
+		for i := 1; i < last-1; i++ {
+			switch c := s[i]; {
+			case c == '\\', c == quote && s[i+1] == quote:
+				i++
+			// Accept only escaped quotes and eject otherwise.
+			case c == quote:
+				continue Top
+			}
 		}
+		return true
 	}
 	return false
 }


Test patch:
diff --git a/cmd/atlas/internal/migrate/report_test.go b/cmd/atlas/internal/migrate/report_test.go
index 523e6a40..994dab93 100644
--- a/cmd/atlas/internal/migrate/report_test.go
+++ b/cmd/atlas/internal/migrate/report_test.go
@@ -27,7 +27,7 @@ func TestReporter_Status(t *testing.T) {
 	c, err := sqlclient.Open(ctx, "sqlite://?mode=memory")
 	require.NoError(t, err)
 	defer c.Close()
-	require.NoError(t, (&Reporter{
+	require.NoError(t, (&StatusReporter{
 		Client:       c,
 		Dir:          dir,
 		ReportWriter: &TemplateWriter{T: DefaultTemplate, W: &buf},
@@ -47,7 +47,7 @@ func TestReporter_Status(t *testing.T) {
 	ex, err := migrate.NewExecutor(c.Driver, dir, rrw)
 	require.NoError(t, err)
 	require.NoError(t, ex.ExecuteN(ctx, 1))
-	require.NoError(t, (&Reporter{
+	require.NoError(t, (&StatusReporter{
 		Client:       c,
 		Dir:          dir,
 		ReportWriter: &TemplateWriter{T: DefaultTemplate, W: &buf},
@@ -63,7 +63,7 @@ func TestReporter_Status(t *testing.T) {
 	buf.Reset()
 	require.NoError(t, err)
 	require.NoError(t, ex.ExecuteN(ctx, 1))
-	require.NoError(t, (&Reporter{
+	require.NoError(t, (&StatusReporter{
 		Client:       c,
 		Dir:          dir,
 		ReportWriter: &TemplateWriter{T: DefaultTemplate, W: &buf},
@@ -79,7 +79,7 @@ func TestReporter_Status(t *testing.T) {
 	buf.Reset()
 	require.NoError(t, err)
 	require.Error(t, ex.ExecuteN(ctx, 1))
-	require.NoError(t, (&Reporter{
+	require.NoError(t, (&StatusReporter{
 		Client:       c,
 		Dir:          dir,
 		ReportWriter: &TemplateWriter{T: DefaultTemplate, W: &buf},
@@ -102,7 +102,7 @@ Last migration attempt had errors:
 	*dir = *dir2
 	require.NoError(t, err)
 	require.NoError(t, ex.ExecuteN(ctx, 1))
-	require.NoError(t, (&Reporter{
+	require.NoError(t, (&StatusReporter{
 		Client:       c,
 		Dir:          dir,
 		ReportWriter: &TemplateWriter{T: DefaultTemplate, W: &buf},
diff --git a/sql/internal/sqlx/sqlx_test.go b/sql/internal/sqlx/sqlx_test.go
index 952bb751..72251621 100644
--- a/sql/internal/sqlx/sqlx_test.go
+++ b/sql/internal/sqlx/sqlx_test.go
@@ -144,6 +144,40 @@ func TestExprLastIndex(t *testing.T) {
 	}
 }
 
+func TestIsQuoted(t *testing.T) {
+	tests := []struct {
+		input  string
+		quotes []byte
+		want   bool
+	}{
+		{"''", []byte{'"', '\''}, true},
+		{`""`, []byte{'"', '\''}, true},
+		{"' '' \"\"'' '", []byte{'\''}, true},
+		{"''''''''", []byte{'\''}, true},
+		{"'foo'''", []byte{'\''}, true},
+		{"'foo'''''", []byte{'\''}, true},
+		{"'foo'', '''", []byte{'\''}, true},
+		{"'foo bar'", []byte{'\''}, true},
+		{`"never say \"never\""`, []byte{'"'}, true},
+		{`"never say \"never\'"`, []byte{'"'}, true},
+		{`'never say \"never\''`, []byte{'\''}, true},
+
+		{"'", []byte{'"', '\''}, false},
+		{`"`, []byte{'"', '\''}, false},
+		{"'foo' ''", []byte{'\''}, false},
+		{"'foo' ()  ''", []byte{'\''}, false},
+		{"'foo', ''", []byte{'\''}, false},
+		{"'foo', 'bar'", []byte{'\''}, false},
+		{"'foo',\" 'bar'", []byte{'\''}, false},
+	}
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			quoted := IsQuoted(tt.input, tt.quotes...)
+			require.Equal(t, tt.want, quoted)
+		})
+	}
+}
+
 func TestReverseChanges(t *testing.T) {
 	tests := []struct {
 		input  []schema.Change