diff --git a/api/environments/models.py b/api/environments/models.py index 80f79ef08889..feb5e89fc971 100644 --- a/api/environments/models.py +++ b/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) diff --git a/api/experimentation/services.py b/api/experimentation/services.py index 514246996a84..8496f9c3feaa 100644 --- a/api/experimentation/services.py +++ b/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, diff --git a/api/tests/unit/experimentation/test_services.py b/api/tests/unit/experimentation/test_services.py index 9934d65224c3..e318751dd05e 100644 --- a/api/tests/unit/experimentation/test_services.py +++ b/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 diff --git a/api/tests/unit/experimentation/test_signals.py b/api/tests/unit/experimentation/test_signals.py index b152952ea0fe..9c79f428108e 100644 --- a/api/tests/unit/experimentation/test_signals.py +++ b/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() diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index 8c508716727e..fe00b0203b8f 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md @@ -804,7 +804,7 @@ Attributes: ### `warehouse.connection.connected` Logged at `info` from: - - `api/experimentation/services.py:1147` + - `api/experimentation/services.py:1190` Attributes: - `environment.id` @@ -813,8 +813,8 @@ Attributes: ### `warehouse.connection.event_names_failed` Logged at `warning` from: - - `api/experimentation/services.py:224` - - `api/experimentation/services.py:1247` + - `api/experimentation/services.py:267` + - `api/experimentation/services.py:1290` Attributes: - `environment.id` @@ -824,7 +824,7 @@ Attributes: ### `warehouse.connection.event_stats_failed` Logged at `warning` from: - - `api/experimentation/services.py:1210` + - `api/experimentation/services.py:1253` Attributes: - `environment.id` @@ -833,7 +833,7 @@ Attributes: ### `warehouse.connection.test_event_sent` Logged at `info` from: - - `api/experimentation/services.py:921` + - `api/experimentation/services.py:964` Attributes: - `environment.id` @@ -842,7 +842,7 @@ Attributes: ### `warehouse.connection.verification_failed` Logged at `warning` from: - - `api/experimentation/services.py:1122` + - `api/experimentation/services.py:1165` Attributes: - `environment.id` @@ -852,7 +852,7 @@ Attributes: ### `warehouse.connection.verification_succeeded` Logged at `info` from: - - `api/experimentation/services.py:1132` + - `api/experimentation/services.py:1175` Attributes: - `environment.id` @@ -861,7 +861,7 @@ Attributes: ### `warehouse.delivery.all_objects_rejected` Logged at `error` from: - - `api/experimentation/services.py:1077` + - `api/experimentation/services.py:1120` Attributes: - `connection.id` @@ -872,7 +872,7 @@ Attributes: ### `warehouse.delivery.budget_exhausted` Logged at `info` from: - - `api/experimentation/services.py:966` + - `api/experimentation/services.py:1009` Attributes: - `connection.id` @@ -883,7 +883,7 @@ Attributes: ### `warehouse.delivery.completed` Logged at `info` from: - - `api/experimentation/services.py:1087` + - `api/experimentation/services.py:1130` Attributes: - `connection.id` @@ -896,7 +896,7 @@ Attributes: ### `warehouse.delivery.failed` Logged at `error` from: - - `api/experimentation/services.py:1060` + - `api/experimentation/services.py:1103` Attributes: - `connection.id` @@ -907,7 +907,7 @@ Attributes: ### `warehouse.delivery.object_rejected` Logged at `error` from: - - `api/experimentation/services.py:995` + - `api/experimentation/services.py:1038` Attributes: - `connection.id` @@ -919,7 +919,7 @@ Attributes: ### `warehouse.srm.overallocated` Logged at `error` from: - - `api/experimentation/services.py:517` + - `api/experimentation/services.py:560` Attributes: - `environment.id` @@ -929,7 +929,7 @@ Attributes: ### `warehouse.srm.unkeyed_variant` Logged at `error` from: - - `api/experimentation/services.py:503` + - `api/experimentation/services.py:546` Attributes: - `environment.id`