feat(CockroachDB support): Switch to JSONField for SDK metrics labels

api/app_analytics/migrations/0006_add_labels.py

@@ -1,9 +1,10 @@
 # Generated by Django 4.2.21 on 2025-06-16 16:55
 
-import django.contrib.postgres.fields.hstore
-from django.contrib.postgres.operations import HStoreExtension
+# import django.contrib.postgres.fields.hstore
+# from django.contrib.postgres.operations import HStoreExtension
 
 from django.db import migrations
+from django.db import models
 
 
 class Migration(migrations.Migration):
@@ -13,25 +14,37 @@ class Migration(migrations.Migration):
     ]
 
     operations = [
-        HStoreExtension(),
+        # # The extension usage is annulated by the subsequent migration.
+        # # To avoid the operational overhead of enabling it in some environments,
+        # # we do not create it here anymore.
+        # # However, leave it commented out for reference/history.
+        # HStoreExtension(),
         migrations.AddField(
             model_name="apiusagebucket",
             name="labels",
-            field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            # This field is now a JSONField instead of HStoreField.
+            # field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            field=models.JSONField(default=dict),
         ),
         migrations.AddField(
             model_name="apiusageraw",
             name="labels",
-            field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            # This field is now a JSONField instead of HStoreField.
+            # field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            field=models.JSONField(default=dict),
         ),
         migrations.AddField(
             model_name="featureevaluationbucket",
             name="labels",
-            field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            # This field is now a JSONField instead of HStoreField.
+            # field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            field=models.JSONField(default=dict),
         ),
         migrations.AddField(
             model_name="featureevaluationraw",
             name="labels",
-            field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            # This field is now a JSONField instead of HStoreField.
+            # field=django.contrib.postgres.fields.hstore.HStoreField(default=dict),
+            field=models.JSONField(default=dict),
         ),
     ]

api/app_analytics/migrations/0008_labels_jsonb.py

@@ -0,0 +1,75 @@
+# Generated by Django 4.2.22 on 2025-07-25 14:09
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("app_analytics", "0007_rename_environment_id_created_at_index"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            state_operations=[
+                migrations.AlterField(
+                    model_name="apiusagebucket",
+                    name="labels",
+                    field=models.JSONField(default=dict),
+                ),
+                migrations.AlterField(
+                    model_name="apiusageraw",
+                    name="labels",
+                    field=models.JSONField(default=dict),
+                ),
+                migrations.AlterField(
+                    model_name="featureevaluationbucket",
+                    name="labels",
+                    field=models.JSONField(default=dict),
+                ),
+                migrations.AlterField(
+                    model_name="featureevaluationraw",
+                    name="labels",
+                    field=models.JSONField(default=dict),
+                ),
+            ],
+            database_operations=[
+                migrations.RunSQL(
+                    # Only alter the columns that are currently hstore.
+                    # See 0006_add_labels to understand why we are doing this.
+                    sql="""
+                    DO $$
+                    DECLARE relname text;
+                    BEGIN
+                        FOR relname IN
+                            SELECT c.relname
+                            FROM pg_class c
+                            JOIN pg_attribute a ON a.attrelid = c.oid
+                            JOIN pg_type t ON a.atttypid = t.oid
+                            WHERE c.relname IN (
+                                'app_analytics_apiusagebucket',
+                                'app_analytics_apiusageraw',
+                                'app_analytics_featureevaluationbucket',
+                                'app_analytics_featureevaluationraw'
+                            )
+                            AND a.attname = 'labels'
+                            AND t.typname = 'hstore'
+                        LOOP
+                            EXECUTE format(
+                                'ALTER TABLE %I
+                                 ALTER COLUMN labels DROP DEFAULT,
+                                 ALTER COLUMN labels TYPE jsonb USING hstore_to_json(labels),
+                                 ALTER COLUMN labels SET DEFAULT ''{}''::jsonb',
+                                relname
+                            );
+                        END LOOP;
+                    END
+                    $$;
+                    """,
+                    # We don't want hstore in the database at all, 
+                    # so don't do anything for reverse SQL.
+                    reverse_sql="",
+                ),
+            ],
+        ),
+    ]

api/app_analytics/models.py

@@ -1,6 +1,5 @@
 from datetime import timedelta
 
