Add correction location.country for award_types

apps/location/transfer_data.py

@@ -8,6 +8,7 @@ from django.conf import settings
 from django.core.exceptions import MultipleObjectsReturned
 from collection.models import Collection
 from requests import get
+from main.models import AwardType
 
 
 def transfer_countries():
@@ -391,6 +392,24 @@ def fix_location_collection():
         collection.save()
 
 
+def fix_award_type():
+    award_types = AwardType.objects.filter(old_id__isnull=False)
+    for award_type in award_types:
+        mysql_award_type = transfer_models.AwardTypes.objects.\
+            raw(f"""SELECT at.id, s.country_code_2 AS country_code
+                            FROM award_types as at
+                            JOIN sites s on s.id = at.site_id
+                            WHERE at.id={award_type.old_id}""")[0]
+
+        try:
+            correct_country = Country.objects.get(code=mysql_award_type.country_code, mysql_ids__isnull=False)
+        except Country.DoesNotExist:
+            continue
+
+        award_type.country = correct_country
+        award_type.save()
+
+
 def fix_location_models():
 
     # try:
@@ -402,7 +421,8 @@ def fix_location_models():
     # add_correct_location_models(ruby_data)
     #
     # fix_location_address()
-    fix_location_collection()
+    # fix_location_collection()
+    fix_award_type()
 
 
 data_types = {