source: main/waeup.sirp/branches/ulif-extimgstore/src/waeup/sirp/imagestorage.py @ 7004

##
## imagestorage.py
## Login : <uli@pu.smp.net>
## Started on  Mon Jul  4 16:02:14 2011 Uli Fouquet
## $Id$
## 
## Copyright (C) 2011 Uli Fouquet
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 2 of the License, or
## (at your option) any later version.
## 
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
## 
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##
"""A storage for image files.
"""
import grok
import hashlib
import os
import tempfile
import transaction
import warnings
from StringIO import StringIO
from ZODB.blob import Blob
from persistent import Persistent
from hurry.file import HurryFile
from hurry.file.interfaces import IFileRetrieval
from zope.component import queryUtility
from zope.interface import Interface
from waeup.sirp.image import WAeUPImageFile
from waeup.sirp.interfaces import (
    IFileStoreNameChooser, IExtFileStore, IFileStoreHandler,)
from waeup.sirp.utils.helpers import cmp_files

def md5digest(fd):
    """Get an MD5 hexdigest for the file stored in `fd`.

    `fd`
      a file object open for reading.

    """
    return hashlib.md5(fd.read()).hexdigest()

class FileStoreNameChooser(grok.Adapter):
    grok.context(Interface)
    grok.implements(IFileStoreNameChooser)

    def checkName(self, name):
        """Check whether an object name is valid.

        Raises a user error if the name is not valid.
        """
        pass

    def chooseName(self, name):
        """Choose a unique valid name for the object.

        The given name and object may be taken into account when
        choosing the name.

        chooseName is expected to always choose a valid name (that
        would pass the checkName test) and never raise an error.
        """
        return u'unknown_file'

class Basket(grok.Container):
    """A basket holds a set of image files with same hash.
    """

    def _del(self):
        """Remove temporary files associated with local blobs.

        A basket holds files as Blob objects. Unfortunately, if a
        basket was not committed (put into ZODB), those blobs linger
        around as real files in some temporary directory and won't be
        removed.

        This is a helper function to remove all those uncommitted
        blobs that has to be called explicitly, for instance in tests.
        """
        key_list = self.keys()
        for key in key_list:
            item = self[key]
            if getattr(item, '_p_oid', None):
                # Don't mess around with blobs in ZODB
                continue
            fd = item.open('r')
            name = getattr(fd, 'name', None)
            fd.close()
            if name is not None and os.path.exists(name):
                os.unlink(name)
            del self[key]
        return

    def getInternalId(self, fd):
        """Get the basket-internal id for the file stored in `fd`.

        `fd` must be a file open for reading. If an (byte-wise) equal
        file can be found in the basket, its internal id (basket id)
        is returned, ``None`` otherwise.
        """
        fd.seek(0)
        for key, val in self.items():
            fd_stored = val.open('r')
            file_len = os.stat(fd_stored.name)[6]
            if file_len == 0:
                # Nasty workaround. Blobs seem to suffer from being emptied
                # accidentally.
                site = grok.getSite()
                if site is not None:
                    site.logger.warn(
                        'Empty Blob detected: %s' % fd_stored.name)
                warnings.warn("EMPTY BLOB DETECTED: %s" % fd_stored.name)
                fd_stored.close()
                val.open('w').write(fd.read())
                return key
            fd_stored.seek(0)
            if cmp_files(fd, fd_stored):
                fd_stored.close()
                return key
            fd_stored.close()
        return None

    @property
    def curr_id(self):
        """The current basket id.

        An integer number which is not yet in use. If there are
        already `maxint` entries in the basket, a :exc:`ValueError` is
        raised. The latter is _highly_ unlikely. It would mean to have
        more than 2**32 hash collisions, i.e. so many files with the
        same MD5 sum.
        """
        num = 1
        while True:
            if str(num) not in self.keys():
                return str(num)
            num += 1
            if num <= 0:
                name = getattr(self, '__name__', None)
                raise ValueError('Basket full: %s' % name)

    def storeFile(self, fd, filename):
        """Store the file in `fd` into the basket.

        The file will be stored in a Blob.
        """
        fd.seek(0)
        internal_id = self.getInternalId(fd) # Moves file pointer!
        if internal_id is None:
            internal_id = self.curr_id
            fd.seek(0)
            self[internal_id] = Blob()
            transaction.commit() # Urgently needed to make the Blob
                                 # persistent. Took me ages to find
                                 # out that solution, which makes some
                                 # design flaw in ZODB Blobs likely.
            self[internal_id].open('w').write(fd.read())
            fd.seek(0)
            self._p_changed = True
        return internal_id

    def retrieveFile(self, basket_id):
        """Retrieve a file open for reading with basket id `basket_id`.

        If there is no such id, ``None`` is returned. It is the
        callers responsibility to close the open file.
        """
        if basket_id in self.keys():
            return self[basket_id].open('r')
        return None

