From a1b06a3ef8dce0725ddd9a67b9700830ec053b28 Mon Sep 17 00:00:00 2001
From: ~wispem-wantex <wispem-wantex@playfulpachyderm.com>
Date: Tue, 8 Sep 2026 18:18:31 -0700
Subject: [PATCH] codegen: use `must.Be` helper in Save function

---
 pkg/codegen/modelgenerate/generate_model.go | 50 ++++++++++++++++-----
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/pkg/codegen/modelgenerate/generate_model.go b/pkg/codegen/modelgenerate/generate_model.go
index 96e1074..b4dca52 100644
--- a/pkg/codegen/modelgenerate/generate_model.go
+++ b/pkg/codegen/modelgenerate/generate_model.go
@@ -53,20 +53,48 @@ func GoTypeForColumn(c schema.Column) ast.Expr {
 	}
 }
 
-func PanicIfRowsAffected(tbl schema.Table) *ast.IfStmt {
-	return &ast.IfStmt{
-		Cond: &ast.BinaryExpr{
+// MustBeRowsAffected produces an AST for a `must.Be(...)` call asserting that exactly one row
+// was affected by the preceding statement, e.g.:
+//
+//	must.Be(must.Get(result.RowsAffected()) == 1, "%w: Food ID=%d", ErrNotInDB, f.ID)
+//
+// For "without rowid" tables, the message includes the table's primary key column(s) instead of ID.
+func MustBeRowsAffected(tbl schema.Table) *ast.ExprStmt {
+	pkParts := []string{}
+	pkArgs := []ast.Expr{}
+	if tbl.IsWithoutRowid {
+		for _, col := range tbl.PrimaryKeyColumns() {
+			verb := "%v"
+			if col.Type == "integer" || col.Type == "int" || col.IsNonCodeTableForeignKey() {
+				verb = "%d"
+			}
+			pkParts = append(pkParts, fmt.Sprintf("%s=%s", col.Name, verb))
+			pkArgs = append(pkArgs, &ast.SelectorExpr{X: ast.NewIdent(tbl.VarName), Sel: ast.NewIdent(col.GoFieldName())})
+		}
+	} else {
+		pkParts = append(pkParts, "ID=%d")
+		pkArgs = append(pkArgs, &ast.SelectorExpr{X: ast.NewIdent(tbl.VarName), Sel: ast.NewIdent("ID")})
+	}
+
+	msg := fmt.Sprintf("%%w: %s %s", tbl.GoTypeName, strings.Join(pkParts, ", "))
+
+	args := append([]ast.Expr{
+		&ast.BinaryExpr{
 			X: mustCall(&ast.CallExpr{
 				Fun:  &ast.SelectorExpr{X: ast.NewIdent("result"), Sel: ast.NewIdent("RowsAffected")},
 				Args: []ast.Expr{},
 			}),
-			Op: token.NEQ,
+			Op: token.EQL,
 			Y:  &ast.BasicLit{Kind: token.INT, Value: "1"},
 		},
-		Body: &ast.BlockStmt{List: []ast.Stmt{
-			&ast.ExprStmt{X: &ast.CallExpr{Fun: ast.NewIdent("panic"), Args: []ast.Expr{ast.NewIdent(tbl.VarName)}}},
-		}},
-	}
+		&ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("%q", msg)},
+		ast.NewIdent("ErrNotInDB"),
+	}, pkArgs...)
+
+	return &ast.ExprStmt{X: &ast.CallExpr{
+		Fun:  &ast.SelectorExpr{X: ast.NewIdent("must"), Sel: ast.NewIdent("Be")},
+		Args: args,
+	}}
 }
 
 // ---------------
@@ -418,7 +446,7 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl {
 					})
 				}
 				ret = append(ret, namedExecStmt(upsertStmt)...)
-				ret = append(ret, PanicIfRowsAffected(tbl))
+				ret = append(ret, MustBeRowsAffected(tbl))
 			} else {
 				// if item.ID == 0 {...} else {...}
 				ret = append(ret, &ast.IfStmt{
@@ -472,7 +500,7 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl {
 							[]ast.Stmt{Comment("Do update")},
 							append(
 								namedExecStmt(updateStmt),
-								PanicIfRowsAffected(tbl),
+								MustBeRowsAffected(tbl),
 							)...,
 						),
 					},
@@ -720,7 +748,7 @@ func GenerateDeleteItemFunc(tbl schema.Table) *ast.FuncDecl {
 					},
 				})},
 			},
-			PanicIfRowsAffected(tbl),
+			MustBeRowsAffected(tbl),
 		},
 	}