path: root/alot/db/manager.py

# Copyright (C) 2011-2012  Patrick Totzke <patricktotzke@gmail.com>
# This file is released under the GNU GPL, version 3 or a later revision.
# For further details see the COPYING file

import abc
import asyncio
from   contextlib import closing
from   functools  import partialmethod
import logging
import os

from notmuch2 import Database, NotmuchError

from .errors import (DatabaseError, DatabaseROError, NonexistantObjectError,
                     QueryError)
from .sort   import ORDER
from .thread import Thread
from ..settings.const import settings

def _is_subdir_of(subpath, superpath):
    # make both absolute
    superpath = os.path.realpath(superpath)
    subpath = os.path.realpath(subpath)

    # return true, if the common prefix of both is equal to directory
    # e.g. /a/b/c/d.rst and directory is /a/b, the common prefix is /a/b
    return os.path.commonprefix([subpath, superpath]) == superpath

# DB write operations
class _DBOperation(abc.ABC):
    _dbman = None
    _tags  = None

    future = None

    def __init__(self, dbman, tags):
        self._dbman = dbman
        self._tags  = tags

        self.future = dbman._loop.create_future()

    def apply(self):
        logging.debug('Performing DB write: %s', self)

        try:
            with Database(self._dbman.path, Database.MODE.READ_WRITE) as db:
                logging.debug('got writeable DB')

                with db.atomic():
                    self._apply(db)
        except Exception as e:
            logging.exception(e)
            self.future.set_exception(e)
        else:
            logging.debug('DB write completed: %s', self)
            self.future.set_result(True)

    @abc.abstractmethod
    def _apply(self, db):
        pass

    def __str__(self):
        return '%s:%s' % (self.__class__.__name__, self._tags)

class _DBOperationTagAdd(_DBOperation):
    _query = None

    def __init__(self, dbman, tags, query):
        self._query = query

        super().__init__(dbman, tags)

    def _apply(self, db):
        for msg in db.messages(self._query):
            msg_tags  = msg.tags
            msg_tags |= self._tags

    def __str__(self):
        return '%s:%s' % (super().__str__(), self._query)

class _DBOperationTagRemove(_DBOperationTagAdd):
    def _apply(self, db):
        for msg in db.messages(self._query):
            msg_tags  = msg.tags
            msg_tags -= self._tags

class _DBOperationTagSet(_DBOperationTagAdd):
    def _apply(self, db):
        for msg in db.messages(self._query):
            msg_tags      = msg.tags
            property_tags = msg_tags & self._dbman._property_tags

            msg_tags.clear()
            msg_tags |= self._tags | property_tags

class _DBOperationMsgAdd(_DBOperation):
    _path = None

    def __init__(self, dbman, tags, path):
        self._path = path

        super().__init__(dbman, tags)

    def _apply(self, db):
        msg, _ = db.add(self._path, sync_flags = self._dbman._sync_flags)

        msg_tags  = msg.tags
        msg_tags |= self._tags

    def __str__(self):
        return '%s:%s' % (super().__str__(), self._path)

