Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with
or
.
Download ZIP
Tree: 79fe213d7d
Fetching contributors…

Cannot retrieve contributors at this time

493 lines (421 sloc) 15.6 kB
"""
These classes implement ways to store files in the server.
"""
from util import log
class FileStore(object):
"""
An object that abstracts a filesystem. This is the base class for filestores.
"""
def __init__(self):
raise NotImplementedError
def new_file(self, **kwargs):
"""
Return a file handle for a new write-only file with the
given properties. If the file already exists, it will
overwritten.
:arg \*\*kwargs: the properties of the new file (one should be
``filename="[filename]"``)
:returns: an open file handle for the new file
:rtype: file handle
"""
raise NotImplementedError
def delete_files(self, **kwargs):
"""
Delete every file in the filestore whose properties match
the keyword arguments.
:arg \*\*kwargs: all files whose MongoDB properties match these
will be deleted
"""
raise NotImplementedError
def get_file(self, **kwargs):
"""
Return a read-only file handle for a given file
with the properties given by the keyword arguments.
If the file does not exist, return ``None``.
:arg \*\*kwargs: the properties of the desired file
:returns: the opened file, or ``None`` if no file exists
:rtype: file handle
"""
raise NotImplementedError
def create_file(self, file_handle, **kwargs):
"""
Copy an existing file into the filestore.
:arg file file_handle: a file handle open for reading
:arg \*\*kwargs: labels for the new file (one shoud be
``filename="[filename]"``)
"""
raise NotImplementedError
def copy_file(self, file_handle, **kwargs):
"""Copy a file from the filestore into another file.
:arg file file_handle: a file handle open for writing
:arg \*\*kwargs: labels to identify the file to copy
"""
raise NotImplementedError
def create_secret(self, session):
"""
Generate a new :mod:`hmac` object and associate it
with the session. Used only with "untrusted" database
adaptors. (See :ref:`trusted`.)
:arg str session: the ID of the new session
"""
raise NotImplementedError
def new_context(self):
"""
Reconnect to the filestore. This function should be
called before the first filestore access in each new process.
"""
def new_context_copy(self):
"""
Create a copy of this object for use in a single thread.
:returns: a new filestore object
:rtype: FileStore
"""
try:
from sagecell_config import mongo_config
except ImportError:
# we may not be able to import sagecell_config if we are untrusted
mongo_config=None
DEBUG = False
def Debugger(func):
if DEBUG:
def decorated(*args, **kwargs):
print "****************Entering ",func.func_name
print " args ",args, kwargs
#try:
# print "filename: %s"%(args[0]._filename(**kwargs),)
#except Exception as a:
# print "Couldn't get filename", a
ret = func(*args, **kwargs)
print ret
return ret
return decorated
else:
return func
from sqlalchemy import create_engine, Column, Integer, String
from sqlalchemy.types import Binary
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from StringIO import StringIO
class FileStoreSQLAlchemy(FileStore):
"""
A filestore in a SQLAlchemy database.
:arg str fs_file: the SQLAlchemy URI for a database file
"""
def __init__(self, fs_file=None):
if fs_file is not None:
engine = create_engine(fs_file)
self.SQLSession = sessionmaker(bind=engine)
FileStoreSQLAlchemy.Base.metadata.create_all(engine)
self.new_context()
@Debugger
def new_file(self, session, filename, **kwargs):
"""
See :meth:`FileStore.new_file`
"""
self.delete_files(session, filename)
log("FS Creating %s/%s"%(session, filename))
return FileStoreSQLAlchemy.DBFileWriter(self, session, filename)
@Debugger
def delete_files(self, session=None, filename=None, **kwargs):
"""
See :meth:`FileStore.new_file`
"""
q = self.dbsession.query(FileStoreSQLAlchemy.StoredFile)
if session is not None:
q = q.filter_by(session=session)
if filename is not None:
q = q.filter_by(filename=filename)
q.delete()
self.dbsession.commit()
@Debugger
def get_file(self, session, filename, **kwargs):
"""
See :meth:`FileStore.get_file`
"""
return StringIO(self.dbsession.query(FileStoreSQLAlchemy.StoredFile.contents) \
.filter_by(session=session, filename=filename).first().contents)
@Debugger
def create_file(self, file_handle, session, filename, **kwargs):
"""
See :meth:`FileStore.create_file`
"""
f = FileStoreSQLAlchemy.StoredFile(session=session, filename=filename)
if type(file_handle) is FileStoreSQLAlchemy.DBFileWriter:
contents = file_handle.getvalue()
else:
contents = file_handle.read()
f.contents = contents
self.dbsession.add(f)
self.dbsession.commit()
@Debugger
def copy_file(self, file_handle, session, filename, **kwargs):
"""
See :meth:`FileStore.copy_file`
"""
self.dbsession.add(FileStoreSQLAlchemy.StoredFile(session=session,
filename=filename, contents=file_handle.read()))
self.dbsession.commit()
@Debugger
def new_context(self):
"""
See :meth:`FileStore.new_context`
"""
self.dbsession = self.SQLSession()
@Debugger
def new_context_copy(self):
"""
See :meth:`FileStore.new_context_copy`
"""
new = type(self)()
new.SQLSession = self.SQLSession
new.new_context()
return new
Base = declarative_base()
class StoredFile(Base):
"""A file stored in the database"""
__tablename__ = 'filestore'
n = Column(Integer, primary_key=True)
session = Column(String)
filename = Column(String)
contents = Column(Binary)
class DBFileWriter(StringIO, object):
"""
A file-like object that writes its contents to the database when it is
closed.
:arg FileStoreSQLAlchemy filestore: the filestore object to write to
:arg str session: the ID of the session that is the source of this file
:arg str filename: the name of the file
"""
def __init__(self, filestore, session, filename):
self.filestore = filestore
self.session = session
self.filename = filename
super(type(self), self).__init__()
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
def close(self):
self.filestore.create_file(self, self.session, self.filename)
super(type(self), self).close()
try:
from gridfs import GridFS
import pymongo
except ImportError:
pass
class FileStoreMongo(FileStore):
"""
Filestore database using GridFS (see :mod:`gridfs`)
:arg pymongo.database.Database connection: MongoDB database object
"""
def __init__(self, connection):
self._conn=connection
self.new_context()
self._fs=GridFS(self.database)
def _filename(self, **kwargs):
return {'session': kwargs.get('session', kwargs.get('cell_id', 'SESSION NOT FOUND')), 'filename': kwargs['filename']}
@Debugger
def new_file(self, **kwargs):
"""
See :meth:`FileStore.new_file`
:rtype: :class:`gridfs.grid_file.GridIn`
"""
self.delete_files(**kwargs)
log("FS Creating %s"%self._filename(**kwargs))
return self._fs.new_file(**self._filename(**kwargs))
@Debugger
def delete_files(self, **kwargs):
"""
See :meth:`FileStore.delete_files`
"""
while self._fs.exists(self._filename(**kwargs)):
self._fs.delete(self._fs.get_last_version(**self._filename(**kwargs))._id)
@Debugger
def get_file(self, **kwargs):
"""
See :meth:`FileStore.get_file`
:rtype: :class:`gridfs.grid_file.GridOut`
"""
if self._fs.exists(self._filename(**kwargs)):
return self._fs.get(self._fs.get_last_version(**self._filename(**kwargs))._id)
else:
return None
@Debugger
def create_file(self, file_handle, **kwargs):
"""
See :meth:`FileStore.create_file`
"""
with self.new_file(**kwargs) as f:
f.write(file_handle.read())
@Debugger
def copy_file(self, file_handle, **kwargs):
"""
See :meth:`FileStore.copy_file`
"""
file_handle.write(self.get_file(**kwargs).read())
@Debugger
def new_context(self):
"""
See :meth:`FileStore.new_context`
"""
self.database=pymongo.database.Database(self._conn, mongo_config['mongo_db'])
uri=mongo_config['mongo_uri']
if '@' in uri:
# strip off optional mongodb:// part
if uri.startswith('mongodb://'):
uri=uri[len('mongodb://'):]
result=self.database.authenticate(uri[:uri.index(':')],uri[uri.index(':')+1:uri.index('@')])
if result==0:
raise Exception("MongoDB authentication problem")
@Debugger
def new_context_copy(self):
"""
See :meth:`FileStore.new_context_copy`
"""
return type(self)(self._conn)
valid_untrusted_methods=()
from flask import safe_join
import os
class FileStoreFilesystem(FileStore):
"""
Filestore using the file system
The levels parameter controls how the session is split up to give
subdirectories. For example, if ``levels=4``, then session
``0c490701-b1b0-40b8-88ea-70b61a580cf2`` files are stored in
subdirectory :file:`0/c/4/9/0c490701-b1b0-40b8-88ea-70b61a580cf2`.
This prevents having too many directories in the root directory.
:arg str dir: A directory in which to store files
:arg int levels: The number of levels to use for splitting up session directories
"""
def __init__(self, dir, levels=0):
self._dir = dir
self._levels=levels
def _filename(self, **kwargs):
if 'session' in kwargs:
session=kwargs['session']
elif 'cell_id' in kwargs:
session = kwargs['cell_id']
else:
session = "SESSION_NOT_FOUND"
session_subdir = list(str(session)[:self._levels])+[str(session)]
# Use Flask's safe_join to make sure we don't overwrite something crucial
session_dir = safe_join(self._dir, os.path.join(*session_subdir))
if not os.path.isdir(session_dir):
os.makedirs(session_dir)
return safe_join(session_dir, kwargs['filename'])
@Debugger
def new_file(self, **kwargs):
"""
See :meth:`FileStore.new_file`
"""
return open(self._filename(**kwargs), 'w')
@Debugger
def delete_files(self, session, filename):
"""
See :meth:`FileStore.delete_files`
"""
os.path.remove(self._filename(session=session, filename=filename))
@Debugger
def get_file(self, session, filename):
"""
See :meth:`FileStore.get_file`
"""
f=self._filename(session=session, filename=filename)
if os.path.exists(f):
return open(f, 'r')
else:
return None
@Debugger
def create_file(self, file_handle, **kwargs):
"""
See :meth:`FileStore.create_file`
"""
with self.new_file(**kwargs) as f:
f.write(file_handle.read())
@Debugger
def copy_file(self, file_handle, **kwargs):
"""
See :meth:`FileStore.copy_file`
"""
file_handle.write(self.get_file(**kwargs).read())
@Debugger
def new_context(self):
"""
Empty function
"""
pass
valid_untrusted_methods=()
import zmq
from db_zmq import db_method
from uuid import uuid4
from json import dumps
from os import fstat
import mmap
class FileStoreZMQ(FileStoreMongo):
u"""
A connection to a filestore database over \xd8MQ.
This can be used in the same way as a normal filestore,
but without risk of compromising the database.
:arg str address: the address the database should connect with
"""
def __init__(self, address):
self.address=address
self._xreq=None
@property
def socket(self):
"""
The ``socket`` property is automatically initialized the first
time it is called. We do this since we shouldn't create a
context in a parent process. Instead, we'll wait until we
actually start using the db api to create a context. If you
use the same class in a child process, you should first call
the :meth:`new_context` method.
"""
if self._xreq is None:
self.new_context()
return self._xreq
def new_context(self):
u"""
Reconnect to \xd8MQ. This function should be
called before the first database access in each new process.
"""
self._context=zmq.Context()
self._xreq=self._context.socket(zmq.XREQ)
self._xreq.connect(self.address)
log(u"ZMQ connecting to %s"%self.address)
def create_file(self, file_handle, hmac, **kwargs):
"""
See :meth:`FileStore.create_file`
:arg hmac: an object to be updated with the contents
of the message to be sent
:type hmac: :mod:`hmac` object
"""
# Use mmap if the filesize is larger than 1MiB;
# otherwise just copy the string to memory before sending it
if fstat(file_handle.fileno()).st_size>2**20:
f=mmap.mmap(file_handle.fileno(),0,access=mmap.ACCESS_READ)
else:
f=file_handle.read()
msg_str=dumps({'msg_type':'create_file',"header":str(uuid4()),
'content':kwargs})
log("Sending: msg_str: %r, old_digest: %r"%(msg_str, hmac.digest()))
hmac.update(msg_str)
log("New digest: %r"%hmac.digest())
message=[msg_str, hmac.digest(), f]
self.socket.send_multipart(message,copy=False,track=True).wait()
self.socket.recv()
def copy_file(self, file_handle, hmac, **kwargs):
"""
See :meth:`FileStore.copy_file`
:arg hmac: an object to be updated with the contents
of the message to be sent
:type hmac: :mod:`hmac` object
"""
msg_str=dumps({'msg_type':'copy_file','content':kwargs})
hmac.update(msg_str)
self.socket.send_multipart([msg_str, hmac.digest()])
file_handle.write(self.socket.recv())
create_secret=db_method('create_secret',['session'], True)
new_file=db_method('new_file',['session', 'filename'], True)
delete_files=db_method('delete_files',['session', 'filename'], True)
get_file=db_method('get_file',['session', 'filename'], True)
Jump to Line
Something went wrong with that request. Please try again.