fix(compiler): Fix column expansion to work with quoted non-keyword identifiers

close #2575

Author: orisano

internal/compiler/expand.go

@@ -38,12 +38,7 @@ func (c *Compiler) expand(qc *QueryCatalog, raw *ast.RawStmt) ([]source.Edit, er
 
 func (c *Compiler) quoteIdent(ident string) string {
 	if c.parser.IsReservedKeyword(ident) {
-		switch c.conf.Engine {
-		case config.EngineMySQL:
-			return "`" + ident + "`"
-		default:
-			return "\"" + ident + "\""
-		}
+		return c.quote(ident)
 	}
 	if c.conf.Engine == config.EnginePostgreSQL {
 		// camelCase means the column is also camelCase
@@ -54,6 +49,15 @@ func (c *Compiler) quoteIdent(ident string) string {
 	return ident
 }
 
+func (c *Compiler) quote(x string) string {
+	switch c.conf.Engine {
+	case config.EngineMySQL:
+		return "`" + x + "`"
+	default:
+		return "\"" + x + "\""
+	}
+}
+
 func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node) ([]source.Edit, error) {
 	tables, err := c.sourceTables(qc, node)
 	if err != nil {
@@ -132,16 +136,36 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 		for _, p := range parts {
 			old = append(old, c.quoteIdent(p))
 		}
-		oldString := strings.Join(old, ".")
+
+		var oldString string
+		var oldFunc func(string) int
 
 		// use the sqlc.embed string instead
 		if embed, ok := qc.embeds.Find(ref); ok {
 			oldString = embed.Orig()
+		} else {
+			oldFunc = func(s string) int {
+				length := 0
+				for i, o := range old {
+					if hasSeparator := i > 0; hasSeparator {
+						length++
+					}
+					if strings.HasPrefix(s[length:], o) {
+						length += len(o)
+					} else if quoted := c.quote(o); strings.HasPrefix(s[length:], quoted) {
+						length += len(quoted)
+					} else {
+						length += len(o)
+					}
+				}
+				return length
+			}
 		}
 
 		edits = append(edits, source.Edit{
 			Location: res.Location - raw.StmtLocation,
 			Old:      oldString,
+			OldFunc:  oldFunc,
 			New:      strings.Join(cols, ", "),
 		})
 	}

internal/engine/sqlite/convert.go

@@ -767,7 +767,7 @@ func (c *cc) convertInsert_stmtContext(n *parser.Insert_stmtContext) ast.Node {
 		rel.Schemaname = &schemaName
 	}
 	if n.Table_alias() != nil {
-		tableAlias := n.Table_alias().GetText()
+		tableAlias := identifier(n.Table_alias().GetText())
 		rel.Alias = &ast.Alias{
 			Aliasname: &tableAlias,
 		}
@@ -837,7 +837,7 @@ func (c *cc) convertTablesOrSubquery(n []parser.ITable_or_subqueryContext) []ast
 				rv.Schemaname = &schema
 			}
 			if from.Table_alias() != nil {
-				alias := from.Table_alias().GetText()
+				alias := identifier(from.Table_alias().GetText())
 				rv.Alias = &ast.Alias{Aliasname: &alias}
 			}
 			if from.Table_alias_fallback() != nil {
@@ -870,7 +870,7 @@ func (c *cc) convertTablesOrSubquery(n []parser.ITable_or_subqueryContext) []ast
 			}
 
 			if from.Table_alias() != nil {
-				alias := from.Table_alias().GetText()
+				alias := identifier(from.Table_alias().GetText())
 				rf.Alias = &ast.Alias{Aliasname: &alias}
 			}
 
@@ -881,7 +881,7 @@ func (c *cc) convertTablesOrSubquery(n []parser.ITable_or_subqueryContext) []ast
 			}
 
 			if from.Table_alias() != nil {
-				alias := from.Table_alias().GetText()
+				alias := identifier(from.Table_alias().GetText())
 				rs.Alias = &ast.Alias{Aliasname: &alias}
 			}
 

internal/source/code.go

@@ -12,6 +12,7 @@ type Edit struct {
 	Location int
 	Old      string
 	New      string
+	OldFunc  func(string) int
 }
 
 func LineNumber(source string, head int) (int, int) {
@@ -63,8 +64,14 @@ func Mutate(raw string, a []Edit) (string, error) {
 		if start > len(s) || start < 0 {
 			return "", fmt.Errorf("edit start location is out of bounds")
 		}
+		var oldLen int
+		if edit.OldFunc != nil {
+			oldLen = edit.OldFunc(s[start:])
+		} else {
+			oldLen = len(edit.Old)
+		}
 
-		stop := edit.Location + len(edit.Old)
+		stop := edit.Location + oldLen
 		if stop > len(s) {
 			return "", fmt.Errorf("edit stop location is out of bounds")
 		}
@@ -73,7 +80,7 @@ func Mutate(raw string, a []Edit) (string, error) {
 		// this edit overlaps the previous one (and is therefore a developer error)
 		if idx != 0 {
 			prevEdit := a[idx-1]
-			if prevEdit.Location < edit.Location+len(edit.Old) {
+			if prevEdit.Location < edit.Location+oldLen {
 				return "", fmt.Errorf("2 edits overlap")
 			}
 		}