Merge pull request #26 from nezuo/set-user-ids

Add APIs to set the user ids for a document

src/Collection.lua

@@ -35,24 +35,29 @@ end
 --[=[
 	Loads the document with `key`, migrates it, and session locks it.
 
+	If specified, the document's `DataStoreKeyInfo:GetUserIds()` will be set to `defaultUserIds` if the document has
+	never been loaded.
+
 	@param key string
+	@param defaultUserIds {number}?
 	@return Promise<Document>
 ]=]
-function Collection:load(key)
+function Collection:load(key, defaultUserIds)
 	if self.openDocuments[key] == nil then
 		local lockId = HttpService:GenerateGUID(false)
 
 		self.openDocuments[key] = self
 			.data
 			:load(self.dataStore, key, function(value, keyInfo)
 				if value == nil then
-					return "succeed",
-						{
-							compressionScheme = "None",
-							migrationVersion = #self.options.migrations,
-							lockId = lockId,
-							data = self.options.defaultData,
-						}
+					local data = {
+						compressionScheme = "None",
+						migrationVersion = #self.options.migrations,
+						lockId = lockId,
+						data = self.options.defaultData,
+					}
+
+					return "succeed", data, defaultUserIds
 				end
 
 				if value.migrationVersion > #self.options.migrations then
@@ -70,15 +75,16 @@ function Collection:load(key)
 				local migrated = Migration.migrate(self.options.migrations, value.migrationVersion, decompressed)
 				local scheme, compressed = Compression.compress(migrated)
 
-				return "succeed",
-					{
-						compressionScheme = scheme,
-						migrationVersion = #self.options.migrations,
-						lockId = lockId,
-						data = compressed,
-					}
+				local data = {
+					compressionScheme = scheme,
+					migrationVersion = #self.options.migrations,
+					lockId = lockId,
+					data = compressed,
+				}
+
+				return "succeed", data, keyInfo:GetUserIds(), keyInfo:GetMetadata()
 			end)
