refactor (codegen): reduce redundancy, make 'get item(s) by blah' generators more general

pkg/codegen/modelgenerate/generate_model.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"go/ast"
 	"go/token"
+	"slices"
 	"strings"
 
 	"github.com/jinzhu/inflection"
@@ -539,13 +540,48 @@ func getByIDFuncName(tblname string) string {
 	return "Get" + schema.TypenameFromTablename(tblname) + "ByID"
 }
 
+// paramNamesFor picks parameter identifiers for a `Get...By()` function's columns.
+// A single column always uses the short GoVarName() form (e.g. "name").  For multiple
+// columns, GoVarName() is preferred for readability, but if two or more of the given
+// columns would produce the same short name (e.g. two foreign keys that both abbreviate
+// to "uID"), the longer, unambiguous LongGoVarName() form is used for all of them instead,
+// since a collision there would produce invalid (duplicate-parameter) Go code.
+func paramNamesFor(cols []schema.Column) []string {
+	names := make([]string, len(cols))
+	if len(cols) == 1 {
+		names[0] = cols[0].GoVarName()
+		return names
+	}
+
+	hasConflict := false
+	for i, c := range cols {
+		names[i] = c.GoVarName()
+		if slices.Contains(names[:i], names[i]) {
+			hasConflict = true
+			break
+		}
+	}
+	if !hasConflict {
+		return names
+	}
+	for i, c := range cols {
+		names[i] = c.LongGoVarName()
+	}
+	return names
+}
+
+// GenerateGetItemBy produces an AST for a `GetXyzByCol1AndCol2...()` function that returns
+// the single item matching an exact-match lookup over the given columns (or ErrNotInDB).
+// Used for unique indexes (single- or multi-column) and for "without rowid" tables' primary keys.
 func GenerateGetItemBy(tbl schema.Table, cols []schema.Column) *ast.FuncDecl {
+	paramNames := paramNamesFor(cols)
+
 	colNames := []string{}
 	funcNameSuffix := []string{}
 	funcParams := &ast.FieldList{List: []*ast.Field{}}
 	sqlParams := []ast.Expr{}
-	for _, col := range cols {
+	for i, col := range cols {
-		funcParam := ast.NewIdent(col.LongGoVarName())
+		funcParam := ast.NewIdent(paramNames[i])
 		funcParams.List = append(funcParams.List, &ast.Field{Names: []*ast.Ident{funcParam}, Type: GoTypeForColumn(col)})
 		colNames = append(colNames, fmt.Sprintf("%s = ?", col.Name))
 		funcNameSuffix = append(funcNameSuffix, col.GoFieldName())
@@ -635,57 +671,33 @@ func GenerateGetItemByIDFunc(tbl schema.Table) *ast.FuncDecl {
 	}
 }
 
-// GenerateGetItemByUniqColFunc produces an AST for the `GetXyzByID()` function.
+// GenerateGetItemByUniqColFunc produces an AST for the `GetXyzByCol()` function, for a unique
-// E.g., a table with `table.TypeName = "foods"` will produce a "GetFoodByID()" function.
-func GenerateGetItemByUniqColFunc(tbl schema.Table, col schema.Column) *ast.FuncDecl {
-	// Use the xyzSQLFields constant in the select query
-	selectExpr := &ast.BinaryExpr{
-		X: &ast.BinaryExpr{
-			X:  &ast.BasicLit{Kind: token.STRING, Value: "`\n\t    select `"},
-			Op: token.ADD,
-			Y:  SQLFieldsConstIdent(tbl),
-		},
-		Op: token.ADD,
-		Y:  &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("`\n\t      from %s\n\t     where %s = ?\n\t`", tbl.TableName, col.Name)},
-	}
-
-	param := ast.NewIdent(col.GoVarName())
-
-	return &ast.FuncDecl{
-		Recv: dbRecv,
-		Name: ast.NewIdent("Get" + schema.TypenameFromTablename(tbl.TableName) + "By" + col.GoFieldName()),
-		Type: &ast.FuncType{
-			Params: &ast.FieldList{List: []*ast.Field{
-				{Names: []*ast.Ident{param}, Type: GoTypeForColumn(col)},
-			}},
-			Results: &ast.FieldList{List: []*ast.Field{
-				{Names: []*ast.Ident{ast.NewIdent("ret")}, Type: ast.NewIdent(tbl.GoTypeName)},
-				{Names: []*ast.Ident{ast.NewIdent("err")}, Type: ast.NewIdent("error")},
-			}},
-		},
-		Body: &ast.BlockStmt{
-			List: []ast.Stmt{
-				&ast.AssignStmt{
-					Lhs: []ast.Expr{ast.NewIdent("err")},
-					Tok: token.ASSIGN,
-					Rhs: []ast.Expr{&ast.CallExpr{Fun: &ast.SelectorExpr{X: dbDB, Sel: ast.NewIdent("Get")}, Args: []ast.Expr{&ast.UnaryExpr{Op: token.AND, X: ast.NewIdent("ret")}, selectExpr, param}}},
-				},
-				&ast.IfStmt{
-					Cond: &ast.CallExpr{Fun: &ast.SelectorExpr{X: ast.NewIdent("errors"), Sel: ast.NewIdent("Is")}, Args: []ast.Expr{ast.NewIdent("err"), &ast.SelectorExpr{X: ast.NewIdent("sql"), Sel: ast.NewIdent("ErrNoRows")}}},
-					Body: &ast.BlockStmt{List: []ast.Stmt{&ast.ReturnStmt{Results: []ast.Expr{&ast.CompositeLit{Type: ast.NewIdent(tbl.GoTypeName)}, ast.NewIdent("ErrNotInDB")}}}},
-				},
-				&ast.ReturnStmt{},
-			},
-		},
-	}
-}
-
-// GenerateGetItemsByColFunc produces an AST for the `GetXyzsByCol()` function, for a non-unique
 // index on a single column.
+// E.g., a table with `table.TypeName = "foods"` will produce a "GetFoodByName()" function.
+func GenerateGetItemByUniqColFunc(tbl schema.Table, col schema.Column) *ast.FuncDecl {
+	return GenerateGetItemBy(tbl, []schema.Column{col})
+}
+
+// GenerateGetItemsBy produces an AST for a `GetXyzsByCol1AndCol2...()` function that returns
+// all items matching an exact-match lookup over the given columns.  Used for non-unique
+// indexes (single- or multi-column).
 // E.g., a table with `table.TableName = "foods"` and a non-unique index on "category" will
 // produce a "GetFoodsByCategory(category string) []Food" function.
-func GenerateGetItemsByColFunc(tbl schema.Table, col schema.Column) *ast.FuncDecl {
+func GenerateGetItemsBy(tbl schema.Table, cols []schema.Column) *ast.FuncDecl {
-	// Use the xyzSQLFields constant in the select query
+	paramNames := paramNamesFor(cols)
+
+	colNames := []string{}
+	funcNameSuffix := []string{}
+	funcParams := &ast.FieldList{List: []*ast.Field{}}
+	sqlParams := []ast.Expr{}
+	for i, col := range cols {
+		funcParam := ast.NewIdent(paramNames[i])
+		funcParams.List = append(funcParams.List, &ast.Field{Names: []*ast.Ident{funcParam}, Type: GoTypeForColumn(col)})
+		colNames = append(colNames, fmt.Sprintf("%s = ?", col.Name))
+		funcNameSuffix = append(funcNameSuffix, col.GoFieldName())
+		sqlParams = append(sqlParams, funcParam)
+	}
+
 	selectExpr := &ast.BinaryExpr{
 		X: &ast.BinaryExpr{
 			X:  &ast.BasicLit{Kind: token.STRING, Value: "`\n\t    select `"},
@@ -693,23 +705,19 @@ func GenerateGetItemsByColFunc(tbl schema.Table, col schema.Column) *ast.FuncDec
 			Y:  SQLFieldsConstIdent(tbl),
 		},
 		Op: token.ADD,
-		Y:  &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("`\n\t      from %s\n\t     where %s = ?\n\t`", tbl.TableName, col.Name)},
+		Y:  &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("`\n\t      from %s\n\t     where %s\n\t`", tbl.TableName, strings.Join(colNames, " and "))},
 	}
 
-	param := ast.NewIdent(col.GoVarName())
-
 	selectCall := doCall(&ast.CallExpr{
 		Fun:  &ast.SelectorExpr{X: dbDB, Sel: ast.NewIdent("Select")},
-		Args: []ast.Expr{&ast.UnaryExpr{Op: token.AND, X: ast.NewIdent("ret")}, selectExpr, param},
+		Args: append([]ast.Expr{&ast.UnaryExpr{Op: token.AND, X: ast.NewIdent("ret")}, selectExpr}, sqlParams...),
 	})
 
 	return &ast.FuncDecl{
 		Recv: dbRecv,
-		Name: ast.NewIdent("Get" + inflection.Plural(schema.TypenameFromTablename(tbl.TableName)) + "By" + col.GoFieldName()),
+		Name: ast.NewIdent("Get" + inflection.Plural(schema.TypenameFromTablename(tbl.TableName)) + "By" + strings.Join(funcNameSuffix, "And")),
 		Type: &ast.FuncType{
-			Params: &ast.FieldList{List: []*ast.Field{
+			Params: funcParams,
-				{Names: []*ast.Ident{param}, Type: GoTypeForColumn(col)},
-			}},
 			Results: &ast.FieldList{List: []*ast.Field{
 				{Names: []*ast.Ident{ast.NewIdent("ret")}, Type: &ast.ArrayType{Elt: ast.NewIdent(tbl.GoTypeName)}},
 			}},
@@ -723,6 +731,14 @@ func GenerateGetItemsByColFunc(tbl schema.Table, col schema.Column) *ast.FuncDec
 	}
 }
 
+// GenerateGetItemsByColFunc produces an AST for the `GetXyzsByCol()` function, for a non-unique
+// index on a single column.
+// E.g., a table with `table.TableName = "foods"` and a non-unique index on "category" will
+// produce a "GetFoodsByCategory(category string) []Food" function.
+func GenerateGetItemsByColFunc(tbl schema.Table, col schema.Column) *ast.FuncDecl {
+	return GenerateGetItemsBy(tbl, []schema.Column{col})
+}
+
 // GenerateGetAllItemsFunc produces an AST for the `GetAllXyzs()` function.
 // E.g., a table with `table.TypeName = "foods"` will produce a "GetAllFoods()" function.
 func GenerateGetAllItemsFunc(tbl schema.Table) *ast.FuncDecl {