From a3882f9cc623eaa23178db07e116c9f7fb80dcfe Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Fri, 17 Jul 2026 15:22:20 +0530 Subject: [PATCH] feat(experimentation): add per-organisation ingestion infrastructure service --- api/app/settings/common.py | 8 + api/experimentation/dataclasses.py | 6 + .../ingestion_infra_service.py | 187 +++++++++++++++ .../test_ingestion_infra_service.py | 223 ++++++++++++++++++ .../observability/_events-catalogue.md | 19 ++ 5 files changed, 443 insertions(+) create mode 100644 api/experimentation/ingestion_infra_service.py create mode 100644 api/tests/unit/experimentation/test_ingestion_infra_service.py diff --git a/api/app/settings/common.py b/api/app/settings/common.py index c5ff83281bdc..3b63102d67b6 100644 --- a/api/app/settings/common.py +++ b/api/app/settings/common.py @@ -806,6 +806,14 @@ # Redis Cluster URL used to communicate with the event ingestion server. INGESTION_REDIS_URL = env.str("INGESTION_REDIS_URL", default="") +# Name prefix for per-organisation experiment events S3 buckets. +INGESTION_EVENTS_BUCKET_PREFIX = env.str("INGESTION_EVENTS_BUCKET_PREFIX", default="") + +# ARN of the IAM role Firehose assumes to deliver experiment events to S3. +INGESTION_FIREHOSE_DELIVERY_ROLE_ARN = env.str( + "INGESTION_FIREHOSE_DELIVERY_ROLE_ARN", default="" +) + # DSN for the ClickHouse instance holding ingested experimentation events. EXPERIMENTATION_CLICKHOUSE_URL = env.str("EXPERIMENTATION_CLICKHOUSE_URL", default=None) diff --git a/api/experimentation/dataclasses.py b/api/experimentation/dataclasses.py index d8db1624ed34..1766ea5aeb44 100644 --- a/api/experimentation/dataclasses.py +++ b/api/experimentation/dataclasses.py @@ -80,3 +80,9 @@ class MetricResult: class ResultsSummary: srm_p_value: float | None metrics: list[MetricResult] + + +@dataclass(frozen=True) +class IngestionInfrastructure: + bucket_name: str + stream_name: str diff --git a/api/experimentation/ingestion_infra_service.py b/api/experimentation/ingestion_infra_service.py new file mode 100644 index 000000000000..30006b2f88c8 --- /dev/null +++ b/api/experimentation/ingestion_infra_service.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import typing +from functools import lru_cache + +import boto3 +import structlog +from botocore.exceptions import ClientError +from django.conf import settings +from django.core.exceptions import ImproperlyConfigured + +from experimentation.dataclasses import IngestionInfrastructure + +if typing.TYPE_CHECKING: + from typing import Any + +logger = structlog.get_logger("experimentation") + +AWS_REGION = "eu-west-2" +STREAM_NAME_PREFIX = "events-ingestion-org-" + +# S3 object keys are namespaced per environment via Firehose dynamic +# partitioning on each record's environment_key field. +EVENTS_PREFIX = ( + "events/env_key=!{partitionKeyFromQuery:env_key}/" + "year=!{timestamp:yyyy}/month=!{timestamp:MM}/" + "day=!{timestamp:dd}/hour=!{timestamp:HH}/" +) +ERRORS_PREFIX = ( + "errors/!{firehose:error-output-type}/" + "year=!{timestamp:yyyy}/month=!{timestamp:MM}/" + "day=!{timestamp:dd}/hour=!{timestamp:HH}/" +) + +# Firehose requires buffers of at least 64 MiB when dynamic partitioning is +# enabled. +BUFFERING_SIZE_MB = 64 +BUFFERING_INTERVAL_SECONDS = 300 +DYNAMIC_PARTITIONING_RETRY_SECONDS = 300 +ERROR_OBJECT_EXPIRATION_DAYS = 30 + + +@lru_cache(maxsize=1) +def _get_s3_client() -> "Any": + return boto3.client("s3", region_name=AWS_REGION) + + +@lru_cache(maxsize=1) +def _get_firehose_client() -> "Any": + return boto3.client("firehose", region_name=AWS_REGION) + + +def get_bucket_name(organisation_id: int) -> str: + if not settings.INGESTION_EVENTS_BUCKET_PREFIX: + raise ImproperlyConfigured("INGESTION_EVENTS_BUCKET_PREFIX is not set") + return f"{settings.INGESTION_EVENTS_BUCKET_PREFIX}-org-{organisation_id}" + + +def get_stream_name(organisation_id: int) -> str: + return f"{STREAM_NAME_PREFIX}{organisation_id}" + + +def _ensure_events_bucket(bucket_name: str, *, organisation_id: int) -> None: + s3 = _get_s3_client() + try: + s3.create_bucket( + Bucket=bucket_name, + CreateBucketConfiguration={"LocationConstraint": AWS_REGION}, + ) + except ClientError as exc: + if exc.response["Error"]["Code"] != "BucketAlreadyOwnedByYou": + raise + else: + logger.info( + "ingestion_infra.bucket_created", + organisation__id=organisation_id, + bucket__name=bucket_name, + ) + s3.put_public_access_block( + Bucket=bucket_name, + PublicAccessBlockConfiguration={ + "BlockPublicAcls": True, + "IgnorePublicAcls": True, + "BlockPublicPolicy": True, + "RestrictPublicBuckets": True, + }, + ) + s3.put_bucket_lifecycle_configuration( + Bucket=bucket_name, + LifecycleConfiguration={ + "Rules": [ + { + "ID": "expire-delivery-errors", + "Filter": {"Prefix": "errors/"}, + "Status": "Enabled", + "Expiration": {"Days": ERROR_OBJECT_EXPIRATION_DAYS}, + } + ] + }, + ) + + +def _expected_destination_configuration(bucket_name: str) -> dict[str, Any]: + return { + "RoleARN": settings.INGESTION_FIREHOSE_DELIVERY_ROLE_ARN, + "BucketARN": f"arn:aws:s3:::{bucket_name}", + "Prefix": EVENTS_PREFIX, + "ErrorOutputPrefix": ERRORS_PREFIX, + "BufferingHints": { + "SizeInMBs": BUFFERING_SIZE_MB, + "IntervalInSeconds": BUFFERING_INTERVAL_SECONDS, + }, + "CompressionFormat": "GZIP", + "DynamicPartitioningConfiguration": { + "Enabled": True, + "RetryOptions": {"DurationInSeconds": DYNAMIC_PARTITIONING_RETRY_SECONDS}, + }, + "ProcessingConfiguration": { + "Enabled": True, + "Processors": [ + { + "Type": "MetadataExtraction", + "Parameters": [ + { + "ParameterName": "MetadataExtractionQuery", + "ParameterValue": "{env_key:.environment_key}", + }, + { + "ParameterName": "JsonParsingEngine", + "ParameterValue": "JQ-1.6", + }, + ], + }, + { + "Type": "AppendDelimiterToRecord", + "Parameters": [ + {"ParameterName": "Delimiter", "ParameterValue": "\\n"}, + ], + }, + ], + }, + } + + +def _ensure_delivery_stream( + stream_name: str, + *, + bucket_name: str, + organisation_id: int, +) -> None: + try: + _get_firehose_client().create_delivery_stream( + DeliveryStreamName=stream_name, + DeliveryStreamType="DirectPut", + ExtendedS3DestinationConfiguration=_expected_destination_configuration( + bucket_name + ), + ) + except ClientError as exc: + if exc.response["Error"]["Code"] != "ResourceInUseException": + raise + else: + logger.info( + "ingestion_infra.stream_created", + organisation__id=organisation_id, + stream__name=stream_name, + bucket__name=bucket_name, + ) + + +def provision_ingestion_infrastructure( + organisation_id: int, +) -> IngestionInfrastructure: + """Idempotently create the organisation's events S3 bucket and Firehose + delivery stream; existing resources are left untouched.""" + if not settings.INGESTION_FIREHOSE_DELIVERY_ROLE_ARN: + raise ImproperlyConfigured("INGESTION_FIREHOSE_DELIVERY_ROLE_ARN is not set") + bucket_name = get_bucket_name(organisation_id) + stream_name = get_stream_name(organisation_id) + + _ensure_events_bucket(bucket_name, organisation_id=organisation_id) + _ensure_delivery_stream( + stream_name, + bucket_name=bucket_name, + organisation_id=organisation_id, + ) + return IngestionInfrastructure(bucket_name=bucket_name, stream_name=stream_name) diff --git a/api/tests/unit/experimentation/test_ingestion_infra_service.py b/api/tests/unit/experimentation/test_ingestion_infra_service.py new file mode 100644 index 000000000000..dee3686c338d --- /dev/null +++ b/api/tests/unit/experimentation/test_ingestion_infra_service.py @@ -0,0 +1,223 @@ +from collections.abc import Iterator +from typing import Any + +import boto3 +import pytest +from botocore.exceptions import ClientError +from django.core.exceptions import ImproperlyConfigured +from moto import mock_firehose, mock_s3 # type: ignore[import-untyped] +from pytest_django.fixtures import SettingsWrapper +from pytest_mock import MockerFixture +from pytest_structlog import StructuredLogCapture + +from experimentation import ingestion_infra_service +from experimentation.dataclasses import IngestionInfrastructure + +DELIVERY_ROLE_ARN = "arn:aws:iam::123456789012:role/firehose-events-delivery" + + +@pytest.fixture() +def ingestion_infra_settings(settings: SettingsWrapper) -> SettingsWrapper: + settings.INGESTION_EVENTS_BUCKET_PREFIX = "flagsmith-events-lake-test" + settings.INGESTION_FIREHOSE_DELIVERY_ROLE_ARN = DELIVERY_ROLE_ARN + return settings + + +@pytest.fixture() +def aws_backends(aws_credentials: None) -> Iterator[None]: + ingestion_infra_service._get_s3_client.cache_clear() + ingestion_infra_service._get_firehose_client.cache_clear() + with mock_s3(), mock_firehose(): + yield + ingestion_infra_service._get_s3_client.cache_clear() + ingestion_infra_service._get_firehose_client.cache_clear() + + +def test_provision_ingestion_infrastructure__no_bucket_prefix__raises_improperly_configured( + ingestion_infra_settings: SettingsWrapper, +) -> None: + # Given + ingestion_infra_settings.INGESTION_EVENTS_BUCKET_PREFIX = "" + + # When / Then + with pytest.raises(ImproperlyConfigured): + ingestion_infra_service.provision_ingestion_infrastructure(organisation_id=42) + + +def test_provision_ingestion_infrastructure__no_delivery_role_arn__raises_improperly_configured( + ingestion_infra_settings: SettingsWrapper, +) -> None: + # Given + ingestion_infra_settings.INGESTION_FIREHOSE_DELIVERY_ROLE_ARN = "" + + # When / Then + with pytest.raises(ImproperlyConfigured): + ingestion_infra_service.provision_ingestion_infrastructure(organisation_id=42) + + +def test_provision_ingestion_infrastructure__fresh_account__creates_bucket_and_stream( + ingestion_infra_settings: SettingsWrapper, + aws_backends: None, + log: StructuredLogCapture, +) -> None: + # When + result = ingestion_infra_service.provision_ingestion_infrastructure( + organisation_id=42, + ) + + # Then + assert result == IngestionInfrastructure( + bucket_name="flagsmith-events-lake-test-org-42", + stream_name="events-ingestion-org-42", + ) + + s3 = boto3.client("s3", region_name="eu-west-2") + public_access_block = s3.get_public_access_block(Bucket=result.bucket_name) + assert public_access_block["PublicAccessBlockConfiguration"] == { + "BlockPublicAcls": True, + "IgnorePublicAcls": True, + "BlockPublicPolicy": True, + "RestrictPublicBuckets": True, + } + lifecycle = s3.get_bucket_lifecycle_configuration(Bucket=result.bucket_name) + assert lifecycle["Rules"] == [ + { + "ID": "expire-delivery-errors", + "Filter": {"Prefix": "errors/"}, + "Status": "Enabled", + "Expiration": {"Days": 30}, + } + ] + + firehose = boto3.client("firehose", region_name="eu-west-2") + stream = firehose.describe_delivery_stream(DeliveryStreamName=result.stream_name)[ + "DeliveryStreamDescription" + ] + assert stream["DeliveryStreamType"] == "DirectPut" + destination = stream["Destinations"][0]["ExtendedS3DestinationDescription"] + assert destination["RoleARN"] == DELIVERY_ROLE_ARN + assert destination["BucketARN"] == "arn:aws:s3:::flagsmith-events-lake-test-org-42" + assert destination["Prefix"] == ( + "events/env_key=!{partitionKeyFromQuery:env_key}/" + "year=!{timestamp:yyyy}/month=!{timestamp:MM}/" + "day=!{timestamp:dd}/hour=!{timestamp:HH}/" + ) + assert destination["ErrorOutputPrefix"] == ( + "errors/!{firehose:error-output-type}/" + "year=!{timestamp:yyyy}/month=!{timestamp:MM}/" + "day=!{timestamp:dd}/hour=!{timestamp:HH}/" + ) + assert destination["CompressionFormat"] == "GZIP" + assert destination["BufferingHints"] == { + "SizeInMBs": 64, + "IntervalInSeconds": 300, + } + assert destination["DynamicPartitioningConfiguration"] == { + "Enabled": True, + "RetryOptions": {"DurationInSeconds": 300}, + } + assert destination["ProcessingConfiguration"] == { + "Enabled": True, + "Processors": [ + { + "Type": "MetadataExtraction", + "Parameters": [ + { + "ParameterName": "MetadataExtractionQuery", + "ParameterValue": "{env_key:.environment_key}", + }, + { + "ParameterName": "JsonParsingEngine", + "ParameterValue": "JQ-1.6", + }, + ], + }, + { + "Type": "AppendDelimiterToRecord", + "Parameters": [ + {"ParameterName": "Delimiter", "ParameterValue": "\\n"}, + ], + }, + ], + } + + assert log.events == [ + { + "level": "info", + "event": "ingestion_infra.bucket_created", + "organisation__id": 42, + "bucket__name": "flagsmith-events-lake-test-org-42", + }, + { + "level": "info", + "event": "ingestion_infra.stream_created", + "organisation__id": 42, + "stream__name": "events-ingestion-org-42", + "bucket__name": "flagsmith-events-lake-test-org-42", + }, + ] + + +def test_provision_ingestion_infrastructure__already_provisioned__is_idempotent( + ingestion_infra_settings: SettingsWrapper, + aws_backends: None, + log: StructuredLogCapture, +) -> None: + # Given + first = ingestion_infra_service.provision_ingestion_infrastructure( + organisation_id=42, + ) + events_after_first_run = list(log.events) + + # When + second = ingestion_infra_service.provision_ingestion_infrastructure( + organisation_id=42, + ) + + # Then + assert second == first + assert log.events == events_after_first_run + + +def test_provision_ingestion_infrastructure__bucket_creation_fails__propagates_client_error( + ingestion_infra_settings: SettingsWrapper, + mocker: MockerFixture, +) -> None: + # Given + mock_s3 = mocker.Mock() + mock_s3.create_bucket.side_effect = ClientError( + {"Error": {"Code": "AccessDenied", "Message": "nope"}}, + "CreateBucket", + ) + mocker.patch( + "experimentation.ingestion_infra_service._get_s3_client", + return_value=mock_s3, + ) + + # When / Then + with pytest.raises(ClientError, match="AccessDenied"): + ingestion_infra_service.provision_ingestion_infrastructure(organisation_id=42) + + +def test_provision_ingestion_infrastructure__stream_creation_fails__propagates_client_error( + ingestion_infra_settings: SettingsWrapper, + mocker: MockerFixture, +) -> None: + # Given + mocker.patch( + "experimentation.ingestion_infra_service._get_s3_client", + return_value=mocker.Mock(), + ) + mock_firehose_client: Any = mocker.Mock() + mock_firehose_client.create_delivery_stream.side_effect = ClientError( + {"Error": {"Code": "LimitExceededException", "Message": "too many"}}, + "CreateDeliveryStream", + ) + mocker.patch( + "experimentation.ingestion_infra_service._get_firehose_client", + return_value=mock_firehose_client, + ) + + # When / Then + with pytest.raises(ClientError, match="LimitExceededException"): + ingestion_infra_service.provision_ingestion_infrastructure(organisation_id=42) diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index 7458c075e032..2128946bca77 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md @@ -91,6 +91,25 @@ Attributes: - `experiment.id` - `organisation.id` +### `experimentation.ingestion_infra.bucket_created` + +Logged at `info` from: + - `api/experimentation/ingestion_infra_service.py:74` + +Attributes: + - `bucket.name` + - `organisation.id` + +### `experimentation.ingestion_infra.stream_created` + +Logged at `info` from: + - `api/experimentation/ingestion_infra_service.py:163` + +Attributes: + - `bucket.name` + - `organisation.id` + - `stream.name` + ### `experimentation.results.compute_failed` Logged at `error` from: