Merge pull request #25 from nezuo/beforeClose

Add beforeClose

src/Document.lua

@@ -1,5 +1,6 @@
 local Compression = require(script.Parent.Compression)
 local freezeDeep = require(script.Parent.freezeDeep)
+local Promise = require(script.Parent.Parent.Promise)
 
 --[=[
 	@class Document
@@ -15,6 +16,7 @@ function Document.new(collection, key, validate, lockId, data)
 		lockId = lockId,
 		data = data,
 		closed = false,
+		callingBeforeClose = false,
 	}, Document)
 end
 
@@ -55,7 +57,7 @@ end
 	@return Promise<()>
 ]=]
 function Document:save()
-	assert(not self.closed, "Cannot save a closed document")
+	assert(not self.closed and not self.callingBeforeClose, "Cannot save a closed document")
 
 	return self.collection.data:save(self.collection.dataStore, self.key, function(value)
 		if value.lockId ~= self.lockId then
@@ -79,30 +81,70 @@ end
 	Throws an error if the document was closed.
 	:::
 
+	:::warning
+	If the beforeClose callback errors, the returned promise will reject and the data will not be saved.
+	:::
+
 	@return Promise<()>
 ]=]
 function Document:close()
-	assert(not self.closed, "Cannot close a closed document")
+	assert(not self.closed and not self.callingBeforeClose, "Cannot close a closed document")
 
-	self.closed = true
+	local promise = Promise.resolve()
 
-	self.collection.openDocuments[self.key] = nil
+	if self.beforeCloseCallback ~= nil then
+		self.callingBeforeClose = true
 
-	self.collection.autoSave:removeDocument(self)
+		promise = Promise.new(function(resolve, reject)
+			local ok, err = pcall(self.beforeCloseCallback)
 
-	return self.collection.data:save(self.collection.dataStore, self.key, function(value)
-		if value.lockId ~= self.lockId then
-			return "fail", "The session lock was stolen"
-		end
+			if not ok then
+				reject(`beforeClose callback threw error: {tostring(err)}`)
+			else
+				resolve()
+			end
+		end)
+	end
 
-		local scheme, compressed = Compression.compress(self.data)
+	return promise
+		:finally(function()
+			self.closed = true
 
-		value.compressionScheme = scheme
-		value.data = compressed
-		value.lockId = nil
+			self.collection.openDocuments[self.key] = nil
 
-		return "succeed", value
-	end)
+			self.collection.autoSave:removeDocument(self)
+		end)
+		:andThen(function()
+			return self.collection.data:save(self.collection.dataStore, self.key, function(value)
+				if value.lockId ~= self.lockId then
+					return "fail", "The session lock was stolen"
+				end
+
+				local scheme, compressed = Compression.compress(self.data)
+
+				value.compressionScheme = scheme
+				value.data = compressed
+				value.lockId = nil
+
+				return "succeed", value
+			end)
+		end)
+end
+
+--[=[
+	Sets a callback that is run inside `document:close` before it saves. The document can be read and written to in the
+	callback.
+
+	:::warning
+	Throws an error if it was called previously.
+	:::
+
+	@param callback () -> ()
+]=]
+function Document:beforeClose(callback)
+	assert(self.beforeCloseCallback == nil, "Document:beforeClose can only be called once")
+
+	self.beforeCloseCallback = callback
 end
 
 return Document

src/Document.test.lua

@@ -10,6 +10,7 @@ local DEFAULT_OPTIONS = {
 }
 
 return function(x)