-from django.contrib.postgres.fields import HStoreField
 from django.core.exceptions import ValidationError
 from django.db import models
 from django_lifecycle import (  # type: ignore[import-untyped]
@@ -59,7 +58,7 @@ class APIUsageRaw(models.Model):
     host = models.CharField(max_length=255)
     resource = models.IntegerField(choices=Resource.choices)
     count = models.PositiveIntegerField(default=1)
-    labels = HStoreField(default=dict)
+    labels = models.JSONField(default=dict)
 
     class Meta:
         indexes = [models.Index(fields=["environment_id", "created_at"])]
@@ -70,7 +69,7 @@ class AbstractBucket(LifecycleModelMixin, models.Model):  # type: ignore[misc]
     created_at = models.DateTimeField()
     total_count = models.PositiveIntegerField()
     environment_id = models.PositiveIntegerField()
-    labels = HStoreField(default=dict)
+    labels = models.JSONField(default=dict)
 
     class Meta:
         abstract = True
@@ -107,7 +106,7 @@ class FeatureEvaluationRaw(models.Model):
     environment_id = models.PositiveIntegerField()
     evaluation_count = models.IntegerField(default=0)
     created_at = models.DateTimeField(auto_now_add=True)
-    labels = HStoreField(default=dict)
+    labels = models.JSONField(default=dict)
 
     # Both stored for tracking multivariate split testing.
     identity_identifier = models.CharField(max_length=2000, null=True, default=None)

api/tests/conftest.py

@@ -160,3 +160,11 @@ def migrator(migrator_factory: MigratorFactory) -> Migrator:
         pytest.skip("Skip migration tests to speed up tests where necessary")
     migrator: Migrator = migrator_factory()
     return migrator
+
+
+@pytest.fixture()
+def analytics_migrator(migrator_factory: MigratorFactory) -> Migrator:
+    if settings.SKIP_MIGRATION_TESTS:  # pragma: no cover
+        pytest.skip("Skip migration tests to speed up tests where necessary")
+    migrator: Migrator = migrator_factory("analytics")
+    return migrator

api/tests/unit/app_analytics/test_migrations.py

@@ -0,0 +1,204 @@
+import pytest
+from django.db import connections
+from django_test_migrations.migrator import Migrator
+
+pytestmark = pytest.mark.use_analytics_db
+
+
+def test_0008_labels_jsonb__fresh_install__preserves_data(
+    analytics_migrator: Migrator,
+) -> None:
+    """Test migration on fresh install where labels are already JSONField.
+
+    0006 creates labels as JSONField, so 0008's SQL is a no-op
+    (the conditional loop finds no hstore columns). This test verifies
+    the PL/pgSQL block is syntactically valid and data is preserved.
+    """
+    # Given - state at 0007 (labels columns exist as JSONField)
+    old_state = analytics_migrator.apply_initial_migration(
+        ("app_analytics", "0007_rename_environment_id_created_at_index"),
+    )
+
+    APIUsageRaw = old_state.apps.get_model("app_analytics", "APIUsageRaw")
+    APIUsageBucket = old_state.apps.get_model("app_analytics", "APIUsageBucket")
+    FeatureEvaluationRaw = old_state.apps.get_model(
+        "app_analytics", "FeatureEvaluationRaw"
+    )
+    FeatureEvaluationBucket = old_state.apps.get_model(
+        "app_analytics", "FeatureEvaluationBucket"
+    )
+
+    # Create records with labels
+    labels = {"sdk_type": "python", "sdk_version": "3.0.0"}
+    api_raw = APIUsageRaw.objects.using("analytics").create(
+        environment_id=1, host="test", resource=1, labels=labels
+    )
+    api_bucket = APIUsageBucket.objects.using("analytics").create(
+        environment_id=1,
+        bucket_size=15,
+        created_at="2025-01-01T00:00:00Z",
+        total_count=10,
+        resource=1,
+        labels=labels,
+    )
+    fe_raw = FeatureEvaluationRaw.objects.using("analytics").create(
+        feature_name="test_feature",
+        environment_id=1,
+        evaluation_count=5,
+        labels=labels,
+    )
+    fe_bucket = FeatureEvaluationBucket.objects.using("analytics").create(
+        environment_id=1,
+        bucket_size=15,
+        created_at="2025-01-01T00:00:00Z",
+        total_count=10,
+        feature_name="test_feature",
+        labels=labels,
+    )
+
+    # When - apply the jsonb migration
+    new_state = analytics_migrator.apply_tested_migration(
+        ("app_analytics", "0008_labels_jsonb"),
+    )
+
+    # Then - all records and their labels are preserved
+    NewAPIUsageRaw = new_state.apps.get_model("app_analytics", "APIUsageRaw")
+    NewAPIUsageBucket = new_state.apps.get_model("app_analytics", "APIUsageBucket")
+    NewFeatureEvaluationRaw = new_state.apps.get_model(
+        "app_analytics", "FeatureEvaluationRaw"
+    )
+    NewFeatureEvaluationBucket = new_state.apps.get_model(
+        "app_analytics", "FeatureEvaluationBucket"
+    )
+
+    assert NewAPIUsageRaw.objects.using("analytics").get(id=api_raw.id).labels == labels
+    assert (
+        NewAPIUsageBucket.objects.using("analytics").get(id=api_bucket.id).labels
+        == labels
+    )
+    assert (
+        NewFeatureEvaluationRaw.objects.using("analytics").get(id=fe_raw.id).labels
+        == labels
+    )
+    assert (
+        NewFeatureEvaluationBucket.objects.using("analytics")
+        .get(id=fe_bucket.id)
+        .labels
+        == labels
+    )
+
+
+def test_0008_labels_jsonb__hstore_columns__converts_to_jsonb(
+    analytics_migrator: Migrator,
+) -> None:
+    """Test migration converts existing hstore columns to jsonb.
+
+    Simulates the upgrade path for installations that ran the original
+    0006 migration which created labels as HStoreField.
+    """
+    # Given - state at 0007 (labels columns exist as JSONField in Django state)
+    expected_tables = [
+        "app_analytics_apiusagebucket",
+        "app_analytics_apiusageraw",
+        "app_analytics_featureevaluationbucket",
+        "app_analytics_featureevaluationraw",
+    ]
+
+    analytics_migrator.apply_initial_migration(
+        ("app_analytics", "0007_rename_environment_id_created_at_index"),
+    )
+
+    # Simulate the original 0006 migration having created hstore columns
+    # by converting the jsonb columns back to hstore at the database level.
+    connection = connections["analytics"]
+    with connection.cursor() as cursor:
+        cursor.execute("CREATE EXTENSION IF NOT EXISTS hstore")
+        for table in expected_tables:
+            cursor.execute(
+                f"ALTER TABLE {table} "
+                f"ALTER COLUMN labels TYPE hstore USING labels::text::hstore, "
+                f"ALTER COLUMN labels SET DEFAULT ''::hstore"
+            )
+
+        # Insert data as hstore values
+        cursor.execute(
+            "INSERT INTO app_analytics_apiusageraw "
+            "(environment_id, host, resource, count, labels, created_at) "
+            'VALUES (1, \'test\', 1, 1, \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore, NOW()) '
+            "RETURNING id"
+        )
+        api_raw_id = cursor.fetchone()[0]
+
+        cursor.execute(
+            "INSERT INTO app_analytics_apiusagebucket "
+            "(environment_id, bucket_size, created_at, total_count, resource, labels) "
+            'VALUES (1, 15, \'2025-01-01T00:00:00Z\', 10, 1, \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore) '
+            "RETURNING id"
+        )
+        api_bucket_id = cursor.fetchone()[0]
+
+        cursor.execute(
+            "INSERT INTO app_analytics_featureevaluationraw "
+            "(feature_name, environment_id, evaluation_count, labels, created_at) "
+            'VALUES (\'test_feature\', 1, 5, \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore, NOW()) '
+            "RETURNING id"
+        )
+        fe_raw_id = cursor.fetchone()[0]
+
+        cursor.execute(
+            "INSERT INTO app_analytics_featureevaluationbucket "
+            "(environment_id, bucket_size, created_at, total_count, feature_name, labels) "
+            'VALUES (1, 15, \'2025-01-01T00:00:00Z\', 10, \'test_feature\', \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore) '
+            "RETURNING id"
+        )
+        fe_bucket_id = cursor.fetchone()[0]
+
+    # When - apply the jsonb migration
+    new_state = analytics_migrator.apply_tested_migration(
+        ("app_analytics", "0008_labels_jsonb"),
+    )
+
+    # Then - columns are now jsonb and data is preserved
+    expected_labels = {"sdk_type": "python", "sdk_version": "3.0.0"}
+
+    with connection.cursor() as cursor:
+        for table in expected_tables:
+            cursor.execute(
+                """
+                SELECT t.typname
+                FROM pg_class c
+                JOIN pg_attribute a ON a.attrelid = c.oid
+                JOIN pg_type t ON a.atttypid = t.oid
+                WHERE c.relname = %s AND a.attname = %s
+                """,
+                [table, "labels"],
+            )
+            assert cursor.fetchone()[0] == "jsonb"
+
+    NewAPIUsageRaw = new_state.apps.get_model("app_analytics", "APIUsageRaw")
+    NewAPIUsageBucket = new_state.apps.get_model("app_analytics", "APIUsageBucket")
+    NewFeatureEvaluationRaw = new_state.apps.get_model(
+        "app_analytics", "FeatureEvaluationRaw"
+    )
+    NewFeatureEvaluationBucket = new_state.apps.get_model(
+        "app_analytics", "FeatureEvaluationBucket"
+    )
+
+    assert (
+        NewAPIUsageRaw.objects.using("analytics").get(id=api_raw_id).labels
+        == expected_labels
+    )
+    assert (
+        NewAPIUsageBucket.objects.using("analytics").get(id=api_bucket_id).labels
+        == expected_labels
+    )
+    assert (
+        NewFeatureEvaluationRaw.objects.using("analytics").get(id=fe_raw_id).labels
+        == expected_labels
+    )
+    assert (
+        NewFeatureEvaluationBucket.objects.using("analytics")
+        .get(id=fe_bucket_id)
+        .labels
+        == expected_labels
+    )