Implemented optional data validation

DataStore2/init.lua

@@ -179,13 +179,43 @@ function DataStore:GetTableAsync(default, ...)
 	end)
 end
 
+function DataStore:SetValidator(validator)
+	assert(
+		type(validator) == "function",
+		"function expected, got " .. typeof(validator)
+	)
+	self.validator = validator
+end
+
+local function assertValidatorWithDefaultError(validator, input, defaultError)
+	local isValid, message = validator(input)
+	if not isValid then
+		error(message or defaultError) 
+	end
+end
+
 function DataStore:Set(value, _dontCallOnUpdate)
+	if self.validator ~= nil then
+		assertValidatorWithDefaultError(
+			self.validator,
+			value,
+			"Attempted to set data store to an invalid value during :Set"
+		)
+	end
 	self.value = clone(value)
 	self:_Update(_dontCallOnUpdate)
 end
 
 function DataStore:Update(updateFunc)
-	self.value = updateFunc(self.value)
+	local updateFuncReturn = updateFunc(self.value)
+	if self.validator ~= nil then
+		assertValidatorWithDefaultError(
+			self.validator,
+			updateFuncReturn,
+			"Attempted to set data store to an invalid value during :Update"
+		)
+	end
+	self.value = updateFuncReturn
 	self:_Update()
 end
 

Tests/tests/DataStore2.spec.lua

@@ -79,6 +79,47 @@ return function()
 				expect(DataStore2(nonNilKey, fakePlayer):Get("badDefault")).to.equal("abc")
 			end)
 
+			it("should validate Set", function()
+				local dataStore = DataStore2(UUID(), fakePlayer)
+				local function testValidator(dataToValidate)
+					if dataToValidate == "yepp" then
+						return true
+					elseif dataToValidate == "definitelyNot" then
+						return false, "A validation error message"
+					end
+					return false
+				end
+				dataStore:SetValidator(testValidator)
+				expect(dataStore:Set("nope")).to.throw("Attempted to set datastore to an invalid value")
+				expect(dataStore:Set("definitelyNot")).to.throw("A validation error message")
+				expect(dataStore:Set("yepp")).to.be.ok()
+			end)
+
+			it("should validate Update", function()
+				local dataStore = DataStore2(UUID(), fakePlayer)
+				local function testValidator(dataToValidate)
+					if dataToValidate == "yepp" then
+						return true
+					elseif dataToValidate == "definitelyNot" then
+						return false, "A validation error message"
+					end
+					return false
+				end
+				dataStore:SetValidator(testValidator)
+
+				expect(dataStore:Update(function()
+					return "nope"
+				end)).to.throw("Attempted to set datastore to an invalid value")
+
+				expect(dataStore:Update(function()
+					return "definitelyNot"
+				end)).to.throw("A validation error message")
+
+				expect(dataStore:Update(function()
+					return "yepp"
+				end)).to.be.ok()
+			end)
+
 			it("should set", function()
 				local dataStore = DataStore2(UUID(), fakePlayer)
 				dataStore:Set(1)
@@ -562,4 +603,4 @@ return function()
 			expect(called2).to.equal(1)
 		end)
 	end)
-end
+end