quay/data/userfiles.py

import boto
import os
import logging

from boto.s3.key import Key
from uuid import uuid4


logger = logging.getLogger(__name__)


class S3FileWriteException(Exception):
  pass


class UserRequestFiles(object):
  def __init__(self, s3_access_key, s3_secret_key, bucket_name):
    self._initialized = False
    self._bucket_name = bucket_name
    self._access_key = s3_access_key
    self._secret_key = s3_secret_key
    self._prefix = 'userfiles'
    self._s3_conn = None
    self._bucket = None

  def _initialize_s3(self):
    if not self._initialized:
      self._s3_conn = boto.connect_s3(self._access_key, self._secret_key)
      self._bucket = self._s3_conn.get_bucket(self._bucket_name)
      self._initialized = True

  def prepare_for_drop(self, mime_type):
    """ Returns a signed URL to upload a file to our bucket. """
    self._initialize_s3()
    logger.debug('Requested upload url with content type: %s' % mime_type)
    file_id = str(uuid4())
    full_key = os.path.join(self._prefix, file_id)
    k = Key(self._bucket, full_key)
    url = k.generate_url(300, 'PUT', headers={'Content-Type': mime_type},
                         encrypt_key=True)
    return (url, file_id)

  def store_file(self, file_like_obj, content_type):
    self._initialize_s3()
    file_id = str(uuid4())
    full_key = os.path.join(self._prefix, file_id)
    k = Key(self._bucket, full_key)
    logger.debug('Setting s3 content type to: %s' % content_type)
    k.set_metadata('Content-Type', content_type)
    bytes_written = k.set_contents_from_file(file_like_obj, encrypt_key=True,
                                             rewind=True)

    if bytes_written == 0:
      raise S3FileWriteException('Unable to write file to S3')

    return file_id

  def get_file_url(self, file_id, expires_in=300, mime_type=None):
    self._initialize_s3()
    full_key = os.path.join(self._prefix, file_id)
    k = Key(self._bucket, full_key)
    headers = None
    if mime_type:
      headers={'Content-Type': mime_type}
    
    return k.generate_url(expires_in, headers=headers)

  def get_file_checksum(self, file_id):
    self._initialize_s3()
    full_key = os.path.join(self._prefix, file_id)
    k = self._bucket.lookup(full_key)
    return k.etag[1:-1][:7]