+	local assertEqual = x.assertEqual
 	local shouldThrow = x.shouldThrow
 
 	x.test("it should not merge close into save when save is running", function(context)
@@ -43,11 +44,14 @@ return function(x)
 		local pendingSave = document:save()
 		local pendingClose = document:close() -- This should override the pending save.
 
+		context.dataStoreService.budget.budgets[Enum.DataStoreRequestType.UpdateAsync] = 1
+
 		context.dataStoreService.yield:stopYield()
 
-		assert(pendingSave == pendingClose, "save and close didn't merge")
+		local values = Promise.all({ ongoingSave, pendingSave }):expect()
 
-		local values = Promise.all({ ongoingSave, pendingSave, pendingClose }):expect()
+		-- If the save and close don't merge, this will throw because there isn't enough budget to close.
+		pendingClose:now("save and close didn't merge"):expect()
 
 		-- save and close should never resolve with a value.
 		-- It's checked in this test to make sure it works with save merging.
@@ -200,4 +204,80 @@ return function(x)
 
 		promise:expect()
 	end)
+
+	x.nested("Document:beforeClose", function()
+		x.test("throws when setting twice", function(context)
+			local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect()
+
+			document:beforeClose(function() end)
+
+			shouldThrow(function()
+				document:beforeClose(function() end)
+			end, "Document:beforeClose can only be called once")
+		end)
+
+		x.test("throws when calling close in callback", function(context)
+			local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect()
+
+			document:beforeClose(function()
+				document:close()
+			end)
+
+			shouldThrow(function()
+				document:close():expect()
+			end, "beforeClose callback threw error")
+		end)
+
+		x.test("throws when calling save in callback", function(context)
+			local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect()
+
+			document:beforeClose(function()
+				document:save()
+			end)
+
+			shouldThrow(function()
+				document:close():expect()
+			end, "beforeClose callback threw error")
+		end)
+
+		x.test("closes document even if beforeClose errors", function(context)
+			local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
+
+			local promise = collection:load("document")
+			local document = promise:expect()
+
+			document:beforeClose(function()
+				error("error")
+			end)
+
+			shouldThrow(function()
+				document:close():expect()
+			end)
+
+			local secondPromise = collection:load("document")
+
+			assert(secondPromise ~= promise, "collection:load should return a new promise")
+
+			shouldThrow(function()
+				document:write({ foo = "baz" })
+			end, "Cannot write to a closed document")
+
+			-- Ignore the could not acquire lock error.
+			secondPromise:catch(function() end)
+		end)
+
+		x.test("saves new data", function(context)
+			local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect()
+
+			document:beforeClose(function()
+				document:read() -- This checks that read doesn't error in the callback.
+
+				document:write({ foo = "new" })
+			end)
+
+			document:close():expect()
+
+			assertEqual(context.read("collection", "document").data.foo, "new")
+		end)
+	end)
 end

src/init.lua

@@ -33,6 +33,7 @@ export type Document<T> = {
 	write: (self: Document<T>, T) -> (),
 	save: (self: Document<T>) -> PromiseTypes.TypedPromise<()>,
 	close: (self: Document<T>) -> PromiseTypes.TypedPromise<()>,
+	beforeClose: (self: Document<T>, callback: () -> ()) -> (),
 }
 
 --[=[

wally.lock

@@ -14,10 +14,10 @@ dependencies = []
 
 [[package]]
 name = "nezuo/lapis"
-version = "0.2.4"
-dependencies = [["Promise", "evaera/[email protected]"], ["DataStoreServiceMock", "nezuo/[email protected]"], ["Midori", "nezuo/[email protected].2"]]
+version = "0.2.5"
+dependencies = [["Promise", "evaera/[email protected]"], ["DataStoreServiceMock", "nezuo/[email protected]"], ["Midori", "nezuo/[email protected].3"]]
 
 [[package]]
 name = "nezuo/midori"
-version = "0.1.2"
+version = "0.1.3"
 dependencies = []

wally.toml

@@ -10,5 +10,5 @@ realm = "server"
 Promise = "evaera/[email protected]"
 
 [dev-dependencies]
-Midori = "nezuo/[email protected].2"
+Midori = "nezuo/[email protected].3"
 DataStoreServiceMock = "nezuo/[email protected]"