feat: userId and customId specific feature gate overrides

statsig-io · Dec 8, 2024 · d9ecf8f · d9ecf8f
1 parent 2469701
commit d9ecf8f
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 3 deletions.
diff --git a/src/main/kotlin/com/statsig/sdk/Evaluator.kt b/src/main/kotlin/com/statsig/sdk/Evaluator.kt
@@ -38,7 +38,7 @@ internal class Evaluator(
         }
     }
     private val persistentStore: UserPersistentStorageHandler
-    private var gateOverrides: MutableMap<String, Boolean> = HashMap()
+    private var gateOverrides: MutableMap<String, MutableMap<String, Boolean>> = HashMap()
     private var configOverrides: MutableMap<String, Map<String, Any>> = HashMap()
     private var layerOverrides: MutableMap<String, Map<String, Any>> = HashMap()
     private var hashLookupTable: MutableMap<String, ULong> = HashMap()
@@ -148,7 +148,17 @@ internal class Evaluator(
     }
 
     fun overrideGate(gateName: String, gateValue: Boolean) {
-        gateOverrides[gateName] = gateValue
+        if (gateOverrides[gateName] == null) {
+            gateOverrides[gateName] = HashMap()
+        }
+        gateOverrides[gateName]?.set("", gateValue)
+    }
+
+    fun overrideGate(gateName: String, gateValue: Boolean, userId: String) {
+        if (gateOverrides[gateName] == null) {
+            gateOverrides[gateName] = HashMap()
+        }
+        gateOverrides[gateName]?.set(userId, gateValue)
     }
 
     fun overrideConfig(configName: String, configValue: Map<String, Any>) {
@@ -171,6 +181,10 @@ internal class Evaluator(
         gateOverrides.remove(gateName)
     }
 
+    fun removeGateOverride(gateName: String, userId: String) {
+        gateOverrides[gateName]?.remove(userId)
+    }
+
     fun getConfig(ctx: EvaluationContext, dynamicConfigName: String) {
         if (configOverrides.containsKey(dynamicConfigName)) {
             ctx.evaluation.jsonValue = configOverrides[dynamicConfigName] ?: mapOf<String, Any>()
@@ -308,7 +322,17 @@ internal class Evaluator(
     @JvmOverloads
     fun checkGate(ctx: EvaluationContext, gateName: String) {
         if (gateOverrides.containsKey(gateName)) {
-            val value = gateOverrides[gateName] ?: false
+            val userIds = mutableListOf<String>()
+            ctx.user.userID?.let { userIds.add(it) }
+            ctx.user.customIDs?.let { customIdMap -> userIds.addAll(customIdMap.values ) }
+            userIds.add("")
+
+            val value: Boolean = userIds
+                .stream()
+                .map { gateOverrides[gateName]?.get(it) }
+                .filter { it != null }
+                .findFirst().orElse(false) ?: false
+
             ctx.evaluation.booleanValue = value
             ctx.evaluation.jsonValue = value
             ctx.evaluation.evaluationDetails = createEvaluationDetails(EvaluationReason.LOCAL_OVERRIDE)

diff --git a/src/main/kotlin/com/statsig/sdk/Statsig.kt b/src/main/kotlin/com/statsig/sdk/Statsig.kt
@@ -300,6 +300,21 @@ class Statsig {
             statsigServer.overrideGate(gateName, gateValue)
         }
 
+        /**
+         * Sets a value to be returned for the given gate instead of the actual evaluated value.
+         *
+         * @param gateName The name of the gate to be overridden
+         * @param gateValue The value that will be returned
+         * @param userId The user ID to override the gate for
+         */
+        @JvmStatic
+        fun overrideGate(gateName: String, gateValue: Boolean, userId: String) {
+            if (!checkInitialized()) {
+                return
+            }
+            statsigServer.overrideGate(gateName, gateValue, userId)
+        }
+
         /**
          * Removes the given gate override
          *
@@ -312,6 +327,19 @@ class Statsig {
             }
         }
 
+        /**
+         * Removes the given gate override for the given user ID
+         *
+         * @param gateName
+         * @param userId
+         */
+        @JvmStatic
+        fun removeGateOverride(gateName: String, userId: String) {
+            if (checkInitialized()) {
+                statsigServer.removeGateOverride(gateName, userId)
+            }
+        }
+
         /**
          * Sets a value to be returned for the given dynamic config/experiment instead of the actual evaluated value.
          *

diff --git a/src/main/kotlin/com/statsig/sdk/StatsigServer.kt b/src/main/kotlin/com/statsig/sdk/StatsigServer.kt
@@ -85,8 +85,12 @@ sealed class StatsigServer {
 
     abstract fun overrideGate(gateName: String, gateValue: Boolean)
 
+    abstract fun overrideGate(gateName: String, gateValue: Boolean, userId: String)
+
     abstract fun removeGateOverride(gateName: String)
 
+    abstract fun removeGateOverride(gateName: String, userId: String)
+
     abstract fun overrideConfig(configName: String, configValue: Map<String, Any>)
 
     abstract fun removeConfigOverride(configName: String)
@@ -943,6 +947,16 @@ private class StatsigServerImpl() :
         }, { return@captureSync })
     }
 
+    override fun overrideGate(gateName: String, gateValue: Boolean, userId: String) {
+        if (!isSDKInitialized()) {
+            return
+        }
+        errorBoundary.captureSync("overrideGate", {
+            isSDKInitialized()
+            evaluator.overrideGate(gateName, gateValue, userId)
+        }, { return@captureSync })
+    }
+
     override fun removeGateOverride(gateName: String) {
         if (!isSDKInitialized()) {
             return
@@ -953,6 +967,16 @@ private class StatsigServerImpl() :
         }, { return@captureSync })
     }
 
+    override fun removeGateOverride(gateName: String, userId: String) {
+        if (!isSDKInitialized()) {
+            return
+        }
+        errorBoundary.captureSync("removeGateOverride", {
+            isSDKInitialized()
+            evaluator.removeGateOverride(gateName, userId)
+        }, { return@captureSync })
+    }
+
     override fun overrideConfig(configName: String, configValue: Map<String, Any>) {
         if (!isSDKInitialized()) {
             return

diff --git a/src/test/java/com/statsig/sdk/LocalOverridesTest.kt b/src/test/java/com/statsig/sdk/LocalOverridesTest.kt
@@ -39,6 +39,21 @@ class LocalOverridesTest {
         assertFalse(Statsig.checkGate(user, "override_me"))
     }
 
+    @Test
+    fun testGateOverridesWithUserId() = runBlocking {
+        users.forEach { user -> testGateOverridesWithUserIdHelper(user) }
+    }
+
+    private fun testGateOverridesWithUserIdHelper(user: StatsigUser) = runBlocking {
+        assertFalse(Statsig.checkGate(user, "override_me"))
+
+        Statsig.overrideGate("override_me", true, user.userID ?: user.customIDs?.get("customID") ?: "")
+        assertTrue(Statsig.checkGate(user, "override_me"))
+
+        Statsig.removeGateOverride("override_me")
+        assertFalse(Statsig.checkGate(user, "override_me"))
+    }
+
     @Test
     fun testConfigOverrides() = runBlocking {
         users.forEach { user -> testConfigOverridesHelper(user) }