Add an admin API endpoint to support per-user feature flags (#15344)

2 files changed, 77 insertions, 0 deletions

diff --git a/synapse/storage/databases/main/__init__.py b/synapse/storage/databases/main/__init__.py
index dc3948c170..0032a92f49 100644
--- a/synapse/storage/databases/main/__init__.py
+++ b/synapse/storage/databases/main/__init__.py
@@ -43,6 +43,7 @@ from .event_federation import EventFederationStore
 from .event_push_actions import EventPushActionsStore
 from .events_bg_updates import EventsBackgroundUpdatesStore
 from .events_forward_extremities import EventForwardExtremitiesStore
+from .experimental_features import ExperimentalFeaturesStore
 from .filtering import FilteringWorkerStore
 from .keys import KeyStore
 from .lock import LockStore
@@ -82,6 +83,7 @@ logger = logging.getLogger(__name__)
 
 class DataStore(
     EventsBackgroundUpdatesStore,
+    ExperimentalFeaturesStore,
     DeviceStore,
     RoomMemberStore,
     RoomStore,
diff --git a/synapse/storage/databases/main/experimental_features.py b/synapse/storage/databases/main/experimental_features.py
new file mode 100644
index 0000000000..cf3226ae5a
--- /dev/null
+++ b/synapse/storage/databases/main/experimental_features.py
@@ -0,0 +1,75 @@
+# Copyright 2023 The Matrix.org Foundation C.I.C
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import TYPE_CHECKING, Dict
+
+from synapse.storage.database import DatabasePool, LoggingDatabaseConnection
+from synapse.storage.databases.main import CacheInvalidationWorkerStore
+from synapse.types import StrCollection
+from synapse.util.caches.descriptors import cached
+
+if TYPE_CHECKING:
+    from synapse.rest.admin.experimental_features import ExperimentalFeature
+    from synapse.server import HomeServer
+
+
+class ExperimentalFeaturesStore(CacheInvalidationWorkerStore):
+    def __init__(
+        self,
+        database: DatabasePool,
+        db_conn: LoggingDatabaseConnection,
+        hs: "HomeServer",
+    ) -> None:
+        super().__init__(database, db_conn, hs)
+
+    @cached()
+    async def list_enabled_features(self, user_id: str) -> StrCollection:
+        """
+        Checks to see what features are enabled for a given user
+        Args:
+            user:
+                the user to be queried on
+        Returns:
+            the features currently enabled for the user
+        """
+        enabled = await self.db_pool.simple_select_list(
+            "per_user_experimental_features",
+            {"user_id": user_id, "enabled": True},
+            ["feature"],
+        )
+
+        return [feature["feature"] for feature in enabled]
+
+    async def set_features_for_user(
+        self,
+        user: str,
+        features: Dict["ExperimentalFeature", bool],
+    ) -> None:
+        """
+        Enables or disables features for a given user
+        Args:
+            user:
+                the user for whom to enable/disable the features
+            features:
+                pairs of features and True/False for whether the feature should be enabled
+        """
+        for feature, enabled in features.items():
+            await self.db_pool.simple_upsert(
+                table="per_user_experimental_features",
+                keyvalues={"feature": feature, "user_id": user},
+                values={"enabled": enabled},
+                insertion_values={"user_id": user, "feature": feature},
+            )
+
+            await self.invalidate_cache_and_stream("list_enabled_features", (user,))