feat(api): auto-connect Flagsmith warehouse for new orgs via experimental_flags value

api/environments/models.py

@@ -184,6 +184,12 @@ class Meta:
     def create_feature_states(self) -> None:
         FeatureState.create_initial_feature_states_for_environment(environment=self)
 
+    @hook(AFTER_CREATE)  # type: ignore[misc]
+    def auto_connect_warehouse(self) -> None:
+        from experimentation.services import ensure_flagsmith_warehouse_connection
+
+        ensure_flagsmith_warehouse_connection(self)
+
     @hook(AFTER_UPDATE)  # type: ignore[misc]
     def clear_environment_cache(self) -> None:
         # TODO: this could rebuild the cache itself (using an async task)

api/experimentation/services.py

@@ -13,7 +13,7 @@
 from clickhouse_driver.util.helpers import parse_url
 from django.conf import settings
 from django.core.cache import cache
-from django.db import transaction
+from django.db import IntegrityError, transaction
 from django.db.models import Q
 from django.utils import timezone
 from flag_engine.segments.constants import ALL_RULE, PERCENTAGE_SPLIT
@@ -58,6 +58,7 @@
     ExperimentStatus,
     MetricAggregation,
     MetricDirection,
+    WarehouseConnection,
     WarehouseConnectionStatus,
     WarehouseDeliveryLog,
     WarehouseDeliveryOutcome,
@@ -91,7 +92,8 @@
 
     from clickhouse_connect.driver.client import Client as ClickHouseHTTPClient
 
-    from experimentation.models import Metric, WarehouseConnection
+    from environments.models import Environment
+    from experimentation.models import Metric
     from experimentation.types import ExposureGranularity
     from features.feature_states.models import FeatureValueType
     from features.models import FeatureStateValue
@@ -147,6 +149,47 @@ def is_experiment_feature_enabled(organisation: Organisation) -> bool:
     )
 
 
