quay/util/migrate/backfill_user_uuids.py

import logging
import uuid

from data.database import User, db
from app import app

LOGGER = logging.getLogger(__name__)

def backfill_user_uuids():
  """ Generates UUIDs for any Users without them. """
  LOGGER.setLevel(logging.DEBUG)
  LOGGER.debug('User UUID Backfill: Began execution')


  # Check to see if any users are missing uuids.
  has_missing_uuids = True
  try:
    User.select(User.id).where(User.uuid >> None).get()
  except User.DoesNotExist:
    has_missing_uuids = False

  if not has_missing_uuids:
    LOGGER.debug('User UUID Backfill: No migration needed')
    return

  LOGGER.debug('User UUID Backfill: Starting migration')
  while True:
    batch_user_ids = list(User
                          .select(User.id)
                          .where(User.uuid >> None)
                          .limit(100))

    if len(batch_user_ids) == 0:
      # There are no users left to backfill. We're done!
      LOGGER.debug('User UUID Backfill: Backfill completed')
      return

    LOGGER.debug('User UUID Backfill: Found %s records to update', len(batch_user_ids))
    for user_id in batch_user_ids:
      with app.config['DB_TRANSACTION_FACTORY'](db):
        try:
          user = User.select(User.id, User.uuid).where(User.id == user_id).get()
          user.uuid = str(uuid.uuid4())
          user.save(only=[User.uuid])
        except User.DoesNotExist:
          pass


if __name__ == "__main__":
  logging.basicConfig(level=logging.DEBUG)
  logging.getLogger('boto').setLevel(logging.CRITICAL)
  logging.getLogger('peewee').setLevel(logging.CRITICAL)

  backfill_user_uuids()