Skip to content

Commit 7f0f7a3

Browse files
khvn26khvn26
committed
fix: Fix hstore-to-jsonb migration and add tests
- Add missing DECLARE for PL/pgSQL loop variable - Drop column default before type change to avoid cast error - Renumber migration to 0008 after rebase (0007 conflict) - Add tests for both fresh install and hstore upgrade paths
1 parent acdd59c commit 7f0f7a3

File tree

3 files changed

+215
-1
lines changed

3 files changed

+215
-1
lines changed

api/app_analytics/migrations/0007_labels_jsonb.py renamed to api/app_analytics/migrations/0008_labels_jsonb.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
class Migration(migrations.Migration):
77

88
dependencies = [
9-
("app_analytics", "0006_add_labels"),
9+
("app_analytics", "0007_rename_environment_id_created_at_index"),
1010
]
1111

1212
operations = [
@@ -39,6 +39,7 @@ class Migration(migrations.Migration):
3939
# See 0006_add_labels to understand why we are doing this.
4040
sql="""
4141
DO $$
42+
DECLARE relname text;
4243
BEGIN
4344
FOR relname IN
4445
SELECT c.relname
@@ -56,6 +57,7 @@ class Migration(migrations.Migration):
5657
LOOP
5758
EXECUTE format(
5859
'ALTER TABLE %I
60+
ALTER COLUMN labels DROP DEFAULT,
5961
ALTER COLUMN labels TYPE jsonb USING hstore_to_json(labels),
6062
ALTER COLUMN labels SET DEFAULT ''{}''::jsonb',
6163
relname

‎api/tests/conftest.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,3 +160,11 @@ def migrator(migrator_factory: MigratorFactory) -> Migrator:
160160
pytest.skip("Skip migration tests to speed up tests where necessary")
161161
migrator: Migrator = migrator_factory()
162162
return migrator
163+
164+
165+
@pytest.fixture()
166+
def analytics_migrator(migrator_factory: MigratorFactory) -> Migrator:
167+
if settings.SKIP_MIGRATION_TESTS: # pragma: no cover
168+
pytest.skip("Skip migration tests to speed up tests where necessary")
169+
migrator: Migrator = migrator_factory("analytics")
170+
return migrator

‎api/tests/unit/app_analytics/test_migrations.py‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
import pytest
2+
from django.db import connections
3+
from django_test_migrations.migrator import Migrator
4+
5+
pytestmark = pytest.mark.use_analytics_db
6+
7+
8+
def test_0008_labels_jsonb__fresh_install__preserves_data(
9+
analytics_migrator: Migrator,
10+
) -> None:
11+
"""Test migration on fresh install where labels are already JSONField.
12+
13+
0006 creates labels as JSONField, so 0008's SQL is a no-op
14+
(the conditional loop finds no hstore columns). This test verifies
15+
the PL/pgSQL block is syntactically valid and data is preserved.
16+
"""
17+
# Given - state at 0007 (labels columns exist as JSONField)
18+
old_state = analytics_migrator.apply_initial_migration(
19+
("app_analytics", "0007_rename_environment_id_created_at_index"),
20+
)
21+
22+
APIUsageRaw = old_state.apps.get_model("app_analytics", "APIUsageRaw")
23+
APIUsageBucket = old_state.apps.get_model("app_analytics", "APIUsageBucket")
24+
FeatureEvaluationRaw = old_state.apps.get_model(
25+
"app_analytics", "FeatureEvaluationRaw"
26+
)
27+
FeatureEvaluationBucket = old_state.apps.get_model(
28+
"app_analytics", "FeatureEvaluationBucket"
29+
)
30+
31+
# Create records with labels
32+
labels = {"sdk_type": "python", "sdk_version": "3.0.0"}
33+
api_raw = APIUsageRaw.objects.using("analytics").create(
34+
environment_id=1, host="test", resource=1, labels=labels
35+
)
36+
api_bucket = APIUsageBucket.objects.using("analytics").create(
37+
environment_id=1,
38+
bucket_size=15,
39+
created_at="2025-01-01T00:00:00Z",
40+
total_count=10,
41+
resource=1,
42+
labels=labels,
43+
)
44+
fe_raw = FeatureEvaluationRaw.objects.using("analytics").create(
45+
feature_name="test_feature",
46+
environment_id=1,
47+
evaluation_count=5,
48+
labels=labels,
49+
)
50+
fe_bucket = FeatureEvaluationBucket.objects.using("analytics").create(
51+
environment_id=1,
52+
bucket_size=15,
53+
created_at="2025-01-01T00:00:00Z",
54+
total_count=10,
55+
feature_name="test_feature",
56+
labels=labels,
57+
)
58+
59+
# When - apply the jsonb migration
60+
new_state = analytics_migrator.apply_tested_migration(
61+
("app_analytics", "0008_labels_jsonb"),
62+
)
63+
64+
# Then - all records and their labels are preserved
65+
NewAPIUsageRaw = new_state.apps.get_model("app_analytics", "APIUsageRaw")
66+
NewAPIUsageBucket = new_state.apps.get_model("app_analytics", "APIUsageBucket")
67+
NewFeatureEvaluationRaw = new_state.apps.get_model(
68+
"app_analytics", "FeatureEvaluationRaw"
69+
)
70+
NewFeatureEvaluationBucket = new_state.apps.get_model(
71+
"app_analytics", "FeatureEvaluationBucket"
72+
)
73+
74+
assert NewAPIUsageRaw.objects.using("analytics").get(id=api_raw.id).labels == labels
75+
assert (
76+
NewAPIUsageBucket.objects.using("analytics").get(id=api_bucket.id).labels
77+
== labels
78+
)
79+
assert (
80+
NewFeatureEvaluationRaw.objects.using("analytics").get(id=fe_raw.id).labels
81+
== labels
82+
)
83+
assert (
84+
NewFeatureEvaluationBucket.objects.using("analytics")
85+
.get(id=fe_bucket.id)
86+
.labels
87+
== labels
88+
)
89+
90+
91+
def test_0008_labels_jsonb__hstore_columns__converts_to_jsonb(
92+
analytics_migrator: Migrator,
93+
) -> None:
94+
"""Test migration converts existing hstore columns to jsonb.
95+
96+
Simulates the upgrade path for installations that ran the original
97+
0006 migration which created labels as HStoreField.
98+
"""
99+
# Given - state at 0007 (labels columns exist as JSONField in Django state)
100+
expected_tables = [
101+
"app_analytics_apiusagebucket",
102+
"app_analytics_apiusageraw",
103+
"app_analytics_featureevaluationbucket",
104+
"app_analytics_featureevaluationraw",
105+
]
106+
107+
analytics_migrator.apply_initial_migration(
108+
("app_analytics", "0007_rename_environment_id_created_at_index"),
109+
)
110+
111+
# Simulate the original 0006 migration having created hstore columns
112+
# by converting the jsonb columns back to hstore at the database level.
113+
connection = connections["analytics"]
114+
with connection.cursor() as cursor:
115+
cursor.execute("CREATE EXTENSION IF NOT EXISTS hstore")
116+
for table in expected_tables:
117+
cursor.execute(
118+
f"ALTER TABLE {table} "
119+
f"ALTER COLUMN labels TYPE hstore USING labels::text::hstore, "
120+
f"ALTER COLUMN labels SET DEFAULT ''::hstore"
121+
)
122+
123+
# Insert data as hstore values
124+
cursor.execute(
125+
"INSERT INTO app_analytics_apiusageraw "
126+
"(environment_id, host, resource, count, labels, created_at) "
127+
'VALUES (1, \'test\', 1, 1, \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore, NOW()) '
128+
"RETURNING id"
129+
)
130+
api_raw_id = cursor.fetchone()[0]
131+
132+
cursor.execute(
133+
"INSERT INTO app_analytics_apiusagebucket "
134+
"(environment_id, bucket_size, created_at, total_count, resource, labels) "
135+
'VALUES (1, 15, \'2025-01-01T00:00:00Z\', 10, 1, \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore) '
136+
"RETURNING id"
137+
)
138+
api_bucket_id = cursor.fetchone()[0]
139+
140+
cursor.execute(
141+
"INSERT INTO app_analytics_featureevaluationraw "
142+
"(feature_name, environment_id, evaluation_count, labels, created_at) "
143+
'VALUES (\'test_feature\', 1, 5, \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore, NOW()) '
144+
"RETURNING id"
145+
)
146+
fe_raw_id = cursor.fetchone()[0]
147+
148+
cursor.execute(
149+
"INSERT INTO app_analytics_featureevaluationbucket "
150+
"(environment_id, bucket_size, created_at, total_count, feature_name, labels) "
151+
'VALUES (1, 15, \'2025-01-01T00:00:00Z\', 10, \'test_feature\', \'"sdk_type"=>"python", "sdk_version"=>"3.0.0"\'::hstore) '
152+
"RETURNING id"
153+
)
154+
fe_bucket_id = cursor.fetchone()[0]
155+
156+
# When - apply the jsonb migration
157+
new_state = analytics_migrator.apply_tested_migration(
158+
("app_analytics", "0008_labels_jsonb"),
159+
)
160+
161+
# Then - columns are now jsonb and data is preserved
162+
expected_labels = {"sdk_type": "python", "sdk_version": "3.0.0"}
163+
164+
with connection.cursor() as cursor:
165+
for table in expected_tables:
166+
cursor.execute(
167+
"""
168+
SELECT t.typname
169+
FROM pg_class c
170+
JOIN pg_attribute a ON a.attrelid = c.oid
171+
JOIN pg_type t ON a.atttypid = t.oid
172+
WHERE c.relname = %s AND a.attname = %s
173+
""",
174+
[table, "labels"],
175+
)
176+
assert cursor.fetchone()[0] == "jsonb"
177+
178+
NewAPIUsageRaw = new_state.apps.get_model("app_analytics", "APIUsageRaw")
179+
NewAPIUsageBucket = new_state.apps.get_model("app_analytics", "APIUsageBucket")
180+
NewFeatureEvaluationRaw = new_state.apps.get_model(
181+
"app_analytics", "FeatureEvaluationRaw"
182+
)
183+
NewFeatureEvaluationBucket = new_state.apps.get_model(
184+
"app_analytics", "FeatureEvaluationBucket"
185+
)
186+
187+
assert (
188+
NewAPIUsageRaw.objects.using("analytics").get(id=api_raw_id).labels
189+
== expected_labels
190+
)
191+
assert (
192+
NewAPIUsageBucket.objects.using("analytics").get(id=api_bucket_id).labels
193+
== expected_labels
194+
)
195+
assert (
196+
NewFeatureEvaluationRaw.objects.using("analytics").get(id=fe_raw_id).labels
197+
== expected_labels
198+
)
199+
assert (
200+
NewFeatureEvaluationBucket.objects.using("analytics")
201+
.get(id=fe_bucket_id)
202+
.labels
203+
== expected_labels
204+
)

0 commit comments

Comments
 (0)