diff --git a/alembic/versions/dae93f1d9110_populate_tag_and_study_tag_tables_with_.py b/alembic/versions/dae93f1d9110_populate_tag_and_study_tag_tables_with_.py index 3818e3eb52..c5a255a223 100644 --- a/alembic/versions/dae93f1d9110_populate_tag_and_study_tag_tables_with_.py +++ b/alembic/versions/dae93f1d9110_populate_tag_and_study_tag_tables_with_.py @@ -6,6 +6,7 @@ Create Date: 2024-02-08 10:30:20.590919 """ +import itertools import json import secrets @@ -30,18 +31,21 @@ def upgrade(): # retrieve the tags and the study-tag pairs from the db study_tags = connexion.execute("SELECT study_id,patch FROM study_additional_data") - study_tags = { - study_id: set(json.loads(patch).get("study", dict()).get("tags", [])) for study_id, patch in study_tags - } - labels = set(label for s_tags in study_tags.values() for label in s_tags) + tags_by_ids = {} + for study_id, patch in study_tags: + obj = json.loads(patch or "{}") + tags = frozenset(obj.get("study", {}).get("tags", ())) + tags_by_ids[study_id] = tags + labels = set(itertools.chain.from_iterable(tags_by_ids.values())) tags = {label: secrets.choice(COLOR_NAMES) for label in labels} + for label, color in tags.items(): connexion.execute( sa.text("INSERT INTO tag (label, color) VALUES (:label, :color)"), {"label": label, "color": color} ) # Create relationships between studies and tags in the `study_tag` table - study_tag_data = {(study_id, label) for study_id, tags in study_tags.items() for label in tags} + study_tag_data = {(study_id, label) for study_id, tags in tags_by_ids.items() for label in tags} for study_id, label in study_tag_data: connexion.execute( sa.text("INSERT INTO study_tag (study_id, tag_label) VALUES (:study_id, :tag_label)"),