Merge lp:~jelmer/bzr-search/hpss into lp:bzr-search

[Jelmer Vernooij]

Add HPSS calls to allow remote access to the bzr-search index, and matching client side calls.

 * Branch.open_index
 * Branch.init_index
 * Index.index_revisions
 * Index.indexed_revisions
 * Index.suggest
 * Index.search

(now updated to not have lp:~jelmer/bzr-search/lazy merged)

[Robert Collins (lifeless)]

Shiny. +1

[Andrew Bennetts (spiv)]

Robert Collins wrote:
> Shiny. +1

I haven't had time to look at the code, but the description is very shiny.
Thanks for making extending HPSS via plugins a reality! :)

-Andrew.

=== modified file '__init__.py'
--- __init__.py	2008-08-28 02:13:47 +0000
+++ __init__.py	2011-11-23 21:50:26 +0000
@@ -38,6 +38,7 @@
 import commands
 import errors
 import index
+from bzrlib.smart.request import request_handlers as smart_request_handlers
 
 
 for command in [
@@ -79,6 +80,25 @@
         index.make_log_search_filter)
 
 
+smart_request_handlers.register_lazy(
+    "Branch.open_index", 'bzrlib.plugins.search.remote',
+    'SmartServerBranchRequestOpenIndex')
+smart_request_handlers.register_lazy(
+    "Branch.init_index", 'bzrlib.plugins.search.remote',
+    'SmartServerBranchRequestInitIndex')
+smart_request_handlers.register_lazy(
+    "Index.index_revisions", 'bzrlib.plugins.search.remote',
+    'SmartServerIndexRequestIndexRevisions')
+smart_request_handlers.register_lazy(
+    "Index.indexed_revisions", 'bzrlib.plugins.search.remote',
+    'SmartServerIndexRequestIndexedRevisions')
+smart_request_handlers.register_lazy(
+    "Index.suggest", 'bzrlib.plugins.search.remote',
+    'SmartServerIndexRequestSuggest')
+smart_request_handlers.register_lazy(
+    "Index.search", 'bzrlib.plugins.search.remote',
+    'SmartServerIndexRequestSearch')
+
 def test_suite():
     # Thunk across to load_tests for niceness with older bzr versions
     from bzrlib.tests import TestLoader

=== modified file 'index.py'
--- index.py	2011-09-19 23:34:37 +0000
+++ index.py	2011-11-23 21:50:26 +0000
@@ -31,8 +31,8 @@
     NotBranchError,
     NoSuchFile,
     UnknownFormatError,
-    IncompatibleAPI,
-)
+    UnknownSmartMethod,
+    )
 from bzrlib.index import CombinedGraphIndex, GraphIndex, InMemoryGraphIndex
 from bzrlib.lockdir import LockDir
 try:
@@ -120,6 +120,14 @@
             transport = transport.clone(path)
             transport.ensure_base()
         index_transport = transport
+    elif getattr(branch.bzrdir, "_call", None) is not None:
+        # FIXME 2011-11-17 JRV: Is there a better way to probe
+        # for smart server branches ?
+        from bzrlib.plugins.search.remote import RemoteIndex
+        try:
+            return RemoteIndex.init(branch)
+        except UnknownSmartMethod:
+            raise errors.CannotIndex(branch)
     else:
         raise errors.CannotIndex(branch)
     lockdir = LockDir(index_transport, 'names-lock')
@@ -205,6 +213,16 @@
         path = 'bzr-search/svn-lookaside/' + uuid + '/' + branch_path
         transport = transport.clone(path)
         commits_only = False
+    elif getattr(branch.bzrdir, "_call", None) is not None:
+        # FIXME 2011-11-17 JRV: Is there a better way to probe
+        # for smart server branches ?
+        from bzrlib.plugins.search.remote import RemoteIndex
+        try:
+            return RemoteIndex.open(branch)
+        except UnknownSmartMethod:
+            # Fall back to traditional methods...
+            transport = branch.bzrdir.transport.clone('bzr-search')
+            commits_only = False
     else:
         transport = branch.bzrdir.transport.clone('bzr-search')
         commits_only = False
