Repository: dynamoosejs/dynamoose
Pull request #770 (2020-04-04 03:32:13)

Custom method improvements
This PR improves support for custom methods, including better documentation, and being able to pass parameters into the custom methods.


Patch:
diff --git a/docs/api/Model.md b/docs/api/Model.md
index aeab8e3f..2520a15f 100644
--- a/docs/api/Model.md
+++ b/docs/api/Model.md
@@ -430,7 +430,7 @@ User.methods.set("scanAll", async function () {
 	let results = await this.scan().exec();
 	lastKey = results.lastKey;
 	do {
-		const newResult = await this.scan().startKey(lastKey).exec();
+		const newResult = await this.scan().startAt(lastKey).exec();
 		results = [...results, ...newResult];
 		lastKey = newResult.lastKey;
 	} while (lastKey)
@@ -440,7 +440,11 @@ User.methods.set("scanAll", async function () {
 User.methods.set("scanAll", function (cb) {
 	let result = [];
 	const main = (lastKey) => {
-		this.scan().exec((err, newResult) => {
+		let scan = this.scan();
+		if (lastKey) {
+			scan.startAt(lastKey);
+		}
+		scan.exec((err, newResult) => {
 			if (err) {
 				cb(err);
 			} else {
@@ -469,6 +473,62 @@ models.forEach((model) => {
 });
 ```
 
+You can also pass parameters into your custom method. It is important to note that if you decide to pass custom parameters into your custom method, the `callback` parameter will always be passed in as the last parameter. This means it's highly recommended that you always pass in the same number of parameters every time to your custom method. In the event you are unable to do this (dynamic/custom parameter length), you can use the JavaScript `arguments` variable to retrieve the last argument that was passed into the function.
+
+```js
+// Setup:
+const User = new dynamoose.model("Model", ModelSchema);
+User.methods.set("scanAll", async function (startAt) {
+	let scan = this.scan();
+	if (startAt) {
+		scan.startAt(startAt);
+	}
+	let results = await scan.exec();
+	lastKey = results.lastKey;
+	do {
+		const newResult = await this.scan().startAt(lastKey).exec();
+		results = [...results, ...newResult];
+		lastKey = newResult.lastKey;
+	} while (lastKey)
+	return results;
+});
+// OR
+User.methods.set("scanAll", function (startAt, cb) {
+	let result = [];
+	const main = (lastKey) => {
+		let scan = this.scan();
+		if (lastKey) {
+			scan.startAt(lastKey);
+		}
+		scan.exec((err, newResult) => {
+			if (err) {
+				cb(err);
+			} else {
+				result = [...result, ...newResult];
+				if (newResult.lastKey) {
+					main(newResult.lastKey);
+				} else {
+					cb(result);
+				}
+			}
+		});
+	};
+	main(startAt);
+});
+
+// Using:
+User.scanAll({"id": 1024}, (err, models) => {
+	models.forEach((model) => {
+		console.log(model);
+	});
+});
+// OR
+const models = await User.scanAll({"id": 1024});
+models.forEach((model) => {
+	console.log(model);
+});
+```
+
 ## Model.methods.delete(name)
 
 This allows you to delete an existing method from the model. If no existing method is assigned for that name, the function will do nothing and no error will be thrown.
@@ -492,7 +552,7 @@ User.methods.document.set("setName", async function () {
 	this.name = await getRandomName();
 });
 // OR
-User.methods.set("scanAll", function (cb) {
+User.methods.document.set("setName", function (cb) {
 	getRandomName((err, name) => {
 		if (err) {
 			cb(err);
@@ -514,6 +574,37 @@ await user.setName();
 console.log("Set name");
 ```
 
+You can also pass parameters into your custom method. It is important to note that if you decide to pass custom parameters into your custom method, the `callback` parameter will always be passed in as the last parameter. This means it's highly recommended that you always pass in the same number of parameters every time to your custom method. In the event you are unable to do this (dynamic/custom parameter length), you can use the JavaScript `arguments` variable to retrieve the last argument that was passed into the function.
+
+```js
+// Setup:
+const User = new dynamoose.model("Model", ModelSchema);
+User.methods.document.set("setName", async function (firstName, lastName) {
+	this.name = await verifyName(`${firstName} ${lastName}`);
+});
+// OR
+User.methods.document.set("scanAll", function (firstName, lastName, cb) {
+	verifyName(`${firstName} ${lastName}`, (err, name) => {
+		if (err) {
+			cb(err);
+		} else {
+			this.name = name;
+			cb();
+		}
+	});
+});
+
+// Using:
+const user = new User();
+
+user.setName("Charlie", "Fish", (err) => {
+	console.log("Set name");
+});
+// OR
+await user.setName("Charlie", "Fish");
+console.log("Set name");
+```
+
 ## Model.methods.document.delete(name)
 
 This allows you to delete an existing method from the document. If no existing method is assigned for that name, the function will do nothing and no error will be thrown.
diff --git a/lib/Model.js b/lib/Model.js
index b99cb44b..900aed79 100644
--- a/lib/Model.js
+++ b/lib/Model.js
@@ -517,15 +517,17 @@ Model.prototype.query = {"carrier": require("./DocumentRetriever")("query")};
 Model.prototype.methods = {
 	"set": function (name, fn) {
 		if (!this.Document[name] || (this.Document[name][Internal.internalProperties] && this.Document[name][Internal.internalProperties].type === "customMethod")) {
-			this.Document[name] = (cb) => {
+			const self = this;
+			this.Document[name] = function (...args) {
+				const cb = typeof args[args.length - 1] === "function" ? args[args.length - 1] : undefined;
 				if (cb) {
-					const result = fn.bind(this.Document)(cb);
+					const result = fn.bind(self.Document)(...args);
 					if (result instanceof Promise) {
 						result.then((result) => cb(null, result)).catch((err) => cb(err));
 					}
 				} else {
 					return new Promise((resolve, reject) => {
-						const result = fn.bind(this.Document)((err, result) => {
+						const result = fn.bind(self.Document)(...args, (err, result) => {
 							if (err) {
 								reject(err);
 							} else {
@@ -549,15 +551,16 @@ Model.prototype.methods = {
 	"document": {
 		"set": function (name, fn) {
 			if (!this.Document.prototype[name] || (this.Document.prototype[name][Internal.internalProperties] && this.Document.prototype[name][Internal.internalProperties].type === "customMethod")) {
-				this.Document.prototype[name] = function (cb) {
+				this.Document.prototype[name] = function (...args) {
+					const cb = typeof args[args.length - 1] === "function" ? args[args.length - 1] : undefined;
 					if (cb) {
-						const result = fn.bind(this)(cb);
+						const result = fn.bind(this)(...args);
 						if (result instanceof Promise) {
 							result.then((result) => cb(null, result)).catch((err) => cb(err));
 						}
 					} else {
 						return new Promise((resolve, reject) => {
-							const result = fn.bind(this)((err, result) => {
+							const result = fn.bind(this)(...args, (err, result) => {
 								if (err) {
 									reject(err);
 								} else {


Test patch:
diff --git a/test/Model.js b/test/Model.js
index 32c476e1..9fd105a5 100644
--- a/test/Model.js
+++ b/test/Model.js
@@ -2708,6 +2708,23 @@ describe("Model", () => {
 								expect(methodTypeCallDetails[0].arguments[0]).to.be.a("function");
 							});
 
+							it("Should call custom method with correct arguments if passing in one argument", async () => {
+								await callerType.func(User.action)("Hello World");
+								expect(methodTypeCallDetails.length).to.eql(1);
+								expect(methodTypeCallDetails[0].arguments.length).to.eql(2);
+								expect(methodTypeCallDetails[0].arguments[0]).to.eql("Hello World");
+								expect(methodTypeCallDetails[0].arguments[1]).to.be.a("function");
+							});
+
+							it("Should call custom method with correct arguments if passing in two arguments", async () => {
+								await callerType.func(User.action)("Hello", "World");
+								expect(methodTypeCallDetails.length).to.eql(1);
+								expect(methodTypeCallDetails[0].arguments.length).to.eql(3);
+								expect(methodTypeCallDetails[0].arguments[0]).to.eql("Hello");
+								expect(methodTypeCallDetails[0].arguments[1]).to.eql("World");
+								expect(methodTypeCallDetails[0].arguments[2]).to.be.a("function");
+							});
+
 							it("Should call custom method with correct `this`", async () => {
 								methodTypeCallResult.result = "Success";
 								await callerType.func(User.action)();
@@ -2833,6 +2850,23 @@ describe("Model", () => {
 									expect(methodTypeCallDetails[0].arguments[0]).to.be.a("function");
 								});
 
+								it("Should call custom method with correct arguments if passing in one argument", async () => {
+									await callerType.func(user.action)("Hello World");
+									expect(methodTypeCallDetails.length).to.eql(1);
+									expect(methodTypeCallDetails[0].arguments.length).to.eql(2);
+									expect(methodTypeCallDetails[0].arguments[0]).to.eql("Hello World");
+									expect(methodTypeCallDetails[0].arguments[1]).to.be.a("function");
+								});
+
+								it("Should call custom method with correct arguments if passing in two arguments", async () => {
+									await callerType.func(user.action)("Hello", "World");
+									expect(methodTypeCallDetails.length).to.eql(1);
+									expect(methodTypeCallDetails[0].arguments.length).to.eql(3);
+									expect(methodTypeCallDetails[0].arguments[0]).to.eql("Hello");
+									expect(methodTypeCallDetails[0].arguments[1]).to.eql("World");
+									expect(methodTypeCallDetails[0].arguments[2]).to.be.a("function");
+								});
+
 								it("Should call custom method with correct `this`", async () => {
 									methodTypeCallResult.result = "Success";
 									// Not sure why we have to bind `user` here, when we don't have to bind anything for Model.methods.set