Make execs with assert rows more compact

Author: hlubek

README.md

@@ -141,11 +141,9 @@ func FindCustomerByID(ctx context.Context, executor qrbpgx.Executor, id uuid.UUI
     Where(customer.ID.Eq(qrb.Arg(id))).
     GroupBy(customer.ID)
 
-  row, err := qrbpgx.Build(q).WithExecutor(executor).QueryRow(ctx)
-  if err != nil {
-    return result, err
-  }
-  return constructpgx.ScanRow[model.Customer](row)
+  return constructpgx.ScanRow[model.Customer](
+    qrbpgx.Build(q).WithExecutor(executor).QueryRow(ctx),
+  )
 }
 
 // CustomerChangeSet is generated by construct for handling partially filled models
@@ -165,12 +163,9 @@ func UpdateCustomer(ctx context.Context, executor qrbpgx.Executor, id uuid.UUID,
     Where(customer.ID.Eq(qrb.Arg(id))).
     SetMap(changeSet.toMap())
 
-  res, err := qrbpgx.Build(q).WithExecutor(executor).Exec(ctx)
-  if err != nil {
-    return err
-  }
-
-  return constructpgx.AssertRowsAffected(res, "update", 1)
+  return constructpgx.AssertRowsAffected("update", 1)(
+    qrbpgx.Build(q).WithExecutor(executor).Exec(ctx),
+  )
 }
 
 func customerJson() builder.JsonBuildObjectBuilder {

constructpgx/constructpgx.go

@@ -61,13 +61,14 @@ func scanRow[T any](row pgx.Row) (T, error) {
 	return result, nil
 }
 
-// AssertRowsAffected checks if the given result affected exactly the expected number of rows.
-func AssertRowsAffected(result pgconn.CommandTag, operation string, expectedRows int) error {
-	actualRows := result.RowsAffected()
-	if actualRows != int64(expectedRows) {
-		return fmt.Errorf("%s affected %d rows, but expected exactly %d", operation, actualRows, expectedRows)
-	}
-
-	return nil
+// AssertRowsAffected returns a func that checks if the given operation result affected exactly the expected number of rows.
+func AssertRowsAffected(operation string, expectedRows int) func(pgconn.CommandTag) error {
+	return func(result pgconn.CommandTag) error {
+		actualRows := result.RowsAffected()
+		if actualRows != int64(expectedRows) {
+			return fmt.Errorf("%s affected %d rows, but expected exactly %d", operation, actualRows, expectedRows)
+		}
 
+		return nil
+	}
 }

constructpgx/constructpgx_test.go

@@ -88,12 +88,12 @@ func TestScanRow(t *testing.T) {
 
 func TestAssertRowsAffected(t *testing.T) {
 	t.Run("returns nil if rows affected matches expected", func(t *testing.T) {
-		err := constructpgx.AssertRowsAffected(pgconn.NewCommandTag("UPDATE 2"), "update", 2)
+		err := constructpgx.AssertRowsAffected("update", 2)(pgconn.NewCommandTag("UPDATE 2"))
 		require.NoError(t, err)
 	})
 
 	t.Run("returns error if rows affected does not match expected", func(t *testing.T) {
-		err := constructpgx.AssertRowsAffected(pgconn.NewCommandTag("UPDATE 2"), "update", 1)
+		err := constructpgx.AssertRowsAffected("update", 1)(pgconn.NewCommandTag("UPDATE 2"))
 		require.Error(t, err)
 	})
 }

constructsql/constructsql.go

@@ -71,16 +71,20 @@ func scanRow[T any](row RowScanner) (result T, err error) {
 }
 
 // AssertRowsAffected checks if the given result affected exactly the expected number of rows.
-func AssertRowsAffected(result sql.Result, operation string, expectedRows int) error {
-	actualRows, err := result.RowsAffected()
-	if err != nil {
-		return fmt.Errorf("getting affected rows: %w", err)
-	}
-
-	if actualRows != int64(expectedRows) {
-		return fmt.Errorf("%s affected %d rows, but expected exactly %d", operation, actualRows, expectedRows)
-	}
+func AssertRowsAffected(operation string, expectedRows int) func(sql.Result, error) error {
+	return func(result sql.Result, err error) error {
+		if err != nil {
+			return err
+		}
+		actualRows, err := result.RowsAffected()
+		if err != nil {
+			return fmt.Errorf("getting affected rows: %w", err)
+		}
 
-	return nil
+		if actualRows != int64(expectedRows) {
+			return fmt.Errorf("%s affected %d rows, but expected exactly %d", operation, actualRows, expectedRows)
+		}
 
+		return nil
+	}
 }

constructsql/constructsql_test.go

@@ -100,18 +100,22 @@ func (r sqlResult) RowsAffected() (int64, error) {
 
 func TestAssertRowsAffected(t *testing.T) {
 	t.Run("returns nil if rows affected matches expected", func(t *testing.T) {
-
-		err := constructsql.AssertRowsAffected(sqlResult{affectedRows: 2}, "update", 2)
+		err := constructsql.AssertRowsAffected("update", 2)(sqlResult{affectedRows: 2}, nil)
 		require.NoError(t, err)
 	})
 
 	t.Run("returns error if rows affected does not match expected", func(t *testing.T) {
-		err := constructsql.AssertRowsAffected(sqlResult{affectedRows: 2}, "update", 1)
+		err := constructsql.AssertRowsAffected("update", 1)(sqlResult{affectedRows: 2}, nil)
 		require.Error(t, err)
 	})
 
 	t.Run("returns error if rows affected returns error", func(t *testing.T) {
-		err := constructsql.AssertRowsAffected(sqlResult{affectedRowsErr: errors.New("some error")}, "update", 1)
+		err := constructsql.AssertRowsAffected("update", 1)(sqlResult{affectedRowsErr: errors.New("some error")}, nil)
+		require.Error(t, err)
+	})
+
+	t.Run("returns previous error", func(t *testing.T) {
+		err := constructsql.AssertRowsAffected("update", 1)(sqlResult{}, errors.New("some error"))
 		require.Error(t, err)
 	})
 }