class ImageStorage(grok.Container):
    """A container for image files.
    """
    def _del(self):
        for basket in self.values():
            try:
                basket._del()
            except:
                pass

    def storeFile(self, fd, filename):
        fd.seek(0)
        digest = md5digest(fd)
        fd.seek(0)
        if not digest in self.keys():
            self[digest] = Basket()
        basket_id = self[digest].storeFile(fd, filename)
        full_id = "%s-%s" % (digest, basket_id)
        return full_id

    def retrieveFile(self, file_id):
        if not '-' in file_id:
            return None
        full_id, basket_id = file_id.split('-', 1)
        if not full_id in self.keys():
            return None
        return self[full_id].retrieveFile(basket_id)

class ImageStorageFileRetrieval(Persistent):
    grok.implements(IFileRetrieval)

    def getImageStorage(self):
        site = grok.getSite()
        if site is None:
            return None
        return site.get('images', None)

    def isImageStorageEnabled(self):
        site = grok.getSite()
        if site is None:
            return False
        if site.get('images', None) is None:
            return False
        return True

    def getFile(self, data):
        # ImageStorage is disabled, so give fall-back behaviour for
        # testing without ImageStorage
        if not self.isImageStorageEnabled():
            return StringIO(data)
        storage = self.getImageStorage()
        if storage is None:
            raise ValueError('Cannot find an image storage')
        result = storage.retrieveFile(data)
        if result is None:
            return StringIO(data)
        return storage.retrieveFile(data)

    def createFile(self, filename, f):
        if not self.isImageStorageEnabled():
            return WAeUPImageFile(filename, f.read())
        storage = self.getImageStorage()
        if storage is None:
            raise ValueError('Cannot find an image storage')
        file_id = storage.storeFile(f, filename)
        return WAeUPImageFile(filename, file_id)


