Add location.address correction function

apps/location/transfer_data.py

@@ -1,6 +1,6 @@
 from transfer.serializers import location as location_serializers
 from transfer import models as transfer_models
-from location.models import Country, Region, City
+from location.models import Country, Region, City, Address
 from pprint import pprint
 import json
 
@@ -223,11 +223,8 @@ def migrate_city_photos():
         pprint(f"Address serializer errors: {serialized_data.errors}")
 
 
-"""
-Update location models with ruby library
-Utils functions defined before transfer functions
-"""
-
+# Update location models with ruby library
+# Utils functions defined before transfer functions
 def get_ruby_socket(params):
     url = 'http://172.21.0.1:5678'  # docker host
     response = get(url, params=params)
@@ -239,6 +236,8 @@ def get_ruby_socket(params):
     return data
 
 
+# Get data from ruby and save it to file
+# Save errors from ruby to another file
 def get_ruby_data():
     cities = City.objects.filter(old_id__isnull=False)
     ruby_data = {}
@@ -305,14 +304,8 @@ def get_ruby_data():
     return ruby_data
 
 
-def fix_location_models():
-
-    try:
-        ruby_data_file = open(f"{settings.PROJECT_ROOT}/apps/location/ruby_data.py", "r")
-        ruby_data = json.loads(ruby_data_file.read())
-    except FileNotFoundError:
-        ruby_data = get_ruby_data()
-
+# Add correct objects of Country, Region and City with mysql_ids array (Country, Region) and mysql_id (City)
+def add_correct_location_models(ruby_data):
     for mysql_id, city_object in ruby_data.items():
         country_data = city_object["country"]
         try:
@@ -363,6 +356,32 @@ def fix_location_models():
             break
 
 
+def fix_location_address():
+    addresses = Address.objects.filter(old_id__isnull=False)
+    for address in addresses:
+        mysql_location = transfer_models.Locations.objects.get(id=address.old_id)
+        try:
+            correct_city = City.objects.get(mysql_id=mysql_location.city_id)
+        except City.DoesNotExist:
+            continue
+
+        address.city = correct_city
+        address.save()
+
+
+def fix_location_models():
+
+    try:
+        ruby_data_file = open(f"{settings.PROJECT_ROOT}/apps/location/ruby_data.py", "r")
+        ruby_data = json.loads(ruby_data_file.read())
+    except FileNotFoundError:
+        ruby_data = get_ruby_data()
+
+    add_correct_location_models(ruby_data)
+
+    fix_location_address()
+
+
 data_types = {
     "dictionaries": [
         transfer_countries,