class DBManager:
    """
    Keeps track of your index parameters, maintains a write-queue and
    lets you look up threads and messages directly to the persistent wrapper
    classes.
    """
    """constants representing sort orders"""

    _loop          = None

    _sync_flags    = None

    _exclude_tags = None
    _property_tags = None

    _write_task    = None
    _write_queue   = None

    def __init__(self, loop, path=None, ro=False):
        """
        :param path: absolute path to the notmuch index
        :type path: str
        :param ro: open the index in read-only mode
        :type ro: bool
        """
        self.ro = ro
        self.path = path

        self._loop = loop

        # read notmuch's config regarding imap flag synchronization
        self._sync_flags    = settings.get_notmuch_setting('maildir', 'synchronize_flags')

        self._exclude_tags = frozenset(settings.get('exclude_tags'))
        self._property_tags = frozenset(settings.get('property_tags'))

    def _db_ro(self):
        return closing(Database(path = self.path, mode = Database.MODE.READ_ONLY))

    def _count(self, what, querystring):
        try:
            with self._db_ro() as db:
                func = getattr(db, 'count_' +what)
                return func(querystring, exclude_tags = self._exclude_tags)
        except NotmuchError:
            return -1

    count_messages = partialmethod(_count, 'messages')
    """returns number of messages that match `querystring`"""

    count_threads  = partialmethod(_count, 'threads')
    """returns number of threads that match `querystring`"""

    def _get_notmuch_thread(self, db, tid):
        """returns :class:`notmuch.database.Thread` with given id"""
        querystr = 'thread:' + tid
        try:
            return next(db.threads(querystr, exclude_tags = self._exclude_tags))
        except StopIteration:
            errmsg = 'no thread with id %s exists!' % tid
            raise NonexistantObjectError(errmsg)

    def get_thread(self, tid):
        """returns :class:`Thread` with given thread id (str)"""
        with self._db_ro() as db:
            return Thread(self, self._get_notmuch_thread(db, tid))

    def get_all_tags(self):
        """
        returns all tagsstrings used in the database
        :rtype: set of str
        """
        with self._db_ro() as db:
            return set(db.tags)

    def get_named_queries(self):
        """
        returns the named queries stored in the database.
        :rtype: dict (str -> str) mapping alias to full query string
        """
        q_prefix = 'query.'

        with self._db_ro() as db:
            queries = filter(lambda k: k.startswith(q_prefix), db.config)
            return { q[len(q_prefix):] : db.config[q] for q in queries }

    def get_threads(self, querystring, sort = ORDER.NEWEST_FIRST,
                    exclude_tags = frozenset()):
        """
        asynchronously look up thread ids matching `querystring`.

        :param querystring: The query string to use for the lookup
        :type querystring: str.
        :param sort: Sort order.
        :type query: alot.db.sort.ORDER
        :param exclude_tags: Tags to exclude by default unless included in the
                             search
        :type exclude_tags: set of str
        :returns: iterator over thread ids
        """
        with self._db_ro() as db:
            exclude_tags = self._exclude_tags | exclude_tags

            try:
                for t in db.threads(querystring, sort = sort,
                                    exclude_tags = exclude_tags):
                    yield t.threadid
            except NotmuchError as e:
                raise QueryError from e

    async def startup(self):
        self._write_queue = asyncio.Queue()
        self._write_task  = asyncio.create_task(self._db_write_task())

    async def shutdown(self):
        if self._write_task:
            await self._write_queue.put(None)
            await self._write_task

    async def _db_write_task(self):
        # this task serialises write operations on the database and
        # sends them off to a thread so they do not block the event loop
        while True:
            cur_item = await self._write_queue.get()
            if cur_item is None:
                self._write_queue.task_done()
                break

            logging.debug('submitting write task: %s', cur_item)

            await self._loop.run_in_executor(None, cur_item.apply)
            self._write_queue.task_done()

    def tags_add(self, query, tags):
        """
        Asynchronously add tags to messages matching `querystring`.

        :param querystring: notmuch search string
        :type querystring: str
        :param tags: a set of tags to be added
        :type tags: set of str
        """
        if self.ro:
            raise DatabaseROError()

        if not tags:
            ret = self._loop.create_future()
            ret.set_result(True)
            return ret

        op = _DBOperationTagAdd(self, tags, query)
        self._write_queue.put_nowait(op)
        return op.future

    def tags_remove(self, query, tags):
        """
        Asynchronously remove tags to messages matching `querystring`.

        :param querystring: notmuch search string
        :type querystring: str
        :param tags: a set of tags to be added
        :type tags: set of str
        """
        if self.ro:
            raise DatabaseROError()

        if not tags:
            ret = self._loop.create_future()
            ret.set_result(True)
            return ret

        op = _DBOperationTagRemove(self, tags, query)
        self._write_queue.put_nowait(op)
        return op.future

    def tags_set(self, query, tags):
        """
        Asynchronously set tags to messages matching `querystring`.

        :param querystring: notmuch search string
        :type querystring: str
        :param tags: a set of tags to be added
        :type tags: set of str
        """
        if self.ro:
            raise DatabaseROError()

        op = _DBOperationTagSet(self, tags, query)
        self._write_queue.put_nowait(op)
        return op.future

    def msg_add(self, path, tags):
        """
        Asynchronously add a file to the notmuch index.

        :param path: path to the file
        :type path: str
        :param tags: tagstrings to add
        :type tags: list of str
        """
        if self.ro:
            raise DatabaseROError()
        if not _is_subdir_of(path, self.path):
            msg = 'message path %s ' % path
            msg += ' is not below notmuchs '
            msg += 'root path (%s)' % self.path
            raise DatabaseError(msg)

        op = _DBOperationMsgAdd(self, tags, path)
        self._write_queue.put_nowait(op)
        return op.future