quay/util/migrate/migratebitbucketservices.py

import logging
import json

from app import app
from data.database import configure, BaseModel, uuid_generator
from peewee import *
from bitbucket import BitBucket
from buildtrigger.bitbuckethandler import BitbucketBuildTrigger

configure(app.config)

logger = logging.getLogger(__name__)

# Note: We vendor the RepositoryBuildTrigger and its dependencies here
class Repository(BaseModel):
  pass

class BuildTriggerService(BaseModel):
  name = CharField(index=True, unique=True)

class AccessToken(BaseModel):
  pass

class User(BaseModel):
  pass

class RepositoryBuildTrigger(BaseModel):
  uuid = CharField(default=uuid_generator)
  service = ForeignKeyField(BuildTriggerService, index=True)
  repository = ForeignKeyField(Repository, index=True)
  connected_user = ForeignKeyField(User)
  auth_token = CharField(null=True)
  private_key = TextField(null=True)
  config = TextField(default='{}')
  write_token = ForeignKeyField(AccessToken, null=True)
  pull_robot = ForeignKeyField(User, related_name='triggerpullrobot')


def run_bitbucket_migration():
  bitbucket_trigger = BuildTriggerService.get(BuildTriggerService.name == "bitbucket")

  encountered = set()
  while True:
    found = list(RepositoryBuildTrigger.select().where(
                      RepositoryBuildTrigger.service == bitbucket_trigger,
                      RepositoryBuildTrigger.config ** "%\"hook_id%"))

    found = [f for f in found if not f.uuid in encountered]

    if not found:
      logger.debug('No additional records found')
      return

    logger.debug('Found %s records to be changed', len(found))
    for trigger in found:
      encountered.add(trigger.uuid)

      try:
        config = json.loads(trigger.config)
      except:
        logging.error("Cannot parse config for trigger %s", trigger.uuid)
        continue

      logger.debug("Checking trigger %s", trigger.uuid)
      if 'hook_id' in config:
        logger.debug("Updating trigger %s to a webhook", trigger.uuid)

        trigger_handler = BitbucketBuildTrigger(trigger)
        client = trigger_handler._get_repository_client()

        hook_id = config['hook_id']

        # Lookup the old service hook.
        logger.debug("Looking up old service URL for trigger %s", trigger.uuid)
        (result, hook_data, err_msg) = client.services().get(hook_id)
        if not result or not hook_data:
          logger.error('Error when retrieving service hook for trigger %s: %s', trigger.uuid, err_msg)
          continue

        if not 'webhook_id' in config:
          hook_data = hook_data[0]['service']
          webhook_url = [f for f in hook_data['fields'] if f['name'] == 'URL'][0]['value']
          logger.debug("Adding webhook for trigger %s: %s", trigger.uuid, webhook_url)

          # Add the new web hook.
          description = 'Webhook for invoking builds on %s' % app.config['REGISTRY_TITLE_SHORT']
          webhook_events = ['repo:push']
          (result, data, err_msg) = client.webhooks().create(description, webhook_url, webhook_events)
          if not result:
            logger.error('Error when adding webhook for trigger %s: %s', trigger.uuid, err_msg)
            continue

          config['webhook_id'] = data['uuid']
          trigger.config = json.dumps(config)
          trigger.save()

        # Remove the old service hook.
        logger.debug("Deleting old service URL for trigger %s", trigger.uuid)
        (result, _, err_msg) = client.services().delete(hook_id)
        if not result:
          logger.error('Error when deleting service hook for trigger %s: %s', trigger.uuid, err_msg)
          continue

        del config['hook_id']

        # Update the config.
        trigger.config = json.dumps(config)
        trigger.save()
        logger.debug("Trigger %s updated to a webhook", trigger.uuid)


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

  run_bitbucket_migration()