quay/data/userfiles.py

import os
import logging
import magic

from uuid import uuid4
from flask import url_for, request, send_file, make_response, abort
from flask.views import View
from io import BufferedReader


logger = logging.getLogger(__name__)


class UserfilesHandlers(View):
  methods = ['GET', 'PUT']

  def __init__(self, distributed_storage, location, files):
    self._storage = distributed_storage
    self._files = files
    self._locations = {location}
    self._magic = magic.Magic(mime=True)

  def get(self, file_id):
    path = self._files.get_file_id_path(file_id)
    try:
      file_stream = self._storage.stream_read_file(self._locations, path)
      buffered = BufferedReader(file_stream)
      file_header_bytes = buffered.peek(1024)
      return send_file(buffered, mimetype=self._magic.from_buffer(file_header_bytes))
    except IOError:
      abort(404)

  def put(self, file_id):
    input_stream = request.stream
    if request.headers.get('transfer-encoding') == 'chunked':
      # Careful, might work only with WSGI servers supporting chunked
      # encoding (Gunicorn)
      input_stream = request.environ['wsgi.input']

    c_type = request.headers.get('Content-Type', None)

    path = self._files.get_file_id_path(file_id)
    self._storage.stream_write(self._locations, path, input_stream, c_type)

    return make_response('Okay')

  def dispatch_request(self, file_id):
    if request.method == 'GET':
      return self.get(file_id)
    elif request.method == 'PUT':
      return self.put(file_id)


class DelegateUserfiles(object):
  def __init__(self, app, distributed_storage, location, path, handler_name):
    self._app = app
    self._storage = distributed_storage
    self._locations = {location}
    self._prefix = path
    self._handler_name = handler_name

  def _build_url_adapter(self):
    return self._app.url_map.bind(self._app.config['SERVER_HOSTNAME'],
                                  script_name=self._app.config['APPLICATION_ROOT'] or '/',
                                  url_scheme=self._app.config['PREFERRED_URL_SCHEME'])

  def get_file_id_path(self, file_id):
    return os.path.join(self._prefix, file_id)

  def prepare_for_drop(self, mime_type, requires_cors=True):
    """ Returns a signed URL to upload a file to our bucket. """
    logger.debug('Requested upload url with content type: %s' % mime_type)
    file_id = str(uuid4())
    path = self.get_file_id_path(file_id)
    url = self._storage.get_direct_upload_url(self._locations, path, mime_type, requires_cors)

    if url is None:
      with self._app.app_context() as ctx:
        ctx.url_adapter = self._build_url_adapter()
        return (url_for(self._handler_name, file_id=file_id, _external=True), file_id)

    return (url, file_id)

  def store_file(self, file_like_obj, content_type):
    file_id = str(uuid4())
    path = self.get_file_id_path(file_id)
    self._storage.stream_write(self._locations, path, file_like_obj, content_type)
    return file_id

  def get_file_url(self, file_id, expires_in=300, requires_cors=False):
    path = self.get_file_id_path(file_id)
    url = self._storage.get_direct_download_url(self._locations, path, expires_in, requires_cors)

    if url is None:
      with self._app.app_context() as ctx:
        ctx.url_adapter = self._build_url_adapter()
        return url_for(self._handler_name, file_id=file_id, _external=True)

    return url

  def get_file_checksum(self, file_id):
    path = self.get_file_id_path(file_id)
    return self._storage.get_checksum(self._locations, path)


class Userfiles(object):
  def __init__(self, app=None, distributed_storage=None):
    self.app = app
    if app is not None:
      self.state = self.init_app(app, distributed_storage)
    else:
      self.state = None

  def init_app(self, app, distributed_storage):
    location = app.config.get('USERFILES_LOCATION')
    path = app.config.get('USERFILES_PATH', None)

    handler_name = 'userfiles_handlers'

    userfiles = DelegateUserfiles(app, distributed_storage, location, path, handler_name)

    app.add_url_rule('/userfiles/<file_id>',
                     view_func=UserfilesHandlers.as_view(handler_name,
                                                         distributed_storage=distributed_storage,
                                                         location=location,
                                                         files=userfiles))

    # register extension with app
    app.extensions = getattr(app, 'extensions', {})
    app.extensions['userfiles'] = userfiles
    return userfiles

  def __getattr__(self, name):
    return getattr(self.state, name, None)