class ExtFileStore(object): #grok.GlobalUtility):
    """External file store.

    External file stores are meant to store files 'externally' of the
    ZODB, i.e. in filesystem.

    Most important attribute of the external file store is the `root`
    path which gives the path to the location where files will be
    stored within.

    By default `root` is a ``'media/'`` directory in the root of the
    datacenter root of a site.

    The `root` attribute is 'read-only' because you normally don't
    want to change this path -- it is dynamic. That means, if you call
    the file store from 'within' a site, the root path will be located
    inside this site (a :class:`waeup.sirp.University` instance). If
    you call it from 'outside' a site some temporary dir (always the
    same during lifetime of the file store instance) will be used. The
    term 'temporary' tells what you can expect from this path
    persistence-wise.

    If you insist, you can pass a root path on initialization to the
    constructor but when calling from within a site afterwards, the
    site will override your setting for security measures. This way
    you can safely use one file store for different sites in a Zope
    instance simultanously and files from one site won't show up in
    another.

    An ExtFileStore instance is available as a global utility
    implementing :iface:`waeup.sirp.interfaces.IExtFileStore`.

    To add and retrieve files from the storage, use the appropriate
    methods below.
    """

    grok.implements(IExtFileStore)

    _root = None

    @property
    def root(self):
        """Root dir of this storage.

        The root dir is a readonly value determined dynamically. It
        holds media files for sites or other components.

        If a site is available we return a ``media/`` dir in the
        datacenter storage dir.

        Otherwise we create a temporary dir which will be remembered
        on next call.

        If a site exists and has a datacenter, it has always
        precedence over temporary dirs, also after a temporary
        directory was created.

        Please note that retrieving `root` is expensive. You might
        want to store a copy once retrieved in order to minimize the
        number of calls to `root`.

        """
        site = grok.getSite()
        if site is not None:
            root = os.path.join(site['datacenter'].storage, 'media')
            return root
        if self._root is None:
            self._root = tempfile.mkdtemp()
        return self._root

    def __init__(self, root=None):
        self._root = root
        return

    def getFile(self, file_id):
        """Get a file stored under file ID `file_id`.

        If the file cannot be found ``None`` is returned.
        """
        marker, filename, base, ext = self.extractMarker(file_id)
        handler = queryUtility(IFileStoreHandler, name=marker,
                               default=DefaultFileStoreHandler())
        path = handler.pathFromFileID(self, self.root, file_id)
        if not os.path.exists(path):
            return None
        fd = open(path, 'rb')
        return fd

    def createFile(self, filename, f):
        """Store a file.
        """
        file_id = filename
        root = self.root # Calls to self.root are expensive
        marker, filename, base, ext = self.extractMarker(file_id)
        handler = queryUtility(IFileStoreHandler, name=marker,
                               default=DefaultFileStoreHandler())
        f, path, file_obj = handler.createFile(
            self, root, file_id, filename, f)
        dirname = os.path.dirname(path)
        if not os.path.exists(dirname):
            os.makedirs(dirname, 0755)
        open(path, 'wb').write(f.read())
        return file_obj

    def extractMarker(self, file_id):
        """split filename into marker, filename, basename, and extension.

        A marker is a leading part of a string of form
        ``__MARKERNAME__`` followed by the real filename. This way we
        can put markers into a filename to request special processing.

        Returns a quadruple

          ``(marker, filename, basename, extension)``

        where ``marker`` is the marker in lowercase, filename is the
        complete trailing real filename, ``basename`` is the basename
        of the filename and ``extension`` the filename extension of
        the trailing filename. See examples below.

        Example:

           >>> extractMarker('__MaRkEr__sample.jpg')
           ('marker', 'sample.jpg', 'sample', '.jpg')

        If no marker is contained, we assume the whole string to be a
        real filename:

           >>> extractMarker('no-marker.txt')
           ('', 'no-marker.txt', 'no-marker', '.txt')

        Filenames without extension give an empty extension string:

           >>> extractMarker('no-marker')
           ('', 'no-marker', 'no-marker', '')

        """
        if not isinstance(file_id, basestring) or not file_id:
            return ('', '', '', '')
        parts = file_id.split('__', 2)
        marker = ''
        if len(parts) == 3 and parts[0] == '':
            marker = parts[1].lower()
            file_id = parts[2]
        basename, ext = os.path.splitext(file_id)
        return (marker, file_id, basename, ext)

grok.global_utility(ExtFileStore, provides=IExtFileStore)

class DefaultStorage(ExtFileStore):
    grok.provides(IFileRetrieval)

grok.global_utility(DefaultStorage, provides=IFileRetrieval)

class DefaultFileStoreHandler(grok.GlobalUtility):
    grok.implements(IFileStoreHandler)

    def pathFromFileID(self, store, root, file_id):
        return os.path.join(root, file_id)

    def createFile(self, store, root, filename, file_id, f):
        path = self.pathFromFileID(store, root, file_id)
        return f, path, WAeUPImageFile(filename, file_id)
        return path, HurryFile(filename, file_id)