feat(experimentation): Snowflake warehouse connection support (#7565)

Author: Zaimwa9

api/experimentation/migrations/0002_add_config_and_created_status.py

@@ -0,0 +1,32 @@
+# Generated by Django 5.2.14 on 2026-05-21 14:02
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("experimentation", "0001_introduce_warehouse_connection"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="warehouseconnection",
+            name="config",
+            field=models.JSONField(blank=True, null=True),
+        ),
+        migrations.AlterField(
+            model_name="warehouseconnection",
+            name="status",
+            field=models.CharField(
+                choices=[
+                    ("created", "Created"),
+                    ("pending_connection", "Pending Connection"),
+                    ("connected", "Connected"),
+                    ("errored", "Errored"),
+                ],
+                default="created",
+                max_length=50,
+            ),
+        ),
+    ]

api/experimentation/models.py

@@ -12,6 +12,7 @@ class WarehouseType(models.TextChoices):
 
 
 class WarehouseConnectionStatus(models.TextChoices):
+    CREATED = "created", "Created"
     PENDING_CONNECTION = "pending_connection", "Pending Connection"
     CONNECTED = "connected", "Connected"
     ERRORED = "errored", "Errored"
@@ -30,9 +31,12 @@ class WarehouseConnection(LifecycleModelMixin, SoftDeleteExportableModel):  # ty
     status = models.CharField(
         max_length=50,
         choices=WarehouseConnectionStatus.choices,
-        default=WarehouseConnectionStatus.PENDING_CONNECTION,
+        default=WarehouseConnectionStatus.CREATED,
     )
     name = models.CharField(max_length=255)
+    config: models.JSONField[dict[str, object] | None, dict[str, object] | None] = (
+        models.JSONField(null=True, blank=True)
+    )
     created_at = models.DateTimeField(auto_now_add=True)
 
     class Meta:

api/experimentation/serializers.py

@@ -8,13 +8,38 @@
     WarehouseConnectionStatus,
     WarehouseType,
 )
+from experimentation.types import SNOWFLAKE_DEFAULTS, SnowflakeConfig
 
 
 class WarehouseConnectionSerializer(serializers.ModelSerializer):  # type: ignore[type-arg]
+    name = serializers.CharField(max_length=255, required=False)
+    config = serializers.JSONField(default=None, required=False, allow_null=True)
+
     class Meta:
         model = WarehouseConnection