-			:andThen(function(value)
+			:andThen(function(value, keyInfo)
 				local data = Compression.decompress(value.compressionScheme, value.data)
 
 				local ok, message = self.options.validate(data)
@@ -88,7 +94,7 @@ function Collection:load(key)
 
 				freezeDeep(data)
 
-				local document = Document.new(self, key, self.options.validate, lockId, data)
+				local document = Document.new(self, key, self.options.validate, lockId, data, keyInfo:GetUserIds())
 
 				self.autoSave:addDocument(document)
 

src/Data/Throttle.lua

@@ -4,14 +4,16 @@ local Promise = require(script.Parent.Parent.Parent.Promise)
 
 local function updateAsync(request)
 	return Promise.new(function(resolve)
-		local result, transformed
-
+		local resultOutside, transformedOutside, keyInfo
 		local ok, err = pcall(function()
-			request.dataStore:UpdateAsync(request.key, function(...)
-				result, transformed = request.transform(...)
+			_, keyInfo = request.dataStore:UpdateAsync(request.key, function(...)
+				local result, transformed, userIds = request.transform(...)
+
+				resultOutside = result
+				transformedOutside = transformed
 
 				if result == "succeed" then
-					return transformed
+					return transformed, userIds
 				else
 					return nil
 				end
@@ -21,7 +23,7 @@ local function updateAsync(request)
 		if not ok then
 			resolve("retry", err)
 		else
-			resolve(result, transformed)
+			resolve(resultOutside, transformedOutside, keyInfo)
 		end
 	end)
 end
@@ -59,10 +61,10 @@ function Throttle:start()
 				continue
 			end
 
-			local promise = updateAsync(request):andThen(function(result, value)
+			local promise = updateAsync(request):andThen(function(result, value, keyInfo)
 				if result == "succeed" then
 					request.attempts = 0
-					request.resolve(value)
+					request.resolve(value, keyInfo)
 				elseif result == "fail" then
 					request.attempts = 0
 					request.reject(`DataStoreFailure({value})`)

src/Document.lua

@@ -8,13 +8,14 @@ local Promise = require(script.Parent.Parent.Promise)
 local Document = {}
 Document.__index = Document
 
-function Document.new(collection, key, validate, lockId, data)
+function Document.new(collection, key, validate, lockId, data, userIds)
 	return setmetatable({
 		collection = collection,
 		key = key,
 		validate = validate,
 		lockId = lockId,
 		data = data,
+		userIds = userIds,
 		closed = false,
 		callingBeforeClose = false,
 	}, Document)
@@ -47,6 +48,34 @@ function Document:write(data)
 	self.data = data
 end
 
+--[=[
+	Adds a user id to the document's `DataStoreKeyInfo:GetUserIds()`. The change won't apply until the document is saved or closed.
+
+	If the user id is already associated with the document the method won't do anything.
+
+	@param userId number
+]=]
+function Document:addUserId(userId)
+	if table.find(self.userIds, userId) == nil then
+		table.insert(self.userIds, userId)
+	end
+end
+
+--[=[
+	Removes a user id from the document's `DataStoreKeyInfo:GetUserIds()`. The change won't apply until the document is saved or closed.
+
+	If the user id is not associated with the document the method won't do anything.
+
+	@param userId number
+]=]
+function Document:removeUserId(userId)
+	local index = table.find(self.userIds, userId)
+
+	if index ~= nil then
+		table.remove(self.userIds, index)
+	end
+end
+
 --[=[
 	Saves the document's data. If the save is throttled and you call it multiple times, it will save only once with the latest data.
 
@@ -69,7 +98,7 @@ function Document:save()
 		value.compressionScheme = scheme
 		value.data = compressed
 
-		return "succeed", value
+		return "succeed", value, self.userIds
 	end)
 end
 
@@ -126,7 +155,7 @@ function Document:close()
 				value.data = compressed
 				value.lockId = nil
 
-				return "succeed", value
+				return "succeed", value, self.userIds
 			end)
 		end)
 end

src/init.lua

@@ -25,12 +25,14 @@ export type CollectionOptions<T> = {
 }
 
 export type Collection<T> = {
-	load: (self: Collection<T>, key: string) -> PromiseTypes.TypedPromise<Document<T>>,
+	load: (self: Collection<T>, key: string, defaultUserIds: { number }?) -> PromiseTypes.TypedPromise<Document<T>>,
 }
 
 export type Document<T> = {
 	read: (self: Document<T>) -> T,
 	write: (self: Document<T>, T) -> (),
+	addUserId: (self: Document<T>, userId: number) -> (),
+	removeUserId: (self: Document<T>, userId: number) -> (),
 	save: (self: Document<T>) -> PromiseTypes.TypedPromise<()>,
 	close: (self: Document<T>) -> PromiseTypes.TypedPromise<()>,
 	beforeClose: (self: Document<T>, callback: () -> ()) -> (),

src/init.test.lua

@@ -41,6 +41,22 @@ return function(x)
 		context.read = function(name, key)
 			return dataStoreService.dataStores[name]["global"].data[key]
 		end
+
+		context.expectUserIds = function(name, key, targetUserIds)
+			local keyInfo = dataStoreService.dataStores[name]["global"].keyInfos[key]
+
+			local currentUserIds = if keyInfo ~= nil then keyInfo:GetUserIds() else {}
+
+			if #currentUserIds ~= #targetUserIds then
+				error("Incorrect user ids length")
+			end
+
+			for index, value in targetUserIds do
+				if currentUserIds[index] ~= value then
+					error("Invalid user id")
+				end
+			end
+		end
 	end)
 
 	x.test("throws when setting invalid config key", function(context)
@@ -96,7 +112,7 @@ return function(x)
 
 	x.test("should session lock the document", function(context)
 		local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
-		local document = collection:load("doc", DEFAULT_OPTIONS):expect()
+		local document = collection:load("doc"):expect()
 
 		local otherLapis = Internal.new(false)
 		otherLapis.setConfig({ dataStoreService = context.dataStoreService, loadAttempts = 1 })
@@ -122,7 +138,7 @@ return function(x)
 
 	x.test("load should retry when document is session locked", function(context)
 		local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
-		local document = collection:load("doc", DEFAULT_OPTIONS):expect()
+		local document = collection:load("doc"):expect()
 
 		local otherLapis = Internal.new(false)
 		otherLapis.setConfig({
@@ -259,4 +275,60 @@ return function(x)
 
 		assert(coroutine.status(thread) == "dead", "")
 	end)
+
+	x.nested("user ids", function()
+		x.test("it uses defaultUserIds on first load", function(context)
+			local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
+
+			local document = collection:load("document", { 123 }):expect()
+			context.expectUserIds("collection", "document", { 123 })
+			document:close():expect()
+			context.expectUserIds("collection", "document", { 123 })
+
+			-- Since the document has already been created, the defaultUserIds should not override the saved ones.
+			document = collection:load("document", { 321 }):expect()
+			context.expectUserIds("collection", "document", { 123 })
+			document:close():expect()
+			context.expectUserIds("collection", "document", { 123 })
+		end)
+
+		x.test("adds new user ids", function(context)
+			local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
+
+			local document = collection:load("document", {}):expect()
+
+			document:addUserId(111)
+			document:addUserId(111) -- It should not add this user id twice.
+			document:addUserId(222)
+
+			context.expectUserIds("collection", "document", {})
+
+			document:save():expect()
+
+			context.expectUserIds("collection", "document", { 111, 222 })
+
+			document:close():expect()
+
+			context.expectUserIds("collection", "document", { 111, 222 })
+		end)
+
+		x.test("removes new user ids", function(context)
+			local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
+
+			local document = collection:load("document", { 333, 444, 555 }):expect()
+
+			document:removeUserId(111) -- It should do nothing if the user id doesn't exist.
+			document:removeUserId(444)
+
+			context.expectUserIds("collection", "document", { 333, 444, 555 })
+
+			document:save():expect()
+
+			context.expectUserIds("collection", "document", { 333, 555 })
+
+			document:close():expect()
+
+			context.expectUserIds("collection", "document", { 333, 555 })
+		end)
+	end)
 end

wally.lock

@@ -9,13 +9,13 @@ dependencies = []
 
 [[package]]
 name = "nezuo/data-store-service-mock"
-version = "0.3.0"
+version = "0.3.5"
 dependencies = []
 
 [[package]]
 name = "nezuo/lapis"
 version = "0.2.5"
-dependencies = [["Promise", "evaera/[email protected]"], ["DataStoreServiceMock", "nezuo/[email protected].0"], ["Midori", "nezuo/[email protected]"]]
+dependencies = [["Promise", "evaera/[email protected]"], ["DataStoreServiceMock", "nezuo/[email protected].5"], ["Midori", "nezuo/[email protected]"]]
 
 [[package]]
 name = "nezuo/midori"

wally.toml

@@ -11,4 +11,4 @@ Promise = "evaera/[email protected]"
 
 [dev-dependencies]
 Midori = "nezuo/[email protected]"
-DataStoreServiceMock = "nezuo/[email protected].0"
+DataStoreServiceMock = "nezuo/[email protected].5"