path: root/src/carquinyol/indexstore.py

# Copyright (C) 2008, One Laptop Per Child
#
# 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., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

import logging
import os

import gobject
import xapian
from xapian import WritableDatabase, Document, Enquire, Query, QueryParser

import sugar.datastore

from carquinyol import layoutmanager
from carquinyol.layoutmanager import MAX_QUERY_LIMIT

_VALUE_TID = 0
_VALUE_CTIME = 1
_VALUE_VID = 2

_PREFIX_TID = 'Q'
_PREFIX_VID = 'V'
_PREFIX_BUNDLE_ID = 'A'
_PREFIX_ACTIVITY_ID = 'I'
_PREFIX_MIME_TYPE = 'M'
_PREFIX_KEEP = 'K'

# Force a flush every _n_ changes to the db
_FLUSH_THRESHOLD = 20

# Force a flush after _n_ seconds since the last change to the db
_FLUSH_TIMEOUT = 60

_PROPERTIES_NOT_TO_INDEX = ['ctime', 'activity_id', 'keep', 'preview']

_MAX_RESULTS = int(2 ** 31 - 1)

class IndexStore(object):
    """Index metadata and provide rich query facilities on it.
    """ 
    def __init__(self):
        self._database = None
        self._flush_timeout = None
        self._pending_writes = 0

    def open_index(self):
        index_path = layoutmanager.get_instance().get_index_path()
        self._database = WritableDatabase(index_path, xapian.DB_CREATE_OR_OPEN)

    def close_index(self):
        self._database.flush()
        self._database = None

    def remove_index(self):
        index_path = layoutmanager.get_instance().get_index_path()
        if not os.path.exists(index_path):
            return
        for f in os.listdir(index_path):
            os.remove(os.path.join(index_path, f))

    def contains(self, tree_id, version_id):
        postings = self._database.postlist(_PREFIX_TID + tree_id + _PREFIX_VID + version_id)
        try:
            postlist_item = postings.next()
        except StopIteration:
            return False
        return True

    def store(self, tree_id, version_id, properties):
        document = Document()
        document.add_term("%s%s%s%s" % (_PREFIX_TID, tree_id, _PREFIX_VID, version_id))
        document.add_term(_PREFIX_TID + tree_id)
        document.add_term(_PREFIX_VID + version_id)
        document.add_term(_PREFIX_BUNDLE_ID + properties.get('bundle_id', ''))
        document.add_term(_PREFIX_MIME_TYPE + properties.get('mime_type', ''))
        document.add_term(_PREFIX_ACTIVITY_ID +
                          properties.get('activity_id', ''))
        document.add_term(_PREFIX_KEEP + str(properties.get('keep', 0)))

        document.add_value(_VALUE_TID, tree_id)
        document.add_value(_VALUE_VID, version_id)
        document.add_value(_VALUE_CTIME, str(properties['ctime']))

        term_generator = xapian.TermGenerator()

        # TODO: we should do stemming, but in which language?
        #if language is not None:
        #    term_generator.set_stemmer(_xapian.Stem(language))

        # TODO: we should use a stopper
        #if stop is not None:
        #    stopper = _xapian.SimpleStopper()
        #    for term in stop:
        #        stopper.add (term)
        #    term_generator.set_stopper (stopper)

        term_generator.set_document(document)
        term_generator.index_text_without_positions(
                self._extract_text(properties), 1, '')

        if not self.contains(tree_id, version_id):
            self._database.add_document(document)
        else:
            self._database.replace_document("%s%s%s%s" % (_PREFIX_TID, tree_id, _PREFIX_VID, version_id), document)
        self._flush()

    def _extract_text(self, properties):
        text = ''
        for key, value in properties.items():
            if key not in _PROPERTIES_NOT_TO_INDEX:
                if text:
                    text += ' '
                if isinstance(value, unicode):
                    value = value.encode('utf-8')
                elif not isinstance(value, basestring):
                    value = str(value)
                text += value
        return text

    def find(self, query, querystring, options):
        enquire = Enquire(self._database)

        offset = options.pop('offset', 0)
        limit = options.pop('limit', MAX_QUERY_LIMIT)
        all_versions = options.pop('all_versions', False)

        enquire.set_query(self._parse_query(query, querystring))

        # This will assure that the results count is exact.
        check_at_least = offset + limit + 1

        enquire.set_sort_by_value(_VALUE_CTIME, True)

        if not all_versions :
            # only select newest entry (sort order) for each tree_id
            enquire.set_collapse_key(_VALUE_TID)

        query_result = enquire.get_mset(offset, limit, check_at_least)
        total_count = query_result.get_matches_estimated()

        tvids = [(hit.document.get_value(_VALUE_TID),
                  hit.document.get_value(_VALUE_VID))
                 for hit in query_result]

        return (tvids, total_count)

    _queryTermMap = {
        'tree_id': _PREFIX_TID,
        'version_id': _PREFIX_VID,
        'bundle_id': _PREFIX_BUNDLE_ID,
        'activity_id': _PREFIX_ACTIVITY_ID,
        'mime_type': _PREFIX_MIME_TYPE,
        'keep': _PREFIX_KEEP,
    }
    _queryValueMap = {
        'ctime': _VALUE_CTIME,
    }
    def _parse_query_term(self, m_name, prefix, m_value) :
        if isinstance(m_value, list) :
            # list -> multiple terms joined by OR
            return Query(Query.OP_OR,
                [self._parse_query_term(m_name, prefix, word) for word in m_value])

        else :
            # simple query on regular term
            return Query(prefix+str(m_value))

    def _parse_query_value(self, m_name, value_no, m_value) :
        if isinstance(m_value, list) :
            # list -> multiple values joined by OR
            return Query(Query.OP_OR,
                [self._parse_query_value(m_name, value_no, word) for word in m_value])

        elif isinstance(m_value, tuple) :
            # 2-tuple -> range query
            if m_name not in self._valueQueryMap :
                raise NotImplementedError("range queries for %r currently not supported")

            if len(m_value) != 2 :
                raise sugar.datastore.InvalidArgumentError("Only tuples of size 2 have a defined meaning. Did you mean to pass a list instead?")

            start, end = m_value
            return Query(Query.OP_VALUE_RANGE,
                self._queryValueMap[m_name], str(start), str(end))

        else :
            # simple query on value-stored metadata
            return Query(Query.OP_VALUE_RANGE,
                self._queryValueMap[m_name], str(m_value), str(m_value))


    def _parse_query_xapian(self, query_str) :
        # TODO: (re)use long-living instance of QueryParser
        query_parser = QueryParser()
        query_parser.set_database(self._database)
        #query_parser.set_default_op(Query.OP_AND)

        # TODO: we should do stemming, but in which language?
        #query_parser.set_stemmer(_xapian.Stem(lang))
        #query_parser.set_stemming_strategy(qp.STEM_SOME)

        for (m_name, prefix) in self._queryTermMap.items() :
            query_parser.add_prefix(m_name, prefix)

        return query_parser.parse_query(
            query_str,
            QueryParser.FLAG_PHRASE |
                    QueryParser.FLAG_BOOLEAN |
                    QueryParser.FLAG_LOVEHATE |
                    QueryParser.FLAG_WILDCARD,
            '')

    def _parse_query(self, query_dict, query_str):
        logging.debug('_parse_query %r' % query_dict)
        queries = []

        if query_str:
            queries.append(self._parse_query_xapian(query_str))

        # construct queries for term-stored metadata
        queries += [
            self._parse_query_term(m_name, prefix, query_dict.pop(m_name))
            for (m_name, prefix) in self._queryTermMap.items()
            if m_name in query_dict]

        # construct queries for value-stored metadata
        queries += [
            self._parse_query_value(m_name, value_no, query_dict.pop(m_name))
            for (m_name, value_no) in self._queryValueMap.items()
            if m_name in query_dict]

        if not queries:
            queries.append(Query(''))

        if query_dict:
            logging.warning('Unknown term(s): %r' % query_dict)

        return Query(Query.OP_AND, queries)

    def delete(self, tree_id, version_id):
        self._database.delete_document("%s%s%s%s" % (_PREFIX_TID, tree_id, _PREFIX_VID, version_id))

    def get_bundle_ids(self):
        bundle_ids = []
        for term in self._database.allterms(_PREFIX_BUNDLE_ID):
            bundle_ids.append(term.term[len(_PREFIX_BUNDLE_ID):])
        return bundle_ids

    def _flush_timeout_cb(self):
        self._flush(True)
        return False

    def _flush(self, force=False):
        """Called after any database mutation"""
        logging.debug('IndexStore.flush: %r %r' % (force, self._pending_writes))

        if self._flush_timeout is not None:
            gobject.source_remove(self._flush_timeout)
            self._flush_timeout = None

        self._pending_writes += 1
        if force or self._pending_writes > _FLUSH_THRESHOLD:
            self._database.flush()
            self._pending_writes = 0
        else:
            self._flush_timeout = gobject.timeout_add(_FLUSH_TIMEOUT * 1000,
                                                      self._flush_timeout_cb)