feat(experimentation): sync warehouse connection environment keys to ingestion redis (#7588)

Author: gagantrivedi

api/app/settings/common.py

@@ -784,6 +784,9 @@
     "REDIS_CLUSTER_READ_FROM_REPLICAS", default=True
 )
 
+# Redis Cluster URL used to communicate with the event ingestion server.
+INGESTION_REDIS_URL = env.str("INGESTION_REDIS_URL", default="")
+
 CACHES = {
     "default": {
         "BACKEND": "django.core.cache.backends.locmem.LocMemCache",

api/experimentation/ingestion_sync_service.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+from django.conf import settings
+from redis.cluster import RedisCluster
+
+INGESTION_ENVIRONMENT_KEY_PREFIX = "experimentation:environment_keys:"
+
+
+SOCKET_TIMEOUT = 1
+
+
+@lru_cache(maxsize=1)
+def _get_client() -> RedisCluster:
+    return RedisCluster.from_url(  # type: ignore[no-untyped-call,no-any-return]
+        settings.INGESTION_REDIS_URL,
+        socket_timeout=SOCKET_TIMEOUT,
+        socket_keepalive=True,
+    )
+
+
+def set_environment_key(environment_api_key: str) -> None:
+    key = f"{INGESTION_ENVIRONMENT_KEY_PREFIX}{environment_api_key}"
+    _get_client().set(key, "")
+
+
+def delete_environment_key(environment_api_key: str) -> None:
+    key = f"{INGESTION_ENVIRONMENT_KEY_PREFIX}{environment_api_key}"
+    _get_client().delete(key)

api/experimentation/models.py

@@ -1,8 +1,17 @@
 from django.db import models
-from django_lifecycle import LifecycleModelMixin  # type: ignore[import-untyped]
+from django_lifecycle import (  # type: ignore[import-untyped]
+    AFTER_CREATE,
+    AFTER_DELETE,
+    LifecycleModelMixin,
+    hook,
+)
 
 from core.models import SoftDeleteExportableModel
 from environments.models import Environment
+from experimentation.tasks import (
+    add_environment_key_to_ingestion,
+    delete_environment_key_from_ingestion,
+)
 
 
 class WarehouseType(models.TextChoices):
@@ -47,3 +56,15 @@ class Meta:
                 name="unique_active_warehouse_per_env",
             ),
         ]
+
+    @hook(AFTER_CREATE)  # type: ignore[misc]
+    def sync_to_ingestion_on_create(self) -> None:
+        add_environment_key_to_ingestion.delay(
+            kwargs={"environment_api_key": self.environment.api_key},
+        )
+
+    @hook(AFTER_DELETE)  # type: ignore[misc]
+    def sync_to_ingestion_on_delete(self) -> None:
+        delete_environment_key_from_ingestion.delay(
+            kwargs={"environment_api_key": self.environment.api_key},
+        )

api/experimentation/tasks.py

@@ -0,0 +1,13 @@
+from task_processor.decorators import register_task_handler
+
+from experimentation import ingestion_sync_service
+
+
+@register_task_handler()
+def add_environment_key_to_ingestion(environment_api_key: str) -> None:
+    ingestion_sync_service.set_environment_key(environment_api_key)
+
+
+@register_task_handler()
+def delete_environment_key_from_ingestion(environment_api_key: str) -> None:
+    ingestion_sync_service.delete_environment_key(environment_api_key)

api/tests/unit/experimentation/conftest.py

@@ -1,10 +1,18 @@
 import pytest
 from django.urls import reverse
+from pytest_mock import MockerFixture
 
 from environments.models import Environment
+from experimentation import ingestion_sync_service
 from experimentation.models import WarehouseConnection, WarehouseType
 
 