-        fields = ("id", "warehouse_type", "status", "name", "created_at")
-        read_only_fields = ("id", "status", "name", "created_at")
+        fields = ("id", "warehouse_type", "status", "name", "config", "created_at")
+        read_only_fields = ("id", "status", "created_at")
+
+    def validate(self, attrs: dict[str, Any]) -> dict[str, Any]:
+        warehouse_type: str = attrs.get(
+            "warehouse_type",
+            getattr(self.instance, "warehouse_type", ""),
+        )
+
+        if "config" not in attrs and self.instance is not None:
+            return attrs
+
+        config: dict[str, Any] | None = attrs.get("config")
+
+        if warehouse_type == WarehouseType.SNOWFLAKE:
+            attrs["config"] = self._validate_snowflake_config(config or {})
+        elif warehouse_type == WarehouseType.FLAGSMITH:
+            if config:
+                raise serializers.ValidationError(
+                    {"config": "Flagsmith warehouse does not accept configuration."}
+                )
+            attrs["config"] = None
+        return attrs
 
     def create(
         self,
@@ -23,6 +48,9 @@ def create(
         environment: Environment = validated_data["environment"]
         warehouse_type: str = validated_data["warehouse_type"]
 
+        if not validated_data.get("name"):
+            validated_data["name"] = self._generate_name(warehouse_type, environment)
+
         existing: WarehouseConnection | None = (
             WarehouseConnection.objects.all_with_deleted()
             .filter(
@@ -34,16 +62,29 @@ def create(
         )
         if existing:
             existing.deleted_at = None
-            existing.status = WarehouseConnectionStatus.PENDING_CONNECTION
-            existing.name = self._generate_name(warehouse_type, environment)
+            existing.status = WarehouseConnectionStatus.CREATED
+            existing.name = validated_data["name"]
+            existing.config = validated_data.get("config")
             existing.save()
             return existing
 
-        validated_data["name"] = self._generate_name(warehouse_type, environment)
         result: WarehouseConnection = super().create(validated_data)
         return result
 
     @staticmethod
     def _generate_name(warehouse_type: str, environment: Environment) -> str:
         label = WarehouseType(warehouse_type).label
         return f"{label} Warehouse - {environment.name}"
+
+    @staticmethod
+    def _validate_snowflake_config(config: dict[str, Any]) -> SnowflakeConfig:
+        account_identifier = config.get("account_identifier", "")
+        if not account_identifier:
+            raise serializers.ValidationError(
+                {"config": {"account_identifier": "This field is required."}}
+            )
+        merged: SnowflakeConfig = {
+            **SNOWFLAKE_DEFAULTS,
+            **config,  # type: ignore[typeddict-item]
+        }
+        return merged

api/experimentation/types.py

@@ -0,0 +1,20 @@
+from typing import TypedDict
+
+
+class SnowflakeConfig(TypedDict):
+    account_identifier: str
+    warehouse: str
+    database: str
+    schema: str
+    role: str
+    user: str
+
+
+SNOWFLAKE_DEFAULTS: SnowflakeConfig = {
+    "account_identifier": "",
+    "warehouse": "COMPUTE_WH",
+    "database": "FLAGSMITH",
+    "schema": "ANALYTICS",
+    "role": "FLAGSMITH_LOADER",
+    "user": "FLAGSMITH_SERVICE",
+}

api/experimentation/views.py

@@ -17,6 +17,7 @@ class WarehouseConnectionViewSet(
     mixins.ListModelMixin,
     mixins.CreateModelMixin,
     mixins.RetrieveModelMixin,
+    mixins.UpdateModelMixin,
     mixins.DestroyModelMixin,
 ):
     serializer_class = WarehouseConnectionSerializer
@@ -34,6 +35,12 @@ def perform_create(self, serializer: BaseSerializer[WarehouseConnection]) -> Non
             connection, self._get_user(self.request), action="created"
         )
 
+    def perform_update(self, serializer: BaseSerializer[WarehouseConnection]) -> None:
+        connection: WarehouseConnection = serializer.save()
+        create_warehouse_audit_log(
+            connection, self._get_user(self.request), action="updated"
+        )
+
     def perform_destroy(self, instance: WarehouseConnection) -> None:
         create_warehouse_audit_log(
             instance, self._get_user(self.request), action="deleted"

api/tests/unit/experimentation/test_serializers.py

@@ -26,7 +26,7 @@ def test_create__no_existing__creates_new_connection(
     # Then
     assert connection.pk is not None
     assert connection.warehouse_type == WarehouseType.FLAGSMITH
-    assert connection.status == WarehouseConnectionStatus.PENDING_CONNECTION
+    assert connection.status == WarehouseConnectionStatus.CREATED
     assert connection.name == f"Flagsmith Warehouse - {environment.name}"
 
 
@@ -53,7 +53,7 @@ def test_create__soft_deleted_exists__resurrects_record(
     # Then
     assert connection.pk == original_pk
     assert connection.deleted_at is None
-    assert connection.status == WarehouseConnectionStatus.PENDING_CONNECTION
+    assert connection.status == WarehouseConnectionStatus.CREATED
     assert connection.name == f"Flagsmith Warehouse - {environment.name}"
 
 
@@ -72,3 +72,61 @@ def test_create__valid_data__name_uses_warehouse_type_label(
     # Then
     assert connection.name.startswith("Flagsmith Warehouse")
     assert environment.name in connection.name
+
+
+def test_create__snowflake_minimal_config__applies_defaults(
+    environment: Environment,
+) -> None:
+    # Given
+    serializer = WarehouseConnectionSerializer(
+        data={
+            "warehouse_type": "snowflake",
+            "name": "My Snowflake",
+            "config": {"account_identifier": "xy12345.us-east-1"},
+        },
+    )
+    serializer.is_valid(raise_exception=True)
+
+    # When
+    connection = serializer.save(environment=environment)
+
+    # Then
+    assert connection.config["account_identifier"] == "xy12345.us-east-1"
+    assert connection.config["warehouse"] == "COMPUTE_WH"
+    assert connection.config["database"] == "FLAGSMITH"
+    assert connection.config["schema"] == "ANALYTICS"
+    assert connection.config["role"] == "FLAGSMITH_LOADER"
+    assert connection.config["user"] == "FLAGSMITH_SERVICE"
+
+
+def test_create__snowflake_missing_account_identifier__raises_validation_error(
+    environment: Environment,
+) -> None:
+    # Given
+    serializer = WarehouseConnectionSerializer(
+        data={
+            "warehouse_type": "snowflake",
+            "name": "Bad Snowflake",
+            "config": {"warehouse": "MY_WH"},
+        },
+    )
+
+    # When / Then
+    assert not serializer.is_valid()
+    assert "config" in serializer.errors
+
+
+def test_create__flagsmith_with_config__raises_validation_error(
+    environment: Environment,
+) -> None:
+    # Given
+    serializer = WarehouseConnectionSerializer(
+        data={
+            "warehouse_type": "flagsmith",
+            "config": {"account_identifier": "nope"},
+        },
+    )
+
+    # When / Then
+    assert not serializer.is_valid()
+    assert "config" in serializer.errors