@@ -391,7 +409,7 @@
 
     def _add_index(self, builder, to_remove=None, allow_pack=True):
         """Add a new component index to the list of indices.
-        
+
         :param builder: A component builder supporting the upload_index call.
         :param to_remove: An optional iterable of components to remove.
         :param allow_pack: Whether an auto pack is permitted by this operation.

=== added file 'remote.py'
--- remote.py	1970-01-01 00:00:00 +0000
+++ remote.py	2011-11-23 21:50:26 +0000
@@ -0,0 +1,276 @@
+# search, a bzr plugin for searching within bzr branches/repositories.
+# Copyright (C) 2011 Jelmer Vernooij
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as published
+# by the Free Software Foundation.
+#
+# 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
+#
+
+"""Smart server integration for bzr-search."""
+
+from bzrlib import remote
+from bzrlib.controldir import ControlDir
+from bzrlib.errors import (
+    ErrorFromSmartServer,
+    UnexpectedSmartServerResponse,
+    )
+from bzrlib.smart.branch import (
+    SmartServerBranchRequest,
+    )
+from bzrlib.smart.request import SuccessfulSmartServerResponse
+
+from bzrlib.plugins.search import errors, index
+
+
+def _encode_termlist(termlist):
+    return ["\0".join([k.encode('utf-8') for k in term]) for term in termlist]
+
+def _decode_termlist(termlist):
+    return [tuple([k.decode('utf-8') for k in term.split('\0')]) for term in termlist]
+
+
+class RemoteIndex(object):
+    """Index accessed over a smart server."""
+
+    def __init__(self, client, path, branch=None):
+        self._client = client
+        self._path = path
+        self._branch = branch
+
+    def _call(self, method, *args, **err_context):
+        try:
+            return self._client.call(method, *args)
+        except ErrorFromSmartServer, err:
+            self._translate_error(err, **err_context)
+
+    def _call_expecting_body(self, method, *args, **err_context):
+        try:
+            return self._client.call_expecting_body(method, *args)
+        except ErrorFromSmartServer, err:
+            self._translate_error(err, **err_context)
+
+    def _call_with_body_bytes(self, method, args, body_bytes, **err_context):
+        try:
+            return self._client.call_with_body_bytes(method, args, body_bytes)
+        except ErrorFromSmartServer, err:
+            self._translate_error(err, **err_context)
+
+    def _call_with_body_bytes_expecting_body(self, method, args, body_bytes,
+                                             **err_context):
+        try:
+            return self._client.call_with_body_bytes_expecting_body(
+                method, args, body_bytes)
+        except errors.ErrorFromSmartServer, err:
+            self._translate_error(err, **err_context)
+
+    def _translate_error(self, err, **context):
+        remote._translate_error(err, index=self, **context)
+
+    @classmethod
+    def open(cls, branch):
+        # This might raise UnknownSmartMethod,
+        # but the caller should handle that.
+        response = branch._call("Branch.open_index",
+            branch._remote_path())
+        if response == ('no', ):
+            raise errors.NoSearchIndex(branch.user_transport)
+        if response != ('yes', ):
+            raise UnexpectedSmartServerResponse(response)
+        return RemoteIndex(branch._client, branch._remote_path(), branch)
+
+    @classmethod
+    def init(cls, branch):
+        response = branch._call("Branch.init_index",
+            branch._remote_path())
+        if response != ('ok', ):
+            raise errors.UnexpectedSmartServerResponse(response)
+        return RemoteIndex(branch._client, branch._remote_path(), branch)
+
+    def index_branch(self, branch, tip_revision):
+        """Index revisions from a branch.
+
+        :param branch: The branch to index.
+        :param tip_revision: The tip of the branch.
+        """
+        self.index_revisions(branch, [tip_revision])
+
+    def index_revisions(self, branch, revisions_to_index):
+        """Index some revisions from branch.
+
+        :param branch: A branch to index.
+        :param revisions_to_index: A set of revision ids to index.
+        """
+        body = "\n".join(revisions_to_index)
+        response = self._call_with_body_bytes(
+            'Index.index_revisions', (self._path, branch._remote_path(),),
+            body)
+        if response != ('ok', ):
+            raise errors.UnexpectedSmartServerResponse(response)
+
+    def indexed_revisions(self):
+        """Return the revision_keys that this index contains terms for."""
+        response, handler = self._call_expecting_body(
+            'Index.indexed_revisions', self._path)
+        if response != ('ok', ):
+            raise errors.UnexpectedSmartServerResponse(response)
+        byte_stream = handler.read_streamed_body()
+        data = ""
+        for bytes in byte_stream:
+            data += bytes
+            lines = data.split("\n")
+            data = lines.pop()
+            for revid in lines:
+                yield (revid, )
+
+    def search(self, termlist):
+        """Trivial set-based search of the index.
+
+        :param termlist: A list of terms.
+        :return: An iterator of SearchResults for documents indexed by all
+            terms in the termlist.
+        """
+        index._ensure_regexes()
+        response, handler = self._call_expecting_body('Index.search',
+            self._path, _encode_termlist(termlist))
+        if response != ('ok', ):
+            raise errors.UnexpectedSmartServerResponse(response)
+        byte_stream = handler.read_streamed_body()
+        data = ""
+        ret = []
+        for bytes in byte_stream:
+            data += bytes
+            lines = data.split("\n")
+            data = lines.pop()
+            for l in lines:
+                if l[0] == 'r':
+                    hit = index.RevisionHit(self._branch.repository, (l[1:], ))
+                elif l[0] == 't':
+                    hit = index.FileTextHit(self, self._branch.repository,
+                        tuple(l[1:].split("\0")), termlist)
+                elif l[0] == 'p':
+                    hit = index.PathHit(l[1:])
+                else:
+                    raise AssertionError("Unknown hit kind %r" % l[0])
+            # We can't yield, since the caller might try to look up results
+            # over the same medium.
+            ret.append(hit)
+        return iter(ret)
+
+    def suggest(self, termlist):
+        """Generate suggestions for extending a search.
+
+        :param termlist: A list of terms.
+        :return: An iterator of terms that start with the last search term in
+            termlist, and match the rest of the search.
+        """
+        response = self._call('Index.suggest',
+            self._path, _encode_termlist(termlist))
+        if response[0] != 'ok':
+            raise UnexpectedSmartServerResponse(response)
+        return [(suggestion.decode('utf-8'),) for suggestion in response[1]]
+
+
+class SmartServerBranchRequestOpenIndex(SmartServerBranchRequest):
+    """Open an index file."""
+
+    def do_with_branch(self, branch):
+        """open an index."""
+        try:
+            idx = index.open_index_branch(branch)
+        except errors.NoSearchIndex:
+            return SuccessfulSmartServerResponse(('no', ))
+        else:
+            return SuccessfulSmartServerResponse(('yes', ))
+
+
+class SmartServerBranchRequestInitIndex(SmartServerBranchRequest):
+    """Create an index."""
+
+    def do_with_branch(self, branch, format=None):
+        """Create an index."""
+        if format is None:
+            idx = index.init_index(branch)
+        else:
+            idx = index.init_index(branch, format)
+        return SuccessfulSmartServerResponse(('ok', ))
+
+
+class SmartServerIndexRequest(SmartServerBranchRequest):
+    """Base class for index requests."""
+
+    def do_with_branch(self, branch, *args):
+        idx = index.open_index_branch(branch)
+        return self.do_with_index(idx, *args)
+
+    def do_with_index(self, index, *args):
+        raise NotImplementedError(self.do_with_index)
+
+
+class SmartServerIndexRequestIndexRevisions(SmartServerIndexRequest):
+    """Index a set of revisions."""
+
+    def do_body(self, body_bytes):
+        revids = body_bytes.split("\n")
+        self._index.index_revisions(self._branch, revids)
+        return SuccessfulSmartServerResponse(('ok', ))
+
+    def do_with_index(self, index, branch_path):
+        self._index = index
+        transport = self.transport_from_client_path(branch_path)
+        controldir = ControlDir.open_from_transport(transport)
+        if controldir.get_branch_reference() is not None:
+            raise errors.NotBranchError(transport.base)
+        self._branch = controldir.open_branch(ignore_fallbacks=True)
+        # Indicate we want a body
+        return None
+
+
+class SmartServerIndexRequestIndexedRevisions(SmartServerIndexRequest):
+    """Retrieve the set of revisions in the index."""
+
+    def body_stream(self, index):
+        for revid in index.indexed_revisions():
+            yield "%s\n" % "\0".join(revid)
+
+    def do_with_index(self, index):
+        return SuccessfulSmartServerResponse(('ok', ),
+            body_stream=self.body_stream(index))
+
+
+class SmartServerIndexRequestSuggest(SmartServerIndexRequest):
+    """Suggest alternative terms."""
+
+    def do_with_index(self, index, termlist):
+        suggestions = index.suggest(_decode_termlist(termlist))
+        return SuccessfulSmartServerResponse(
+            ('ok',
+                [suggestion.encode('utf-8') for (suggestion,) in suggestions]))
+
+
+class SmartServerIndexRequestSearch(SmartServerIndexRequest):
+    """Search for terms."""
+
+    def body_stream(self, results):
+        for hit in results:
+            if isinstance(hit, index.FileTextHit):
+                yield "t%s\0%s\n" % hit.text_key
+            elif isinstance(hit, index.RevisionHit):
+                yield "r%s\n" % hit.revision_key[0]
+            elif isinstance(hit, index.PathHit):
+                yield "p%s\n" % hit.path_utf8
+            else:
+                raise AssertionError("Unknown hit type %r" % hit)
+
+    def do_with_index(self, index, termlist):
+        results = index.search(_decode_termlist(termlist))
+        return SuccessfulSmartServerResponse(
+            ('ok',), body_stream=self.body_stream(results))

=== modified file 'tests/__init__.py'
--- tests/__init__.py	2008-06-14 05:07:54 +0000
+++ tests/__init__.py	2011-11-23 21:50:26 +0000
@@ -25,6 +25,7 @@
         'errors',
         'index',
         'inventory',
+        'remote',
         'transport',
         ]
     standard_tests.addTests(loader.loadTestsFromModuleNames(

=== added file 'tests/test_remote.py'
--- tests/test_remote.py	1970-01-01 00:00:00 +0000
+++ tests/test_remote.py	2011-11-23 21:50:26 +0000
@@ -0,0 +1,128 @@
+# search, a bzr plugin for searching within bzr branches/repositories.
+# Copyright (C) 2011 Jelmer Vernooij
+# 
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as published
+# by the Free Software Foundation.
+# 
+# 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
+# 
+
+"""Tests for the smart server verbs."""
+
+from bzrlib import tests
+from bzrlib.branch import Branch
+from bzrlib.smart import (
+    request as smart_req,
+    )
+
+from bzrlib.plugins.search import (
+    errors,
+    index,
+    )
+from bzrlib.plugins.search.remote import (
+    RemoteIndex,
+    SmartServerBranchRequestOpenIndex,
+    )
+
+
+class TestSmartServerBranchRequestOpenIndex(
+    tests.TestCaseWithMemoryTransport):
+
+    def test_missing(self):
+        """For an empty branch, the result is ('no', )."""
+        backing = self.get_transport()
+        request = SmartServerBranchRequestOpenIndex(backing)
+        self.make_branch('.')
+        self.assertEqual(smart_req.SmartServerResponse(('no', )),
+            request.execute(''))
+
+    def test_present(self):
+        """For a branch with an index, ('yes', ) is returned."""
+        backing = self.get_transport()
+        request = SmartServerBranchRequestOpenIndex(backing)
+        b = self.make_branch('.')
+        index.init_index(b)
+        self.assertEqual(smart_req.SmartServerResponse(('yes', )),
+            request.execute(''))
+
+
+class TestRemoteIndex(tests.TestCaseWithTransport):
+
+    def test_no_index(self):
+        local_branch = self.make_branch('.')
+        remote_transport = self.make_smart_server('.')
+        remote_branch = Branch.open_from_transport(remote_transport)
+        self.assertRaises(errors.NoSearchIndex, RemoteIndex.open,
+            remote_branch)
+
+    def test_open(self):
+        local_branch = self.make_branch('.')
+        index.init_index(local_branch)
+        remote_transport = self.make_smart_server('.')
+        remote_branch = Branch.open_from_transport(remote_transport)
+        idx = RemoteIndex.open(remote_branch)
+        self.assertIsInstance(idx, RemoteIndex)
+
+    def test_init(self):
+        local_branch = self.make_branch('.')
+        remote_transport = self.make_smart_server('.')
+        remote_branch = Branch.open_from_transport(remote_transport)
+        idx = index.init_index(remote_branch)
+        self.assertIsInstance(idx, RemoteIndex)
+
+    def test_init_exists(self):
+        local_branch = self.make_branch('.')
+        index.init_index(local_branch)
+        remote_transport = self.make_smart_server('.')
+        remote_branch = Branch.open_from_transport(remote_transport)
+        #self.assertRaises( index.init_index, remote_branch)
+
+
+class TestWithRemoteIndex(tests.TestCaseWithTransport):
+
+    def make_remote_index(self):
+        tree = self.make_branch_and_tree('.')
+        local_branch = tree.branch
+        index.init_index(local_branch)
+        remote_transport = self.make_smart_server('.')
+        remote_branch = Branch.open_from_transport(remote_transport)
+        return tree, remote_branch, RemoteIndex.open(remote_branch)
+
+    def test_index_revisions(self):
+        tree, branch, index = self.make_remote_index()
+        tree.commit(message="message", rev_id='revid1')
+        index.index_revisions(branch, ['revid1'])
+        self.assertEquals([('revid1',)], list(index.indexed_revisions()))
+
+    def test_indexed_revisions(self):
+        tree, branch, remote_index = self.make_remote_index()
+        tree.commit(message="message", rev_id='revid1')
+        self.assertEquals([], list(remote_index.indexed_revisions()))
+        local_index = index.open_index_branch(tree.branch)
+        local_index.index_revisions(tree.branch, ['revid1'])
+        self.assertEquals([('revid1',)], list(remote_index.indexed_revisions()))
+
+    def test_suggest(self):
+        tree, branch, remote_index = self.make_remote_index()
+        tree.commit(message="first", rev_id='revid1')
+        local_index = index.open_index_branch(tree.branch)
+        local_index.index_revisions(tree.branch, ['revid1'])
+        self.assertEquals([(u'first',)], list(remote_index.suggest([(u'f',)])))
+
+    def test_search(self):
+        tree, branch, remote_index = self.make_remote_index()
+        # The double-space is a cheap smoke test for the tokeniser.
+        revid = tree.commit('first  post')
+        remote_index.index_revisions(branch, [revid])
+        results = list(remote_index.search([('post',)]))
+        self.assertEqual(1, len(results))
+        self.assertIsInstance(results[0], index.RevisionHit)
+        self.assertEqual((revid,), results[0].revision_key)