quay/storage/swift.py

""" Swift storage driver. Based on: github.com/bacongobbler/docker-registry-driver-swift/ """
from swiftclient.client import Connection, ClientException
from storage.basestorage import BaseStorage
from util.generatorfile import GeneratorFile

from random import SystemRandom
import string
import logging

logger = logging.getLogger(__name__)


class SwiftStorage(BaseStorage):
  def __init__(self, swift_container, storage_path, auth_url, swift_user,
               swift_password, auth_version=None, os_options=None, ca_cert_path=None):
    self._swift_container = swift_container
    self._storage_path = storage_path

    self._auth_url = auth_url
    self._ca_cert_path = ca_cert_path

    self._swift_user = swift_user
    self._swift_password = swift_password

    self._auth_version = auth_version or 2
    self._os_options = os_options or {}

    self._initialized = False
    self._swift_connection = None

  def _initialize(self):
    if self._initialized:
      return

    self._initialized = True
    self._swift_connection = self._get_connection()

  def _get_connection(self):
    return Connection(
        authurl=self._auth_url,
        cacert=self._ca_cert_path,

        user=self._swift_user,
        key=self._swift_password,

        auth_version=self._auth_version,
        os_options=self._os_options)

  def _get_relative_path(self, path):
    if path.startswith(self._storage_path):
      path = path[len(self._storage_path)]

    if path.endswith('/'):
      path = path[:-1]

    return path

  def _normalize_path(self, path=None):
    path = self._storage_path + (path or '')

    # Openstack does not like paths starting with '/' and we always normalize
    # to remove trailing '/'
    if path.startswith('/'):
      path = path[1:]

    if path.endswith('/'):
      path = path[:-1]

    return path

  def _get_container(self, path):
    self._initialize()
    path = self._normalize_path(path)

    if path and not path.endswith('/'):
        path += '/'

    try:
       _, container = self._swift_connection.get_container(
            container=self._swift_container,
            prefix=path, delimiter='/')
       return container
    except:
      logger.exception('Could not get container: %s', path)
      raise IOError('Unknown path: %s' % path)

  def _get_object(self, path, chunk_size=None):
    self._initialize()
    path = self._normalize_path(path)
    try:
      _, obj = self._swift_connection.get_object(self._swift_container, path,
          resp_chunk_size=chunk_size)
      return obj
    except Exception:
      logger.exception('Could not get object: %s', path)
      raise IOError('Path %s not found' % path)

  def _put_object(self, path, content, chunk=None, content_type=None, content_encoding=None):
    self._initialize()
    path = self._normalize_path(path)
    headers = {}

    if content_encoding is not None:
      headers['Content-Encoding'] = content_encoding

    try:
      self._swift_connection.put_object(self._swift_container, path, content,
                                        chunk_size=chunk, content_type=content_type,
                                        headers=headers)
    except ClientException:
      # We re-raise client exception here so that validation of config during setup can see
      # the client exception messages.
      raise
    except Exception:
      logger.exception('Could not put object: %s', path)
      raise IOError("Could not put content: %s" % path)

  def _head_object(self, path):
    self._initialize()
    path = self._normalize_path(path)
    try:
      return self._swift_connection.head_object(self._swift_container, path)
    except Exception:
      logger.exception('Could not head object: %s', path)
      return None

  def get_direct_download_url(self, path, expires_in=60, requires_cors=False):
    if requires_cors:
      return None

    # TODO(jschorr): This method is not strictly necessary but would result in faster operations
    # when using this storage engine. However, the implementation (as seen in the link below)
    # is not clean, so we punt on this for now.
    # http://docs.openstack.org/juno/config-reference/content/object-storage-tempurl.html
    return None

  def get_content(self, path):
    return self._get_object(path)

  def put_content(self, path, content):
    self._put_object(path, content)

  def stream_read(self, path):
    for data in self._get_object(path, self.buffer_size):
      yield data

  def stream_read_file(self, path):
    return GeneratorFile(self.stream_read(path))

  def stream_write(self, path, fp, content_type=None, content_encoding=None):
    self._put_object(path, fp, self.buffer_size, content_type=content_type,
                     content_encoding=content_encoding)

  def list_directory(self, path=None):
    container = self._get_container(path)
    if not container:
      raise OSError('Unknown path: %s' % path)

    for entry in container:
      param = None
      if 'name' in entry:
        param = 'name'
      elif 'subdir' in entry:
        param = 'subdir'
      else:
        continue

      yield self._get_relative_path(entry[param])

  def exists(self, path):
    return bool(self._head_object(path))

  def remove(self, path):
    self._initialize()
    path = self._normalize_path(path)
    try:
      self._swift_connection.delete_object(self._swift_container, path)
    except Exception:
      raise IOError('Cannot delete path: %s' % path)

  def _random_checksum(self, count):
    chars = string.ascii_uppercase + string.digits
    return ''.join(SystemRandom().choice(chars) for _ in range(count))

  def get_checksum(self, path):
    headers = self._head_object(path)
    if not headers:
      raise IOError('Cannot lookup path: %s' % path)

    return headers.get('etag', '')[1:-1][:7] or self._random_checksum(7)