fix(tags-db): update following code review

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)"),