From eb514c0ce27bffddbc5418c5f707093707adbb8e Mon Sep 17 00:00:00 2001
From: littlewolf <nickptic@code-artel.com>
Date: Fri, 13 Dec 2019 12:16:02 +0300
Subject: [PATCH] Add clean db setup

---
 apps/location/transfer_data.py                | 59 ++++++++++++++-----
 apps/transfer/management/commands/transfer.py |  3 +-
 apps/transfer/serializers/location.py         | 21 +++++++
 3 files changed, 68 insertions(+), 15 deletions(-)

diff --git a/apps/location/transfer_data.py b/apps/location/transfer_data.py
index 0b510da1..68ccab59 100644
--- a/apps/location/transfer_data.py
+++ b/apps/location/transfer_data.py
@@ -193,18 +193,23 @@ def update_flags():
             query.save()
 
 
-def migrate_city_map_situation():
-    queryset = transfer_models.Cities.objects.raw("""SELECT cities.id, cities.map1, cities.map2, cities.map_ref, cities.situation
-                                        FROM cities WHERE
-                                        region_code IS NOT NULL AND 
-                                        region_code != "" AND
-                                        country_code_2 IS NOT NULL AND 
-                                        country_code_2 != ""
-                                            """)
+def migrate_city_map_situation(get_exists_cities=False):
+    if get_exists_cities:
+        ids = City.objects.values_list('mysql_id', flat=True)
+        queryset = transfer_models.Cities.objects.filter(id__in=list(ids))
+        queryset = list(queryset.values())
+    else:
+        queryset = transfer_models.Cities.objects.raw("""SELECT cities.id, cities.map1, cities.map2, cities.map_ref, cities.situation
+                                            FROM cities WHERE
+                                            region_code IS NOT NULL AND 
+                                            region_code != "" AND
+                                            country_code_2 IS NOT NULL AND 
+                                            country_code_2 != ""
+                                                """)
 
-    queryset = [vars(query) for query in queryset]
+        queryset = [vars(query) for query in queryset]
 
-    serialized_data = location_serializers.CityMapSerializer(data=queryset, many=True)
+    serialized_data = location_serializers.CityMapCorrectSerializer(data=queryset, many=True)
     if serialized_data.is_valid():
         serialized_data.save()
     else:
@@ -699,12 +704,36 @@ def add_fake_country():
             region.save()
 
 
+@atomic
+def setup_clean_db():
+    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()
+
+    print('add_correct_location_models')
+    add_correct_location_models(ruby_data)
+
+    print('add_fake_country')
+    add_fake_country()
+
+    print('migrate_city_map_situation')
+    migrate_city_map_situation(True)
+
+    print('update_flags')
+    update_flags()
+
+    print('transfer_city_gallery')
+    transfer_city_gallery()
+
+
 data_types = {
     "dictionaries": [
-        transfer_countries,
-        transfer_regions,
-        transfer_cities,
-        transfer_addresses,
+        # transfer_countries,
+        # transfer_regions,
+        # transfer_cities,
+        # transfer_addresses,
         transfer_wine_region,
         transfer_wine_sub_region,
         transfer_wine_village,
@@ -730,4 +759,6 @@ data_types = {
     "add_fake_country": [
         add_fake_country
     ],
+
+    "setup_clean_db": [setup_clean_db],
 }
diff --git a/apps/transfer/management/commands/transfer.py b/apps/transfer/management/commands/transfer.py
index d10807b1..1880917f 100644
--- a/apps/transfer/management/commands/transfer.py
+++ b/apps/transfer/management/commands/transfer.py
@@ -53,7 +53,8 @@ class Command(BaseCommand):
         'migrate_city_gallery',
         'fix_location',
         'remove_old_locations',
-        'add_fake_country'
+        'add_fake_country',
+        'setup_clean_db'
     ]
 
     def handle(self, *args, **options):
diff --git a/apps/transfer/serializers/location.py b/apps/transfer/serializers/location.py
index e7435054..5adaf9fc 100644
--- a/apps/transfer/serializers/location.py
+++ b/apps/transfer/serializers/location.py
@@ -405,6 +405,27 @@ class CityMapSerializer(serializers.ModelSerializer):
         return data
 
 
+class CityMapCorrectSerializer(CityMapSerializer):
+    def set_city(self, data):
+        try:
+            city = models.City.objects.get(mysql_id=data['old_id'])
+        except models.City.DoesNotExist as e:
+            raise ValueError(f"Cannot find city with id = {data['old_id']}: {e}")
+        except MultipleObjectsReturned as e:
+            raise ValueError(f"Find multiple cities with id = {data['old_id']}: {e}")
+
+        data['city'] = city
+        return data
+
+    def create(self, validated_data):
+        city = validated_data.pop('city')
+        city.map1=validated_data['map1'],
+        city.map2=validated_data['map2'],
+        city.map_ref=validated_data['map_ref'],
+        city.situation=validated_data['situation']
+        city.save()
+
+
 class CityGallerySerializer(serializers.ModelSerializer):
     id = serializers.IntegerField()
     city_id = serializers.IntegerField()