quay/buildserver/buildserver.py

import docker
import logging
import shutil
import os
import re
import requests
import json

from flask import Flask, jsonify, url_for, abort, make_response
from zipfile import ZipFile
from tempfile import TemporaryFile, mkdtemp
from uuid import uuid4
from multiprocessing.pool import ThreadPool
from base64 import b64encode


BUFFER_SIZE = 8 * 1024
LOG_FORMAT = '%(asctime)-15s - %(levelname)s - %(pathname)s - ' + \
             '%(funcName)s - %(message)s'

app = Flask(__name__)
logger = logging.getLogger(__name__)


def count_steps(dockerfile_path):
  with open(dockerfile_path, 'r') as dockerfileobj:
    steps = 0
    for line in dockerfileobj.readlines():
      stripped = line.strip()
      if stripped and stripped[0] is not '#':
        steps += 1
    return steps


def prepare_zip(request_file):
  build_dir = mkdtemp(prefix='docker-build-')

  # Save the zip file to temp somewhere
  with TemporaryFile() as zip_file:
    zip_file.write(request_file.content)
    to_extract = ZipFile(zip_file)
    to_extract.extractall(build_dir)

  return build_dir


def prepare_dockerfile(request_file):
  build_dir = mkdtemp(prefix='docker-build-')
  dockerfile_path = os.path.join(build_dir, "Dockerfile")
  with open(dockerfile_path, 'w') as dockerfile:
    dockerfile.write(request_file.content)

  return build_dir


def build_image(build_dir, tag_name, num_steps, result_object):
  try:
    logger.debug('Starting build.')
    docker_cl = docker.Client(version='1.5')
    result_object['status'] = 'building'
    build_status = docker_cl.build(path=build_dir, tag=tag_name)

    current_step = 0
    built_image = None
    for status in build_status:
      step_increment = re.search(r'Step ([0-9]+) :', status)
      if step_increment:
        current_step = int(step_increment.group(1))
        logger.debug('Step now: %s/%s' % (current_step, num_steps))
        result_object['current_command'] = current_step
        continue

      complete = re.match(r'Successfully built ([a-z0-9]+)$', status)
      if complete:
        built_image = complete.group(1)
        logger.debug('Final image ID is: %s' % built_image)
        continue

    shutil.rmtree(build_dir)

    # Get the image count
    if not built_image:
      result_object['status'] = 'error'
      result_object['message'] = 'Unable to build dockerfile.'
      return

    history = docker_cl.history(built_image)
    num_images = len(history)
    result_object['total_images'] = num_images

    result_object['status'] = 'pushing'
    logger.debug('Pushing to tag name: %s' % tag_name)
    resp = docker_cl.push(tag_name)

    current_image = 0
    image_progress = 0
    for status in resp:
      if u'status' in status:
        status_msg = status[u'status']
        next_image = r'(Pushing|Image) [a-z0-9]+( already pushed, skipping)?$'
        match = re.match(next_image, status_msg)
        if match:
          current_image += 1
          image_progress = 0
          logger.debug('Now pushing image %s/%s' %
                       (current_image, num_images))

        elif status_msg == u'Pushing' and u'progress' in status:
          percent = r'\(([0-9]+)%\)'
          match = re.search(percent, status[u'progress'])
          if match:
            image_progress = int(match.group(1))

        result_object['current_image'] = current_image
        result_object['image_completion_percent'] = image_progress

      elif u'errorDetail' in status:
        result_object['status'] = 'error'
        if u'message' in status[u'errorDetail']:
          result_object['message'] = status[u'errorDetail'][u'message']
        return

    result_object['status'] = 'complete'
  except Exception as e:
    logger.exception('Exception when processing request.')
    result_object['status'] = 'error'
    result_object['message'] = str(e.message)


MIME_PROCESSORS = {
  'application/zip': prepare_zip,
  'text/plain': prepare_dockerfile,
  'application/octet-stream': prepare_dockerfile,
}


build = {
  'total_commands': None,
  'total_images': None,
  'current_command': None,
  'current_image': None,
  'image_completion_percent': None,
  'status': 'waiting',
  'message': None,
}
pool = ThreadPool(1)


@app.before_first_request
def start_build():
  resource_url = os.environ['RESOURCE_URL']
  tag_name = os.environ['TAG']
  acccess_token = os.environ['TOKEN']

  logger.debug('Starting job with resource url: %s tag: %s and token: %s' %
               (resource_url, tag_name, acccess_token))

  # Save the token
  host = re.match(r'([a-z0-9.:]+)/.+/.+$', tag_name)
  if host:
    docker_endpoint = 'http://%s/v1/' % host.group(1)
    dockercfg_path = os.path.join(os.environ.get('HOME', '.'), '.dockercfg')
    token = b64encode('$token:%s' % acccess_token)
    with open(dockercfg_path, 'w') as dockercfg:
      payload = {
        docker_endpoint: {
          'auth': token,
          'email': '',
        }
      }
      dockercfg.write(json.dumps(payload))

  else:
    raise Exception('Invalid tag name: %s' % tag_name)

  docker_resource = requests.get(resource_url)
  c_type = docker_resource.headers['content-type']

  logger.info('Request to build file of type: %s with tag: %s' %
              (c_type, tag_name))

  if c_type not in MIME_PROCESSORS:
    raise Exception('Invalid dockerfile content type: %s' % c_type)

  build_dir = MIME_PROCESSORS[c_type](docker_resource)

  dockerfile_path = os.path.join(build_dir, "Dockerfile")
  num_steps = count_steps(dockerfile_path)
  logger.debug('Dockerfile had %s steps' % num_steps)

  logger.info('Sending job to builder pool.')
  build['total_commands'] = num_steps

  pool.apply_async(build_image, [build_dir, tag_name, num_steps,
                                 build])


@app.route('/build/', methods=['GET'])
def get_status():
  if build:
    return jsonify(build)
  abort(404)


@app.route('/status/', methods=['GET'])
def health_check():
  return make_response('Running')


if __name__ == '__main__':
  logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)
  app.run(host='0.0.0.0', port=5002, threaded=True)