def base_directory_model(path):
    m = base_model(path)
    m.update(
        type='directory',
        last_modified=DUMMY_CREATED_DATE,
        created=DUMMY_CREATED_DATE,
    )
    return m def base_model(path):
    return {
        "name": path.rsplit('/', 1)[-1],
        "path": path,
        "writable": True,
        "last_modified": None,
        "created": None,
        "content": None,
        "format": None,
        "mimetype": None,
    } def from_b64(path, bcontent, format):
    """
    Decode base64 content for a file.

    format:
      If 'text', the contents will be decoded as UTF-8.
      If 'base64', do nothing.
      If not specified, try to decode as UTF-8, and fall back to base64

    Returns a triple of decoded_content, format, and mimetype.
    """
    decoders = {
        'base64': lambda path, bcontent: (bcontent.decode('ascii'), 'base64'),
        'text': _decode_text_from_base64,
        None: _decode_unknown_from_base64,
    }

    try:
        content, real_format = decoders[format](path, bcontent)
    except HTTPError:
        # Pass through HTTPErrors, since we intend for them to bubble all the
        # way back to the API layer.
        raise
    except Exception as e:
        # Anything else should be wrapped in a CorruptedFile, since it likely
        # indicates misconfiguration of encryption.
        raise CorruptedFile(e)

    default_mimes = {
        'text': 'text/plain',
        'base64': 'application/octet-stream',
    }
    mimetype = mimetypes.guess_type(path)[0] or default_mimes[real_format]

    return content, real_format, mimetype def outside_root_to_404(fn):
    """
    Decorator for converting PathOutsideRoot errors to 404s.
    """
    @wraps(fn)
    def wrapped(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except PathOutsideRoot as e:
            raise HTTPError(404, "Path outside root: [%s]" % e.args[0])
    return wrapped def reads_base64(nb, as_version=NBFORMAT_VERSION):
    """
    Read a notebook from base64.
    """
    try:
        return reads(b64decode(nb).decode('utf-8'), as_version=as_version)
    except Exception as e:
        raise CorruptedFile(e) def to_api_path(db_path):
    """
    Convert database path into API-style path.
    """
    return db_path.strip('/') def to_b64(content, fmt):
    allowed_formats = {'text', 'base64'}
    if fmt not in allowed_formats:
        raise ValueError(
            "Expected file contents in {allowed}, got {fmt}".format(
                allowed=allowed_formats,
                fmt=fmt,
            )
        )
    if fmt == 'text':
        # Unicode -> bytes -> base64-encoded bytes.
        return b64encode(content.encode('utf8'))
    else:
        return content.encode('ascii') def writes_base64(nb, version=NBFORMAT_VERSION):
    """
    Write a notebook as base64.
    """
    return b64encode(writes(nb, version=version).encode('utf-8')) class PostgresCheckpoints(PostgresManagerMixin,
                          GenericCheckpointsMixin,
                          Checkpoints):
    """
    A Checkpoints implementation that saves checkpoints to a remote database.
    """

    @outside_root_to_404
    def create_notebook_checkpoint(self, nb, path):
        """Create a checkpoint of the current state of a notebook

        Returns a checkpoint_id for the new checkpoint.
        """
        b64_content = writes_base64(nb)
        with self.engine.begin() as db:
            return save_remote_checkpoint(
                db,
                self.user_id,
                path,
                b64_content,
                self.crypto.encrypt,
                self.max_file_size_bytes,
            )

    @outside_root_to_404
    def create_file_checkpoint(self, content, format, path):
        """Create a checkpoint of the current state of a file

        Returns a checkpoint_id for the new checkpoint.
        """
        try:
            b64_content = to_b64(content, format)
        except ValueError as e:
            self.do_400(str(e))
        with self.engine.begin() as db:
            return save_remote_checkpoint(
                db,
                self.user_id,
                path,
                b64_content,
                self.crypto.encrypt,
                self.max_file_size_bytes,
            )

    @outside_root_to_404
    def delete_checkpoint(self, checkpoint_id, path):
        """delete a checkpoint for a file"""
        with self.engine.begin() as db:
            return delete_single_remote_checkpoint(
                db, self.user_id, path, checkpoint_id,
            )

    def get_checkpoint_content(self, checkpoint_id, path):
        """Get the content of a checkpoint."""
        with self.engine.begin() as db:
            return get_remote_checkpoint(
                db,
                self.user_id,
                path,
                checkpoint_id,
                self.crypto.decrypt,
            )['content']

    @outside_root_to_404
    def get_notebook_checkpoint(self, checkpoint_id, path):
        b64_content = self.get_checkpoint_content(checkpoint_id, path)
        return {
            'type': 'notebook',
            'content': reads_base64(b64_content),
        }

    @outside_root_to_404
    def get_file_checkpoint(self, checkpoint_id, path):
        b64_content = self.get_checkpoint_content(checkpoint_id, path)
        content, format = _decode_unknown_from_base64(path, b64_content)
        return {
            'type': 'file',
            'content': content,
            'format': format,
        }

    @outside_root_to_404
    def list_checkpoints(self, path):
        """Return a list of checkpoints for a given file"""
        with self.engine.begin() as db:
            return list_remote_checkpoints(db, self.user_id, path)

    @outside_root_to_404
    def rename_all_checkpoints(self, old_path, new_path):
        """Rename all checkpoints for old_path to new_path."""
        with self.engine.begin() as db:
            return move_remote_checkpoints(
                db,
                self.user_id,
                old_path,
                new_path,
            )

    @outside_root_to_404
    def delete_all_checkpoints(self, path):
        """Delete all checkpoints for the given path."""
        with self.engine.begin() as db:
            delete_remote_checkpoints(db, self.user_id, path)

    def purge_db(self):
        """
        Purge all database records for the current user.
        """
        with self.engine.begin() as db:
            purge_remote_checkpoints(db, self.user_id) class CorruptedFile(Exception):
    pass class DirectoryExists(Exception):
    pass class DirectoryNotEmpty(Exception):
    pass class FileExists(Exception):
    pass class FileTooLarge(Exception):
    pass class NoSuchDirectory(Exception):
    pass class NoSuchFile(Exception):
    pass class PathOutsideRoot(Exception):
    pass class RenameRoot(Exception):
    pass class PostgresManagerMixin(HasTraits):
    """
    Shared behavior for Postgres-backed ContentsManagers.
    """
    db_url = Unicode(
        default_value="postgresql://{user}@/pgcontents".format(
            user=getuser(),
        ),
        config=True,
        help="Connection string for the database.",
    )

    user_id = Unicode(
        default_value=getuser(),
        allow_none=True,
        config=True,
        help="Name for the user whose contents we're managing.",
    )

    create_user_on_startup = Bool(
        default_value=True,
        config=True,
        help="Create a user for user_id automatically?",
    )

    max_file_size_bytes = Integer(
        default_value=UNLIMITED,
        config=True,
        help="Maximum size in bytes of a file that will be saved.",
    )

    crypto = Any(
        default_value=NoEncryption(),
        allow_none=False,
        config=True,
        help=(
            "Object with encrypt() and decrypt() methods to "
            "call on data entering/exiting the database.",
        )
    )

    engine = Instance(Engine)

    def _engine_default(self):
        return create_engine(self.db_url, echo=False)

    def __init__(self, *args, **kwargs):
        super(PostgresManagerMixin, self).__init__(*args, **kwargs)
        if self.create_user_on_startup:
            self.ensure_user()

    def ensure_user(self):
        with self.engine.begin() as db:
            ensure_db_user(db, self.user_id)

    def no_such_entity(self, path):
        self.do_404(
            u"No such entity: [{path}]".format(path=path)
        )

    def not_empty(self, path):
        self.do_400(
            u"Directory not empty: [{path}]".format(path=path)
        )

    def file_too_large(self, path):
        self.do_413(u"File is too large to save: [{path}]".format(path=path))

    def already_exists(self, path):
        self.do_409(u"File already exists: [{path}]".format(path=path))

    def do_400(self, msg):
        raise HTTPError(400, msg)

    def do_404(self, msg):
        raise HTTPError(404, msg)

    def do_409(self, msg):
        raise HTTPError(409, msg)

    def do_413(self, msg):
        raise HTTPError(413, msg)

    def do_500(self, msg):
        raise HTTPError(500, msg) def delete_directory(db, user_id, api_path):
    """
    Delete a directory.
    """
    db_dirname = from_api_dirname(api_path)
    try:
        result = db.execute(
            directories.delete().where(
                and_(
                    directories.c.user_id == user_id,
                    directories.c.name == db_dirname,
                )
            )
        )
    except IntegrityError as error:
        if is_foreign_key_violation(error):
            raise DirectoryNotEmpty(api_path)
        else:
            raise

    rowcount = result.rowcount
    if not rowcount:
        raise NoSuchDirectory(api_path)

    return rowcount def delete_file(db, user_id, api_path):
    """
    Delete a file.

    TODO: Consider making this a soft delete.
    """
    result = db.execute(
        files.delete().where(
            _file_where(user_id, api_path)
        )
    )

    rowcount = result.rowcount
    if not rowcount:
        raise NoSuchFile(api_path)

    return rowcount def dir_exists(db, user_id, api_dirname):
    """
    Check if a directory exists.
    """
    return _dir_exists(db, user_id, from_api_dirname(api_dirname)) def ensure_directory(db, user_id, api_path):
    """
    Ensure that the given user has the given directory.
    """
    with ignore_unique_violation():
        create_directory(db, user_id, api_path) def file_exists(db, user_id, path):
    """
    Check if a file exists.
    """
    try:
        get_file(
            db,
            user_id,
            path,
            include_content=False,
            decrypt_func=unused_decrypt_func,
        )
        return True
    except NoSuchFile:
        return False def get_directory(db, user_id, api_dirname, content):
    """
    Return the names of all files/directories that are direct children of
    api_dirname.

    If content is False, return a bare model containing just a database-style
    name.
    """
    db_dirname = from_api_dirname(api_dirname)
    if not _dir_exists(db, user_id, db_dirname):
        raise NoSuchDirectory(api_dirname)
    if content:
        files = files_in_directory(
            db,
            user_id,
            db_dirname,
        )
        subdirectories = directories_in_directory(
            db,
            user_id,
            db_dirname,
        )
    else:
        files, subdirectories = None, None

    # TODO: Consider using namedtuples for these return values.
    return {
        'name': db_dirname,
        'files': files,
        'subdirs': subdirectories,
    } def get_file(db, user_id, api_path, include_content, decrypt_func):
    """
    Get file data for the given user_id and path.

    Include content only if include_content=True.
    """
    query_fields = _file_default_fields()
    if include_content:
        query_fields.append(files.c.content)

    return _get_file(db, user_id, api_path, query_fields, decrypt_func) def get_file_id(db, user_id, api_path):
    """
    Get the value in the 'id' column for the file with the given
    user_id and path.
    """
    return _get_file(
        db,
        user_id,
        api_path,
        [files.c.id],
        unused_decrypt_func,
    )['id'] def purge_user(db, user_id):
    """
    Delete a user and all of their resources.
    """
    db.execute(files.delete().where(
        files.c.user_id == user_id
    ))
    db.execute(directories.delete().where(
        directories.c.user_id == user_id
    ))
    db.execute(users.delete().where(
        users.c.id == user_id
    )) def rename_directory(db, user_id, old_api_path, new_api_path):
    """
    Rename a directory.
    """
    old_db_path = from_api_dirname(old_api_path)
    new_db_path = from_api_dirname(new_api_path)

    if old_db_path == '/':
        raise RenameRoot('Renaming the root directory is not permitted.')

    # Overwriting existing directories is disallowed.
    if _dir_exists(db, user_id, new_db_path):
        raise DirectoryExists(new_api_path)

    # Set this foreign key constraint to deferred so it's not violated
    # when we run the first statement to update the name of the directory.
    db.execute('SET CONSTRAINTS '
               'pgcontents.directories_parent_user_id_fkey DEFERRED')

    new_api_dir, new_name = split_api_filepath(new_api_path)
    new_db_dir = from_api_dirname(new_api_dir)

    # Update the name and parent_name columns for the directory that is being
    # renamed. The parent_name column will not change for a simple rename, but
    # will if the directory is moving.
    db.execute(
        directories.update().where(
            and_(
                directories.c.user_id == user_id,
                directories.c.name == old_db_path,
            )
        ).values(
            name=new_db_path,
            parent_name=new_db_dir,
        )
    )

    # Update the name and parent_name of any descendant directories. Do this in
    # a single statement so the non-deferrable check constraint is satisfied.
    db.execute(
        directories.update().where(
            and_(
                directories.c.user_id == user_id,
                directories.c.name.startswith(old_db_path),
                directories.c.parent_name.startswith(old_db_path),
            ),
        ).values(
            name=func.concat(
                new_db_path,
                func.right(directories.c.name, -func.length(old_db_path)),
            ),
            parent_name=func.concat(
                new_db_path,
                func.right(
                    directories.c.parent_name,
                    -func.length(old_db_path),
                ),
            ),
        )
    ) def rename_file(db, user_id, old_api_path, new_api_path):
    """
    Rename a file.
    """
    # Overwriting existing files is disallowed.
    if file_exists(db, user_id, new_api_path):
        raise FileExists(new_api_path)

    new_dir, new_name = split_api_filepath(new_api_path)

    db.execute(
        files.update().where(
            _file_where(user_id, old_api_path),
        ).values(
            name=new_name,
            parent_name=new_dir,
            created_at=func.now(),
        )
    ) def save_file(db, user_id, path, content, encrypt_func, max_size_bytes):
    """
    Save a file.

    TODO: Update-then-insert is probably cheaper than insert-then-update.
    """
    content = preprocess_incoming_content(
        content,
        encrypt_func,
        max_size_bytes,
    )
    directory, name = split_api_filepath(path)
    with db.begin_nested() as savepoint:
        try:
            res = db.execute(
                files.insert().values(
                    name=name,
                    user_id=user_id,
                    parent_name=directory,
                    content=content,
                )
            )
        except IntegrityError as error:
            # The file already exists, so overwrite its content with the newer
            # version.
            if is_unique_violation(error):
                savepoint.rollback()
                res = db.execute(
                    files.update().where(
                        _file_where(user_id, path),
                    ).values(
                        content=content,
                        created_at=func.now(),
                    )
                )
            else:
                # Unknown error.  Reraise
                raise

    return res
from itertools import chain
from tornado import web
from traitlets import default
from .api_utils import (
    base_directory_model,
    base_model,
    from_b64,
    outside_root_to_404,
    reads_base64,
    to_api_path,
    to_b64,
    writes_base64,
)
from .checkpoints import PostgresCheckpoints
from .error import (
    CorruptedFile,
    DirectoryExists,
    DirectoryNotEmpty,
    FileExists,
    FileTooLarge,
    NoSuchDirectory,
    NoSuchFile,
    PathOutsideRoot,
    RenameRoot,
)
from .managerbase import PostgresManagerMixin
from .query import (
    delete_directory,
    delete_file,
    dir_exists,
    ensure_directory,
    file_exists,
    get_directory,
    get_file,
    get_file_id,
    purge_user,
    rename_directory,
    rename_file,
    save_file,
)
from .utils.ipycompat import Bool, ContentsManager, from_dict
                )
            except NoSuchDirectory:
                if self.file_exists(path):
                    # TODO: It's awkward/expensive to have to check this to
                    # return a 400 instead of 404. Consider just 404ing.
                    self.do_400("Wrong type: %s" % path)
                else:
                    self.no_such_entity(path)

        return self._directory_model_from_db(record, content)

    def _convert_file_records(self, file_records):
        """
        Apply _notebook_model_from_db or _file_model_from_db to each entry
        in file_records, depending on the result of `guess_type`.
        """
        for record in file_records:
            type_ = self.guess_type(record['name'], allow_directory=False)
            if type_ == 'notebook':
                yield self._notebook_model_from_db(record, False)
            elif type_ == 'file':
                yield self._file_model_from_db(record, False, None)
            else:
                self.do_500("Unknown file type %s" % type_)

    def _directory_model_from_db(self, record, content):
        """
        Build a directory model from database directory record.
        """
        model = base_directory_model(to_api_path(record['name']))
        if content:
            model['format'] = 'json'
            model['content'] = list(
                chain(
                    self._convert_file_records(record['files']),
                    (
                        self._directory_model_from_db(subdir, False)
                        for subdir in record['subdirs']
                    ),
                )
            )
        return model

    def _file_model_from_db(self, record, content, format):
        """
        Build a file model from database record.
        """
        # TODO: Most of this is shared with _notebook_model_from_db.
        path = to_api_path(record['parent_name'] + record['name'])
        model = base_model(path)
        model['type'] = 'file'
        model['last_modified'] = model['created'] = record['created_at']
        if content:
            bcontent = record['content']
            model['content'], model['format'], model['mimetype'] = from_b64(
                path,
                bcontent,
                format,
            )
        return model

    def _get_file(self, path, content, format):
        with self.engine.begin() as db:
            try:
                record = get_file(
                    db,
                    self.user_id,
                    path,
                    content,
                    self.crypto.decrypt,
                )
            except NoSuchFile:
                if self.dir_exists(path):
                    # TODO: It's awkward/expensive to have to check this to
                    # return a 400 instead of 404. Consider just 404ing.
                    self.do_400(u"Wrong type: %s" % path)
                else:
                    self.no_such_entity(path)
        return self._file_model_from_db(record, content, format)

    def _save_notebook(self, db, model, path):
        """
        Save a notebook.

        Returns a validation message.
        """
        nb_contents = from_dict(model['content'])
        self.check_and_sign(nb_contents, path)
        save_file(
            db,
            self.user_id,
            path,
            writes_base64(nb_contents),
            self.crypto.encrypt,
            self.max_file_size_bytes,
        )
        # It's awkward that this writes to the model instead of returning.
        self.validate_notebook_model(model)
        return model.get('message')

    def _save_file(self, db, model, path):
        """
        Save a non-notebook file.
        """
        save_file(
            db,
            self.user_id,
            path,
            to_b64(model['content'], model.get('format', None)),
            self.crypto.encrypt,
            self.max_file_size_bytes,
        )
        return None

    def _save_directory(self, db, path):
        """
        'Save' a directory.
        """
        ensure_directory(db, self.user_id, path)


    @outside_root_to_404