from mongodb_migrations.base import BaseMigration


class Migration(BaseMigration):

    def upgrade(self):
        for doc in self.db['application_rules'].find({'redundancy_level': {'$exists': True}}):
            self.db['application_rules'].update(
                {'_id': doc['_id']},
                {'$unset': {'redundancy_level': 1},
                 '$set': {
                     "hosts_settings.redundancy_level": doc['redundancy_level'],
                     "hubs_settings.redundancy_level": doc['redundancy_level'],
                 }
                 }
            )

    def downgrade(self):
        for doc in self.db['application_rules'].find({'hubs_settings.redundancy_level': {'$exists': True}}):
            self.db['application_rules'].update(
                {'_id': doc['_id']},
                {'$unset': {'hosts_settings': 1, "hubs_settings": 1},
                 '$set': {
                     "redundancy_level": doc['hubs_settings']['redundancy_level'],
                 }
                 }
            )