+def get_experiment_flag_config(
+    organisation: Organisation,
+) -> dict[str, object]:
+    if not is_experiment_feature_enabled(organisation):
+        return {}
+    raw = get_openfeature_client().get_string_value(
+        EXPERIMENT_FLAG,
+        default_value="{}",
+        evaluation_context=organisation.openfeature_evaluation_context,
+    )
+    try:
+        parsed = json.loads(raw)
+    except (json.JSONDecodeError, TypeError):
+        return {}
+    return parsed if isinstance(parsed, dict) else {}
+
+
+def ensure_flagsmith_warehouse_connection(
+    environment: Environment,
+) -> WarehouseConnection | None:
+    config = get_experiment_flag_config(environment.project.organisation)
+    if not config.get("auto_connect_warehouse"):
+        return None
+
+    if WarehouseConnection.objects.filter(
+        environment=environment,
+        deleted_at__isnull=True,
+    ).exists():
+        return None
+
+    try:
+        connection: WarehouseConnection = WarehouseConnection.objects.create(
+            environment=environment,
+            warehouse_type=WarehouseType.FLAGSMITH,
+            name="Flagsmith",
+        )
+        return connection
+    except IntegrityError:
+        return None
+
+
 @lru_cache(maxsize=2)
 def _get_clickhouse_client(
     send_receive_timeout: int = CLICKHOUSE_QUERY_TIMEOUT_SECONDS,

api/tests/unit/experimentation/test_services.py

@@ -3,6 +3,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+from django.db import IntegrityError
 from django.db.models import Q
 from flag_engine.segments.constants import PERCENTAGE_SPLIT
 from prometheus_client import REGISTRY
@@ -48,6 +49,7 @@
 from features.multivariate.models import MultivariateFeatureOption
 from features.value_types import STRING
 from features.versioning.dataclasses import MultivariateValueChangeSet
+from organisations.models import Organisation
 from segments.models import Condition, Segment, SegmentRule
 from users.models import FFAdminUser
 from util.mappers import map_environment_to_environment_document
@@ -2505,3 +2507,184 @@ def test_annotate_warehouse_event_stats__clickhouse_connection__queries_customer
     # Then
     get_client.assert_called_once()
     assert getattr(fresh_connection, "event_stats", None) == expected_stats
+
+
+def test_get_experiment_flag_config__flag_disabled__returns_empty(
+    organisation: Organisation,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = MagicMock()
+    mock_client.get_boolean_value.return_value = False
+    mocker.patch(
+        "experimentation.services.get_openfeature_client",
+        return_value=mock_client,
+    )
+
+    # When
+    result = services.get_experiment_flag_config(organisation)
+
+    # Then
+    assert result == {}
+    mock_client.get_string_value.assert_not_called()
+
+
+def test_get_experiment_flag_config__flag_enabled_with_valid_json__returns_parsed(
+    organisation: Organisation,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = MagicMock()
+    mock_client.get_boolean_value.return_value = True
+    mock_client.get_string_value.return_value = '{"auto_connect_warehouse": true}'
+    mocker.patch(
+        "experimentation.services.get_openfeature_client",
+        return_value=mock_client,
+    )
+
+    # When
+    result = services.get_experiment_flag_config(organisation)
+
+    # Then
+    assert result == {"auto_connect_warehouse": True}
+
+
+@pytest.mark.parametrize(
+    "raw_value",
+    ["not-json", "", None],
+    ids=["invalid-json", "empty-string", "none"],
+)
+def test_get_experiment_flag_config__flag_enabled_with_bad_value__returns_empty(
+    organisation: Organisation,
+    mocker: MockerFixture,
+    raw_value: str | None,
+) -> None:
+    # Given
+    mock_client = MagicMock()
+    mock_client.get_boolean_value.return_value = True
+    mock_client.get_string_value.return_value = raw_value
+    mocker.patch(
+        "experimentation.services.get_openfeature_client",
+        return_value=mock_client,
+    )
+
+    # When
+    result = services.get_experiment_flag_config(organisation)
+
+    # Then
+    assert result == {}
+
+
+def test_get_experiment_flag_config__flag_enabled_with_non_dict_json__returns_empty(
+    organisation: Organisation,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = MagicMock()
+    mock_client.get_boolean_value.return_value = True
+    mock_client.get_string_value.return_value = '["free"]'
+    mocker.patch(
+        "experimentation.services.get_openfeature_client",
+        return_value=mock_client,
+    )
+
+    # When
+    result = services.get_experiment_flag_config(organisation)
+
+    # Then
+    assert result == {}
+
+
+@pytest.mark.django_db()
+def test_ensure_flagsmith_warehouse_connection__auto_connect_disabled__returns_none(
+    environment: Environment,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch(
+        "experimentation.services.get_experiment_flag_config",
+        return_value={},
+    )
+
+    # When
+    result = services.ensure_flagsmith_warehouse_connection(environment)
+
+    # Then
+    assert result is None
+    assert not WarehouseConnection.objects.filter(
+        environment=environment,
+    ).exists()
+
+
+@pytest.mark.django_db()
+def test_ensure_flagsmith_warehouse_connection__auto_connect_enabled__creates_connection(
+    environment: Environment,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch(
+        "experimentation.services.get_experiment_flag_config",
+        return_value={"auto_connect_warehouse": True},
+    )
+
+    # When
+    result = services.ensure_flagsmith_warehouse_connection(environment)
+
+    # Then
+    assert result is not None
+    assert result.warehouse_type == WarehouseType.FLAGSMITH
+    assert result.name == "Flagsmith"
+    assert result.environment == environment
+
+
+@pytest.mark.django_db()
+def test_ensure_flagsmith_warehouse_connection__connection_already_exists__returns_none(
+    environment: Environment,
+    warehouse_connection: WarehouseConnection,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch(
+        "experimentation.services.get_experiment_flag_config",
+        return_value={"auto_connect_warehouse": True},
+    )
+
+    # When
+    result = services.ensure_flagsmith_warehouse_connection(environment)
+
+    # Then
+    assert result is None
+    assert (
+        WarehouseConnection.objects.filter(
+            environment=environment,
+            deleted_at__isnull=True,
+        ).count()
+        == 1
+    )
+
+
+def test_ensure_flagsmith_warehouse_connection__race_condition__handles_integrity_error(
+    environment: Environment,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch(
+        "experimentation.services.get_experiment_flag_config",
+        return_value={"auto_connect_warehouse": True},
+    )
+    mocker.patch.object(
+        WarehouseConnection.objects,
+        "filter",
+        return_value=MagicMock(exists=MagicMock(return_value=False)),
+    )
+    mocker.patch.object(
+        WarehouseConnection.objects,
+        "create",
+        side_effect=IntegrityError("duplicate"),
+    )
+
+    # When
+    result = services.ensure_flagsmith_warehouse_connection(environment)
+
+    # Then
+    assert result is None

api/tests/unit/experimentation/test_signals.py

@@ -1,7 +1,10 @@
+import pytest
 from pytest_mock import MockerFixture
 
 from environments.models import Environment, EnvironmentAPIKey
-from experimentation.models import WarehouseConnection
+from experimentation.models import WarehouseConnection, WarehouseType
+from organisations.models import Organisation
+from projects.models import Project
 
 
 def test_environment_api_key__created_with_warehouse__enqueues_write(
@@ -94,3 +97,55 @@ def test_environment_api_key__deleted_without_warehouse__does_not_enqueue(
 
     # Then
     mock_task.delay.assert_not_called()
+
+
+@pytest.mark.django_db()
+def test_environment_create__auto_connect_enabled__creates_warehouse_connection(
+    organisation: Organisation,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch(
+        "experimentation.services.get_experiment_flag_config",
+        return_value={"auto_connect_warehouse": True},
+    )
+    project: Project = Project.objects.create(
+        name="Test Project",
+        organisation=organisation,
+    )
+
+    # When
+    environment: Environment = Environment.objects.create(
+        name="Test Environment",
+        project=project,
+    )
+
+    # Then
+    connection = WarehouseConnection.objects.get(environment=environment)
+    assert connection.warehouse_type == WarehouseType.FLAGSMITH
+    assert connection.name == "Flagsmith"
+
+
+@pytest.mark.django_db()
+def test_environment_create__auto_connect_disabled__no_warehouse_connection(
+    organisation: Organisation,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch(
+        "experimentation.services.get_experiment_flag_config",
+        return_value={},
+    )
+    project: Project = Project.objects.create(
+        name="Test Project",
+        organisation=organisation,
+    )
+
+    # When
+    environment: Environment = Environment.objects.create(
+        name="Test Environment",
+        project=project,
+    )
+
+    # Then
+    assert not WarehouseConnection.objects.filter(environment=environment).exists()