REBASE_BEFORE_MERGE: Improve the db migration

plugins/pulp_deb/plugins/migrations/0002_add_multiple_hashes.py

@@ -1,11 +1,14 @@
 import logging
 import os
 
-from pulp.server import util
+from pulp.server import util as server_utils
 from pulp.server.db import connection
+from pulp.server.db.migrations.lib import utils as migration_utils
 
 
-CHECKSUM_TYPES = [util.TYPE_MD5, util.TYPE_SHA1, util.TYPE_SHA256]
+CHECKSUM_TYPES = [server_utils.TYPE_MD5,
+                  server_utils.TYPE_SHA1,
+                  server_utils.TYPE_SHA256,]
 
 _logger = logging.getLogger(__name__)
 
@@ -15,26 +18,29 @@ def migrate(*args, **kwargs):
     Add seperate fields for md5sum, sha1, sha256
     """
     deb_package_collection = connection.get_collection('units_deb')
-    for deb_package in deb_package_collection.find({}):
+    deb_package_count = deb_package_collection.count()
 
-        storage_path = deb_package['_storage_path']
-        if not os.path.exists(storage_path):
-            continue
+    with migration_utils.MigrationProgressLog('Deb Package', deb_package_count) as progress_log:
+        for deb_package in deb_package_collection.find({}).batch_size(100):
+            storage_path = deb_package['_storage_path']
+            if not os.path.exists(storage_path):
+                continue
 
-        package_id = deb_package['_id']
+            package_id = deb_package['_id']
 
-        with open(storage_path, 'r') as file_handle:
-            checksums = util.calculate_checksums(file_handle, CHECKSUM_TYPES)
+            with open(storage_path, 'r') as file_handle:
+                checksums = server_utils.calculate_checksums(file_handle, CHECKSUM_TYPES)
 
-        new_fields = {
-            'md5sum': checksums[util.TYPE_MD5],
-            'sha1': checksums[util.TYPE_SHA1],
-            'sha256': checksums[util.TYPE_SHA256],
-        }
+            new_fields = {
+                'md5sum': checksums[server_utils.TYPE_MD5],
+                'sha1': checksums[server_utils.TYPE_SHA1],
+                'sha256': checksums[server_utils.TYPE_SHA256],
+            }
 
-        if checksums[deb_package['checksumtype']] != deb_package['checksum']:
-            raise Exception('New checksum does not match existing checksum for\n'
-                            '_id = {}\nfile = {}'.format(package_id, storage_path))
+            if checksums[deb_package['checksumtype']] != deb_package['checksum']:
+                raise Exception('New checksum does not match existing checksum for\n'
+                                '_id = {}\nfile = {}'.format(package_id, storage_path))
 
-        deb_package_collection.update_one({'_id': package_id},
-                                          {'$set': new_fields},)
+            deb_package_collection.update_one({'_id': package_id},
+                                              {'$set': new_fields},)
+            progress_log.progress()