+@pytest.fixture(autouse=True)
+def mock_ingestion_redis_client(mocker: MockerFixture) -> None:
+    ingestion_sync_service._get_client.cache_clear()
+    mocker.patch("experimentation.ingestion_sync_service.RedisCluster.from_url")
+
+
 @pytest.fixture()
 def warehouse_connection(environment: Environment) -> WarehouseConnection:
     connection: WarehouseConnection = WarehouseConnection.objects.create(

api/tests/unit/experimentation/test_ingestion_sync_service.py

@@ -0,0 +1,98 @@
+import pytest
+from pytest_mock import MockerFixture
+from redis.exceptions import RedisError
+
+from experimentation import ingestion_sync_service
+
+
+def test_get_client__configured_url__builds_redis_cluster_with_socket_options(
+    mocker: MockerFixture,
+    settings: object,
+) -> None:
+    # Given
+    settings.INGESTION_REDIS_URL = "redis://ingestion:6379"  # type: ignore[attr-defined]
+    mock_from_url = mocker.patch(
+        "experimentation.ingestion_sync_service.RedisCluster.from_url",
+    )
+
+    # When
+    client = ingestion_sync_service._get_client()
+
+    # Then
+    mock_from_url.assert_called_once_with(
+        "redis://ingestion:6379",
+        socket_timeout=ingestion_sync_service.SOCKET_TIMEOUT,
+        socket_keepalive=True,
+    )
+    assert client is mock_from_url.return_value
+
+
+def test_set_environment_key__valid_key__writes_to_redis(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = mocker.Mock()
+    mocker.patch(
+        "experimentation.ingestion_sync_service._get_client",
+        return_value=mock_client,
+    )
+
+    # When
+    ingestion_sync_service.set_environment_key("test-env-key-001")
+
+    # Then
+    mock_client.set.assert_called_once_with(
+        "experimentation:environment_keys:test-env-key-001",
+        "",
+    )
+
+
+def test_delete_environment_key__valid_key__deletes_from_redis(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = mocker.Mock()
+    mocker.patch(
+        "experimentation.ingestion_sync_service._get_client",
+        return_value=mock_client,
+    )
+
+    # When
+    ingestion_sync_service.delete_environment_key("test-env-key-001")
+
+    # Then
+    mock_client.delete.assert_called_once_with(
+        "experimentation:environment_keys:test-env-key-001",
+    )
+
+
+def test_set_environment_key__redis_error__propagates(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = mocker.Mock()
+    mock_client.set.side_effect = RedisError("boom")
+    mocker.patch(
+        "experimentation.ingestion_sync_service._get_client",
+        return_value=mock_client,
+    )
+
+    # When / Then
+    with pytest.raises(RedisError, match="boom"):
+        ingestion_sync_service.set_environment_key("test-env-key-001")
+
+
+def test_delete_environment_key__redis_error__propagates(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = mocker.Mock()
+    mock_client.delete.side_effect = RedisError("boom")
+    mocker.patch(
+        "experimentation.ingestion_sync_service._get_client",
+        return_value=mock_client,
+    )
+
+    # When / Then
+    with pytest.raises(RedisError, match="boom"):
+        ingestion_sync_service.delete_environment_key("test-env-key-001")

api/tests/unit/experimentation/test_models.py

@@ -0,0 +1,45 @@
+from pytest_mock import MockerFixture
+
+from environments.models import Environment
+from experimentation.models import WarehouseConnection, WarehouseType
+
+
+def test_warehouse_connection__after_create__enqueues_ingestion_add_task(
+    environment: Environment,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_task = mocker.patch(
+        "experimentation.models.add_environment_key_to_ingestion",
+    )
+
+    # When
+    WarehouseConnection.objects.create(
+        environment=environment,
+        warehouse_type=WarehouseType.FLAGSMITH,
+        name="warehouse",
+    )
+
+    # Then
+    mock_task.delay.assert_called_once_with(
+        kwargs={"environment_api_key": environment.api_key},
+    )
+
+
+def test_warehouse_connection__after_delete__enqueues_ingestion_delete_task(
+    warehouse_connection: WarehouseConnection,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_task = mocker.patch(
+        "experimentation.models.delete_environment_key_from_ingestion",
+    )
+    environment_api_key = warehouse_connection.environment.api_key
+
+    # When
+    warehouse_connection.delete()
+
+    # Then
+    mock_task.delay.assert_called_once_with(
+        kwargs={"environment_api_key": environment_api_key},
+    )

api/tests/unit/experimentation/test_tasks.py

@@ -0,0 +1,36 @@
+from pytest_mock import MockerFixture
+
+from experimentation.tasks import (
+    add_environment_key_to_ingestion,
+    delete_environment_key_from_ingestion,
+)
+
+
+def test_add_environment_key_to_ingestion__valid_key__calls_service(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_set = mocker.patch(
+        "experimentation.tasks.ingestion_sync_service.set_environment_key",
+    )
+
+    # When
+    add_environment_key_to_ingestion(environment_api_key="test-env-key-001")
+
+    # Then
+    mock_set.assert_called_once_with("test-env-key-001")
+
+
+def test_delete_environment_key_from_ingestion__valid_key__calls_service(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_delete = mocker.patch(
+        "experimentation.tasks.ingestion_sync_service.delete_environment_key",
+    )
+
+    # When
+    delete_environment_key_from_ingestion(environment_api_key="test-env-key-001")
+
+    # Then
+    mock_delete.assert_called_once_with("test-env-key-001")

infrastructure/aws/production/ecs-task-definition-admin-api.json

@@ -228,6 +228,10 @@
                 {
                     "name": "DASHBOARD_ENDPOINTS_SENTRY_TRACE_SAMPLE_RATE",
                     "value": "0.5"
+                },
+                {
+                    "name": "INGESTION_REDIS_URL",
+                    "value": "rediss://serverless-redis-cache-c4q8sw.serverless.euw2.cache.amazonaws.com:6379"
                 }
             ],
             "secrets": [

infrastructure/aws/production/ecs-task-definition-task-processor.json

@@ -192,6 +192,10 @@
                 {
                     "name": "LOG_FORMAT",
                     "value": "json"
+                },
+                {
+                    "name": "INGESTION_REDIS_URL",
+                    "value": "rediss://serverless-redis-cache-c4q8sw.serverless.euw2.cache.amazonaws.com:6379"
                 }
             ],
             "secrets": [