~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

  • Committer: Aaron Bentley
  • Date: 2006-06-21 14:30:57 UTC
  • mfrom: (1801.1.1 bzr.dev)
  • mto: This revision was merged to the branch mainline in revision 1803.
  • Revision ID: abentley@panoramicfeedback.com-20060621143057-776e4b8d707e430e
Install benchmarks. (Jelmer Vernooij)

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2006, 2007 Canonical Ltd
2
 
#
3
 
# This program is free software; you can redistribute it and/or modify
4
 
# it under the terms of the GNU General Public License as published by
5
 
# the Free Software Foundation; either version 2 of the License, or
6
 
# (at your option) any later version.
7
 
#
8
 
# This program is distributed in the hope that it will be useful,
9
 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
10
 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11
 
# GNU General Public License for more details.
12
 
#
13
 
# You should have received a copy of the GNU General Public License
14
 
# along with this program; if not, write to the Free Software
15
 
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
16
 
 
17
 
"""Server-side repository related request implmentations."""
18
 
 
19
 
import bz2
20
 
import os
21
 
import Queue
22
 
import sys
23
 
import tarfile
24
 
import tempfile
25
 
import threading
26
 
 
27
 
from bzrlib import (
28
 
    bencode,
29
 
    errors,
30
 
    graph,
31
 
    osutils,
32
 
    pack,
33
 
    versionedfile,
34
 
    )
35
 
from bzrlib.bzrdir import BzrDir
36
 
from bzrlib.smart.request import (
37
 
    FailedSmartServerResponse,
38
 
    SmartServerRequest,
39
 
    SuccessfulSmartServerResponse,
40
 
    )
41
 
from bzrlib.repository import _strip_NULL_ghosts, network_format_registry
42
 
from bzrlib import revision as _mod_revision
43
 
from bzrlib.versionedfile import (
44
 
    NetworkRecordStream,
45
 
    record_to_fulltext_bytes,
46
 
    )
47
 
 
48
 
 
49
 
class SmartServerRepositoryRequest(SmartServerRequest):
50
 
    """Common base class for Repository requests."""
51
 
 
52
 
    def do(self, path, *args):
53
 
        """Execute a repository request.
54
 
 
55
 
        All Repository requests take a path to the repository as their first
56
 
        argument.  The repository must be at the exact path given by the
57
 
        client - no searching is done.
58
 
 
59
 
        The actual logic is delegated to self.do_repository_request.
60
 
 
61
 
        :param client_path: The path for the repository as received from the
62
 
            client.
63
 
        :return: A SmartServerResponse from self.do_repository_request().
64
 
        """
65
 
        transport = self.transport_from_client_path(path)
66
 
        bzrdir = BzrDir.open_from_transport(transport)
67
 
        # Save the repository for use with do_body.
68
 
        self._repository = bzrdir.open_repository()
69
 
        return self.do_repository_request(self._repository, *args)
70
 
 
71
 
    def do_repository_request(self, repository, *args):
72
 
        """Override to provide an implementation for a verb."""
73
 
        # No-op for verbs that take bodies (None as a result indicates a body
74
 
        # is expected)
75
 
        return None
76
 
 
77
 
    def recreate_search(self, repository, search_bytes, discard_excess=False):
78
 
        """Recreate a search from its serialised form.
79
 
 
80
 
        :param discard_excess: If True, and the search refers to data we don't
81
 
            have, just silently accept that fact - the verb calling
82
 
            recreate_search trusts that clients will look for missing things
83
 
            they expected and get it from elsewhere.
84
 
        """
85
 
        lines = search_bytes.split('\n')
86
 
        if lines[0] == 'ancestry-of':
87
 
            heads = lines[1:]
88
 
            search_result = graph.PendingAncestryResult(heads, repository)
89
 
            return search_result, None
90
 
        elif lines[0] == 'search':
91
 
            return self.recreate_search_from_recipe(repository, lines[1:],
92
 
                discard_excess=discard_excess)
93
 
        else:
94
 
            return (None, FailedSmartServerResponse(('BadSearch',)))
95
 
 
96
 
    def recreate_search_from_recipe(self, repository, lines,
97
 
        discard_excess=False):
98
 
        """Recreate a specific revision search (vs a from-tip search).
99
 
 
100
 
        :param discard_excess: If True, and the search refers to data we don't
101
 
            have, just silently accept that fact - the verb calling
102
 
            recreate_search trusts that clients will look for missing things
103
 
            they expected and get it from elsewhere.
104
 
        """
105
 
        start_keys = set(lines[0].split(' '))
106
 
        exclude_keys = set(lines[1].split(' '))
107
 
        revision_count = int(lines[2])
108
 
        repository.lock_read()
109
 
        try:
110
 
            search = repository.get_graph()._make_breadth_first_searcher(
111
 
                start_keys)
112
 
            while True:
113
 
                try:
114
 
                    next_revs = search.next()
115
 
                except StopIteration:
116
 
                    break
117
 
                search.stop_searching_any(exclude_keys.intersection(next_revs))
118
 
            search_result = search.get_result()
119
 
            if (not discard_excess and
120
 
                search_result.get_recipe()[3] != revision_count):
121
 
                # we got back a different amount of data than expected, this
122
 
                # gets reported as NoSuchRevision, because less revisions
123
 
                # indicates missing revisions, and more should never happen as
124
 
                # the excludes list considers ghosts and ensures that ghost
125
 
                # filling races are not a problem.
126
 
                return (None, FailedSmartServerResponse(('NoSuchRevision',)))
127
 
            return (search_result, None)
128
 
        finally:
129
 
            repository.unlock()
130
 
 
131
 
 
132
 
class SmartServerRepositoryReadLocked(SmartServerRepositoryRequest):
133
 
    """Calls self.do_readlocked_repository_request."""
134
 
 
135
 
    def do_repository_request(self, repository, *args):
136
 
        """Read lock a repository for do_readlocked_repository_request."""
137
 
        repository.lock_read()
138
 
        try:
139
 
            return self.do_readlocked_repository_request(repository, *args)
140
 
        finally:
141
 
            repository.unlock()
142
 
 
143
 
 
144
 
class SmartServerRepositoryGetParentMap(SmartServerRepositoryRequest):
145
 
    """Bzr 1.2+ - get parent data for revisions during a graph search."""
146
 
 
147
 
    no_extra_results = False
148
 
 
149
 
    def do_repository_request(self, repository, *revision_ids):
150
 
        """Get parent details for some revisions.
151
 
 
152
 
        All the parents for revision_ids are returned. Additionally up to 64KB
153
 
        of additional parent data found by performing a breadth first search
154
 
        from revision_ids is returned. The verb takes a body containing the
155
 
        current search state, see do_body for details.
156
 
 
157
 
        If 'include-missing:' is in revision_ids, ghosts encountered in the
158
 
        graph traversal for getting parent data are included in the result with
159
 
        a prefix of 'missing:'.
160
 
 
161
 
        :param repository: The repository to query in.
162
 
        :param revision_ids: The utf8 encoded revision_id to answer for.
163
 
        """
164
 
        self._revision_ids = revision_ids
165
 
        return None # Signal that we want a body.
166
 
 
167
 
    def do_body(self, body_bytes):
168
 
        """Process the current search state and perform the parent lookup.
169
 
 
170
 
        :return: A smart server response where the body contains an utf8
171
 
            encoded flattened list of the parents of the revisions (the same
172
 
            format as Repository.get_revision_graph) which has been bz2
173
 
            compressed.
174
 
        """
175
 
        repository = self._repository
176
 
        repository.lock_read()
177
 
        try:
178
 
            return self._do_repository_request(body_bytes)
179
 
        finally:
180
 
            repository.unlock()
181
 
 
182
 
    def _do_repository_request(self, body_bytes):
183
 
        repository = self._repository
184
 
        revision_ids = set(self._revision_ids)
185
 
        include_missing = 'include-missing:' in revision_ids
186
 
        if include_missing:
187
 
            revision_ids.remove('include-missing:')
188
 
        body_lines = body_bytes.split('\n')
189
 
        search_result, error = self.recreate_search_from_recipe(
190
 
            repository, body_lines)
191
 
        if error is not None:
192
 
            return error
193
 
        # TODO might be nice to start up the search again; but thats not
194
 
        # written or tested yet.
195
 
        client_seen_revs = set(search_result.get_keys())
196
 
        # Always include the requested ids.
197
 
        client_seen_revs.difference_update(revision_ids)
198
 
        lines = []
199
 
        repo_graph = repository.get_graph()
200
 
        result = {}
201
 
        queried_revs = set()
202
 
        size_so_far = 0
203
 
        next_revs = revision_ids
204
 
        first_loop_done = False
205
 
        while next_revs:
206
 
            queried_revs.update(next_revs)
207
 
            parent_map = repo_graph.get_parent_map(next_revs)
208
 
            current_revs = next_revs
209
 
            next_revs = set()
210
 
            for revision_id in current_revs:
211
 
                missing_rev = False
212
 
                parents = parent_map.get(revision_id)
213
 
                if parents is not None:
214
 
                    # adjust for the wire
215
 
                    if parents == (_mod_revision.NULL_REVISION,):
216
 
                        parents = ()
217
 
                    # prepare the next query
218
 
                    next_revs.update(parents)
219
 
                    encoded_id = revision_id
220
 
                else:
221
 
                    missing_rev = True
222
 
                    encoded_id = "missing:" + revision_id
223
 
                    parents = []
224
 
                if (revision_id not in client_seen_revs and
225
 
                    (not missing_rev or include_missing)):
226
 
                    # Client does not have this revision, give it to it.
227
 
                    # add parents to the result
228
 
                    result[encoded_id] = parents
229
 
                    # Approximate the serialized cost of this revision_id.
230
 
                    size_so_far += 2 + len(encoded_id) + sum(map(len, parents))
231
 
            # get all the directly asked for parents, and then flesh out to
232
 
            # 64K (compressed) or so. We do one level of depth at a time to
233
 
            # stay in sync with the client. The 250000 magic number is
234
 
            # estimated compression ratio taken from bzr.dev itself.
235
 
            if self.no_extra_results or (
236
 
                first_loop_done and size_so_far > 250000):
237
 
                next_revs = set()
238
 
                break
239
 
            # don't query things we've already queried
240
 
            next_revs.difference_update(queried_revs)
241
 
            first_loop_done = True
242
 
 
243
 
        # sorting trivially puts lexographically similar revision ids together.
244
 
        # Compression FTW.
245
 
        for revision, parents in sorted(result.items()):
246
 
            lines.append(' '.join((revision, ) + tuple(parents)))
247
 
 
248
 
        return SuccessfulSmartServerResponse(
249
 
            ('ok', ), bz2.compress('\n'.join(lines)))
250
 
 
251
 
 
252
 
class SmartServerRepositoryGetRevisionGraph(SmartServerRepositoryReadLocked):
253
 
 
254
 
    def do_readlocked_repository_request(self, repository, revision_id):
255
 
        """Return the result of repository.get_revision_graph(revision_id).
256
 
 
257
 
        Deprecated as of bzr 1.4, but supported for older clients.
258
 
 
259
 
        :param repository: The repository to query in.
260
 
        :param revision_id: The utf8 encoded revision_id to get a graph from.
261
 
        :return: A smart server response where the body contains an utf8
262
 
            encoded flattened list of the revision graph.
263
 
        """
264
 
        if not revision_id:
265
 
            revision_id = None
266
 
 
267
 
        lines = []
268
 
        graph = repository.get_graph()
269
 
        if revision_id:
270
 
            search_ids = [revision_id]
271
 
        else:
272
 
            search_ids = repository.all_revision_ids()
273
 
        search = graph._make_breadth_first_searcher(search_ids)
274
 
        transitive_ids = set()
275
 
        map(transitive_ids.update, list(search))
276
 
        parent_map = graph.get_parent_map(transitive_ids)
277
 
        revision_graph = _strip_NULL_ghosts(parent_map)
278
 
        if revision_id and revision_id not in revision_graph:
279
 
            # Note that we return an empty body, rather than omitting the body.
280
 
            # This way the client knows that it can always expect to find a body
281
 
            # in the response for this method, even in the error case.
282
 
            return FailedSmartServerResponse(('nosuchrevision', revision_id), '')
283
 
 
284
 
        for revision, parents in revision_graph.items():
285
 
            lines.append(' '.join((revision, ) + tuple(parents)))
286
 
 
287
 
        return SuccessfulSmartServerResponse(('ok', ), '\n'.join(lines))
288
 
 
289
 
 
290
 
class SmartServerRepositoryGetRevIdForRevno(SmartServerRepositoryReadLocked):
291
 
 
292
 
    def do_readlocked_repository_request(self, repository, revno,
293
 
            known_pair):
294
 
        """Find the revid for a given revno, given a known revno/revid pair.
295
 
        
296
 
        New in 1.17.
297
 
        """
298
 
        try:
299
 
            found_flag, result = repository.get_rev_id_for_revno(revno, known_pair)
300
 
        except errors.RevisionNotPresent, err:
301
 
            if err.revision_id != known_pair[1]:
302
 
                raise AssertionError(
303
 
                    'get_rev_id_for_revno raised RevisionNotPresent for '
304
 
                    'non-initial revision: ' + err.revision_id)
305
 
            return FailedSmartServerResponse(
306
 
                ('nosuchrevision', err.revision_id))
307
 
        if found_flag:
308
 
            return SuccessfulSmartServerResponse(('ok', result))
309
 
        else:
310
 
            earliest_revno, earliest_revid = result
311
 
            return SuccessfulSmartServerResponse(
312
 
                ('history-incomplete', earliest_revno, earliest_revid))
313
 
 
314
 
 
315
 
class SmartServerRequestHasRevision(SmartServerRepositoryRequest):
316
 
 
317
 
    def do_repository_request(self, repository, revision_id):
318
 
        """Return ok if a specific revision is in the repository at path.
319
 
 
320
 
        :param repository: The repository to query in.
321
 
        :param revision_id: The utf8 encoded revision_id to lookup.
322
 
        :return: A smart server response of ('ok', ) if the revision is
323
 
            present.
324
 
        """
325
 
        if repository.has_revision(revision_id):
326
 
            return SuccessfulSmartServerResponse(('yes', ))
327
 
        else:
328
 
            return SuccessfulSmartServerResponse(('no', ))
329
 
 
330
 
 
331
 
class SmartServerRepositoryGatherStats(SmartServerRepositoryRequest):
332
 
 
333
 
    def do_repository_request(self, repository, revid, committers):
334
 
        """Return the result of repository.gather_stats().
335
 
 
336
 
        :param repository: The repository to query in.
337
 
        :param revid: utf8 encoded rev id or an empty string to indicate None
338
 
        :param committers: 'yes' or 'no'.
339
 
 
340
 
        :return: A SmartServerResponse ('ok',), a encoded body looking like
341
 
              committers: 1
342
 
              firstrev: 1234.230 0
343
 
              latestrev: 345.700 3600
344
 
              revisions: 2
345
 
 
346
 
              But containing only fields returned by the gather_stats() call
347
 
        """
348
 
        if revid == '':
349
 
            decoded_revision_id = None
350
 
        else:
351
 
            decoded_revision_id = revid
352
 
        if committers == 'yes':
353
 
            decoded_committers = True
354
 
        else:
355
 
            decoded_committers = None
356
 
        stats = repository.gather_stats(decoded_revision_id, decoded_committers)
357
 
 
358
 
        body = ''
359
 
        if stats.has_key('committers'):
360
 
            body += 'committers: %d\n' % stats['committers']
361
 
        if stats.has_key('firstrev'):
362
 
            body += 'firstrev: %.3f %d\n' % stats['firstrev']
363
 
        if stats.has_key('latestrev'):
364
 
             body += 'latestrev: %.3f %d\n' % stats['latestrev']
365
 
        if stats.has_key('revisions'):
366
 
            body += 'revisions: %d\n' % stats['revisions']
367
 
        if stats.has_key('size'):
368
 
            body += 'size: %d\n' % stats['size']
369
 
 
370
 
        return SuccessfulSmartServerResponse(('ok', ), body)
371
 
 
372
 
 
373
 
class SmartServerRepositoryIsShared(SmartServerRepositoryRequest):
374
 
 
375
 
    def do_repository_request(self, repository):
376
 
        """Return the result of repository.is_shared().
377
 
 
378
 
        :param repository: The repository to query in.
379
 
        :return: A smart server response of ('yes', ) if the repository is
380
 
            shared, and ('no', ) if it is not.
381
 
        """
382
 
        if repository.is_shared():
383
 
            return SuccessfulSmartServerResponse(('yes', ))
384
 
        else:
385
 
            return SuccessfulSmartServerResponse(('no', ))
386
 
 
387
 
 
388
 
class SmartServerRepositoryLockWrite(SmartServerRepositoryRequest):
389
 
 
390
 
    def do_repository_request(self, repository, token=''):
391
 
        # XXX: this probably should not have a token.
392
 
        if token == '':
393
 
            token = None
394
 
        try:
395
 
            token = repository.lock_write(token=token)
396
 
        except errors.LockContention, e:
397
 
            return FailedSmartServerResponse(('LockContention',))
398
 
        except errors.UnlockableTransport:
399
 
            return FailedSmartServerResponse(('UnlockableTransport',))
400
 
        except errors.LockFailed, e:
401
 
            return FailedSmartServerResponse(('LockFailed',
402
 
                str(e.lock), str(e.why)))
403
 
        if token is not None:
404
 
            repository.leave_lock_in_place()
405
 
        repository.unlock()
406
 
        if token is None:
407
 
            token = ''
408
 
        return SuccessfulSmartServerResponse(('ok', token))
409
 
 
410
 
 
411
 
class SmartServerRepositoryGetStream(SmartServerRepositoryRequest):
412
 
 
413
 
    def do_repository_request(self, repository, to_network_name):
414
 
        """Get a stream for inserting into a to_format repository.
415
 
 
416
 
        :param repository: The repository to stream from.
417
 
        :param to_network_name: The network name of the format of the target
418
 
            repository.
419
 
        """
420
 
        self._to_format = network_format_registry.get(to_network_name)
421
 
        if self._should_fake_unknown():
422
 
            return FailedSmartServerResponse(
423
 
                ('UnknownMethod', 'Repository.get_stream'))
424
 
        return None # Signal that we want a body.
425
 
 
426
 
    def _should_fake_unknown(self):
427
 
        """Return True if we should return UnknownMethod to the client.
428
 
        
429
 
        This is a workaround for bugs in pre-1.19 clients that claim to
430
 
        support receiving streams of CHK repositories.  The pre-1.19 client
431
 
        expects inventory records to be serialized in the format defined by
432
 
        to_network_name, but in pre-1.19 (at least) that format definition
433
 
        tries to use the xml5 serializer, which does not correctly handle
434
 
        rich-roots.  After 1.19 the client can also accept inventory-deltas
435
 
        (which avoids this issue), and those clients will use the
436
 
        Repository.get_stream_1.19 verb instead of this one.
437
 
        So: if this repository is CHK, and the to_format doesn't match,
438
 
        we should just fake an UnknownSmartMethod error so that the client
439
 
        will fallback to VFS, rather than sending it a stream we know it
440
 
        cannot handle.
441
 
        """
442
 
        from_format = self._repository._format
443
 
        to_format = self._to_format
444
 
        if not from_format.supports_chks:
445
 
            # Source not CHK: that's ok
446
 
            return False
447
 
        if (to_format.supports_chks and
448
 
            from_format.repository_class is to_format.repository_class and
449
 
            from_format._serializer == to_format._serializer):
450
 
            # Source is CHK, but target matches: that's ok
451
 
            # (e.g. 2a->2a, or CHK2->2a)
452
 
            return False
453
 
        # Source is CHK, and target is not CHK or incompatible CHK.  We can't
454
 
        # generate a compatible stream.
455
 
        return True
456
 
 
457
 
    def do_body(self, body_bytes):
458
 
        repository = self._repository
459
 
        repository.lock_read()
460
 
        try:
461
 
            search_result, error = self.recreate_search(repository, body_bytes,
462
 
                discard_excess=True)
463
 
            if error is not None:
464
 
                repository.unlock()
465
 
                return error
466
 
            source = repository._get_source(self._to_format)
467
 
            stream = source.get_stream(search_result)
468
 
        except Exception:
469
 
            exc_info = sys.exc_info()
470
 
            try:
471
 
                # On non-error, unlocking is done by the body stream handler.
472
 
                repository.unlock()
473
 
            finally:
474
 
                raise exc_info[0], exc_info[1], exc_info[2]
475
 
        return SuccessfulSmartServerResponse(('ok',),
476
 
            body_stream=self.body_stream(stream, repository))
477
 
 
478
 
    def body_stream(self, stream, repository):
479
 
        byte_stream = _stream_to_byte_stream(stream, repository._format)
480
 
        try:
481
 
            for bytes in byte_stream:
482
 
                yield bytes
483
 
        except errors.RevisionNotPresent, e:
484
 
            # This shouldn't be able to happen, but as we don't buffer
485
 
            # everything it can in theory happen.
486
 
            repository.unlock()
487
 
            yield FailedSmartServerResponse(('NoSuchRevision', e.revision_id))
488
 
        else:
489
 
            repository.unlock()
490
 
 
491
 
 
492
 
class SmartServerRepositoryGetStream_1_19(SmartServerRepositoryGetStream):
493
 
 
494
 
    def _should_fake_unknown(self):
495
 
        """Returns False; we don't need to workaround bugs in 1.19+ clients."""
496
 
        return False
497
 
 
498
 
 
499
 
def _stream_to_byte_stream(stream, src_format):
500
 
    """Convert a record stream to a self delimited byte stream."""
501
 
    pack_writer = pack.ContainerSerialiser()
502
 
    yield pack_writer.begin()
503
 
    yield pack_writer.bytes_record(src_format.network_name(), '')
504
 
    for substream_type, substream in stream:
505
 
        for record in substream:
506
 
            if record.storage_kind in ('chunked', 'fulltext'):
507
 
                serialised = record_to_fulltext_bytes(record)
508
 
            elif record.storage_kind == 'inventory-delta':
509
 
                serialised = record_to_inventory_delta_bytes(record)
510
 
            elif record.storage_kind == 'absent':
511
 
                raise ValueError("Absent factory for %s" % (record.key,))
512
 
            else:
513
 
                serialised = record.get_bytes_as(record.storage_kind)
514
 
            if serialised:
515
 
                # Some streams embed the whole stream into the wire
516
 
                # representation of the first record, which means that
517
 
                # later records have no wire representation: we skip them.
518
 
                yield pack_writer.bytes_record(serialised, [(substream_type,)])
519
 
    yield pack_writer.end()
520
 
 
521
 
 
522
 
class _ByteStreamDecoder(object):
523
 
    """Helper for _byte_stream_to_stream.
524
 
 
525
 
    The expected usage of this class is via the function _byte_stream_to_stream
526
 
    which creates a _ByteStreamDecoder, pops off the stream format and then
527
 
    yields the output of record_stream(), the main entry point to
528
 
    _ByteStreamDecoder.
529
 
 
530
 
    Broadly this class has to unwrap two layers of iterators:
531
 
    (type, substream)
532
 
    (substream details)
533
 
 
534
 
    This is complicated by wishing to return type, iterator_for_type, but
535
 
    getting the data for iterator_for_type when we find out type: we can't
536
 
    simply pass a generator down to the NetworkRecordStream parser, instead
537
 
    we have a little local state to seed each NetworkRecordStream instance,
538
 
    and gather the type that we'll be yielding.
539
 
 
540
 
    :ivar byte_stream: The byte stream being decoded.
541
 
    :ivar stream_decoder: A pack parser used to decode the bytestream
542
 
    :ivar current_type: The current type, used to join adjacent records of the
543
 
        same type into a single stream.
544
 
    :ivar first_bytes: The first bytes to give the next NetworkRecordStream.
545
 
    """
546
 
 
547
 
    def __init__(self, byte_stream):
548
 
        """Create a _ByteStreamDecoder."""
549
 
        self.stream_decoder = pack.ContainerPushParser()
550
 
        self.current_type = None
551
 
        self.first_bytes = None
552
 
        self.byte_stream = byte_stream
553
 
 
554
 
    def iter_stream_decoder(self):
555
 
        """Iterate the contents of the pack from stream_decoder."""
556
 
        # dequeue pending items
557
 
        for record in self.stream_decoder.read_pending_records():
558
 
            yield record
559
 
        # Pull bytes of the wire, decode them to records, yield those records.
560
 
        for bytes in self.byte_stream:
561
 
            self.stream_decoder.accept_bytes(bytes)
562
 
            for record in self.stream_decoder.read_pending_records():
563
 
                yield record
564
 
 
565
 
    def iter_substream_bytes(self):
566
 
        if self.first_bytes is not None:
567
 
            yield self.first_bytes
568
 
            # If we run out of pack records, single the outer layer to stop.
569
 
            self.first_bytes = None
570
 
        for record in self.iter_pack_records:
571
 
            record_names, record_bytes = record
572
 
            record_name, = record_names
573
 
            substream_type = record_name[0]
574
 
            if substream_type != self.current_type:
575
 
                # end of a substream, seed the next substream.
576
 
                self.current_type = substream_type
577
 
                self.first_bytes = record_bytes
578
 
                return
579
 
            yield record_bytes
580
 
 
581
 
    def record_stream(self):
582
 
        """Yield substream_type, substream from the byte stream."""
583
 
        self.seed_state()
584
 
        # Make and consume sub generators, one per substream type:
585
 
        while self.first_bytes is not None:
586
 
            substream = NetworkRecordStream(self.iter_substream_bytes())
587
 
            # after substream is fully consumed, self.current_type is set to
588
 
            # the next type, and self.first_bytes is set to the matching bytes.
589
 
            yield self.current_type, substream.read()
590
 
 
591
 
    def seed_state(self):
592
 
        """Prepare the _ByteStreamDecoder to decode from the pack stream."""
593
 
        # Set a single generator we can use to get data from the pack stream.
594
 
        self.iter_pack_records = self.iter_stream_decoder()
595
 
        # Seed the very first subiterator with content; after this each one
596
 
        # seeds the next.
597
 
        list(self.iter_substream_bytes())
598
 
 
599
 
 
600
 
def _byte_stream_to_stream(byte_stream):
601
 
    """Convert a byte stream into a format and a stream.
602
 
 
603
 
    :param byte_stream: A bytes iterator, as output by _stream_to_byte_stream.
604
 
    :return: (RepositoryFormat, stream_generator)
605
 
    """
606
 
    decoder = _ByteStreamDecoder(byte_stream)
607
 
    for bytes in byte_stream:
608
 
        decoder.stream_decoder.accept_bytes(bytes)
609
 
        for record in decoder.stream_decoder.read_pending_records(max=1):
610
 
            record_names, src_format_name = record
611
 
            src_format = network_format_registry.get(src_format_name)
612
 
            return src_format, decoder.record_stream()
613
 
 
614
 
 
615
 
class SmartServerRepositoryUnlock(SmartServerRepositoryRequest):
616
 
 
617
 
    def do_repository_request(self, repository, token):
618
 
        try:
619
 
            repository.lock_write(token=token)
620
 
        except errors.TokenMismatch, e:
621
 
            return FailedSmartServerResponse(('TokenMismatch',))
622
 
        repository.dont_leave_lock_in_place()
623
 
        repository.unlock()
624
 
        return SuccessfulSmartServerResponse(('ok',))
625
 
 
626
 
 
627
 
class SmartServerRepositorySetMakeWorkingTrees(SmartServerRepositoryRequest):
628
 
 
629
 
    def do_repository_request(self, repository, str_bool_new_value):
630
 
        if str_bool_new_value == 'True':
631
 
            new_value = True
632
 
        else:
633
 
            new_value = False
634
 
        repository.set_make_working_trees(new_value)
635
 
        return SuccessfulSmartServerResponse(('ok',))
636
 
 
637
 
 
638
 
class SmartServerRepositoryTarball(SmartServerRepositoryRequest):
639
 
    """Get the raw repository files as a tarball.
640
 
 
641
 
    The returned tarball contains a .bzr control directory which in turn
642
 
    contains a repository.
643
 
 
644
 
    This takes one parameter, compression, which currently must be
645
 
    "", "gz", or "bz2".
646
 
 
647
 
    This is used to implement the Repository.copy_content_into operation.
648
 
    """
649
 
 
650
 
    def do_repository_request(self, repository, compression):
651
 
        tmp_dirname, tmp_repo = self._copy_to_tempdir(repository)
652
 
        try:
653
 
            controldir_name = tmp_dirname + '/.bzr'
654
 
            return self._tarfile_response(controldir_name, compression)
655
 
        finally:
656
 
            osutils.rmtree(tmp_dirname)
657
 
 
658
 
    def _copy_to_tempdir(self, from_repo):
659
 
        tmp_dirname = osutils.mkdtemp(prefix='tmpbzrclone')
660
 
        tmp_bzrdir = from_repo.bzrdir._format.initialize(tmp_dirname)
661
 
        tmp_repo = from_repo._format.initialize(tmp_bzrdir)
662
 
        from_repo.copy_content_into(tmp_repo)
663
 
        return tmp_dirname, tmp_repo
664
 
 
665
 
    def _tarfile_response(self, tmp_dirname, compression):
666
 
        temp = tempfile.NamedTemporaryFile()
667
 
        try:
668
 
            self._tarball_of_dir(tmp_dirname, compression, temp.file)
669
 
            # all finished; write the tempfile out to the network
670
 
            temp.seek(0)
671
 
            return SuccessfulSmartServerResponse(('ok',), temp.read())
672
 
            # FIXME: Don't read the whole thing into memory here; rather stream
673
 
            # it out from the file onto the network. mbp 20070411
674
 
        finally:
675
 
            temp.close()
676
 
 
677
 
    def _tarball_of_dir(self, dirname, compression, ofile):
678
 
        filename = os.path.basename(ofile.name)
679
 
        tarball = tarfile.open(fileobj=ofile, name=filename,
680
 
            mode='w|' + compression)
681
 
        try:
682
 
            # The tarball module only accepts ascii names, and (i guess)
683
 
            # packs them with their 8bit names.  We know all the files
684
 
            # within the repository have ASCII names so the should be safe
685
 
            # to pack in.
686
 
            dirname = dirname.encode(sys.getfilesystemencoding())
687
 
            # python's tarball module includes the whole path by default so
688
 
            # override it
689
 
            if not dirname.endswith('.bzr'):
690
 
                raise ValueError(dirname)
691
 
            tarball.add(dirname, '.bzr') # recursive by default
692
 
        finally:
693
 
            tarball.close()
694
 
 
695
 
 
696
 
class SmartServerRepositoryInsertStreamLocked(SmartServerRepositoryRequest):
697
 
    """Insert a record stream from a RemoteSink into a repository.
698
 
 
699
 
    This gets bytes pushed to it by the network infrastructure and turns that
700
 
    into a bytes iterator using a thread. That is then processed by
701
 
    _byte_stream_to_stream.
702
 
 
703
 
    New in 1.14.
704
 
    """
705
 
 
706
 
    def do_repository_request(self, repository, resume_tokens, lock_token):
707
 
        """StreamSink.insert_stream for a remote repository."""
708
 
        repository.lock_write(token=lock_token)
709
 
        self.do_insert_stream_request(repository, resume_tokens)
710
 
 
711
 
    def do_insert_stream_request(self, repository, resume_tokens):
712
 
        tokens = [token for token in resume_tokens.split(' ') if token]
713
 
        self.tokens = tokens
714
 
        self.repository = repository
715
 
        self.queue = Queue.Queue()
716
 
        self.insert_thread = threading.Thread(target=self._inserter_thread)
717
 
        self.insert_thread.start()
718
 
 
719
 
    def do_chunk(self, body_stream_chunk):
720
 
        self.queue.put(body_stream_chunk)
721
 
 
722
 
    def _inserter_thread(self):
723
 
        try:
724
 
            src_format, stream = _byte_stream_to_stream(
725
 
                self.blocking_byte_stream())
726
 
            self.insert_result = self.repository._get_sink().insert_stream(
727
 
                stream, src_format, self.tokens)
728
 
            self.insert_ok = True
729
 
        except:
730
 
            self.insert_exception = sys.exc_info()
731
 
            self.insert_ok = False
732
 
 
733
 
    def blocking_byte_stream(self):
734
 
        while True:
735
 
            bytes = self.queue.get()
736
 
            if bytes is StopIteration:
737
 
                return
738
 
            else:
739
 
                yield bytes
740
 
 
741
 
    def do_end(self):
742
 
        self.queue.put(StopIteration)
743
 
        if self.insert_thread is not None:
744
 
            self.insert_thread.join()
745
 
        if not self.insert_ok:
746
 
            exc_info = self.insert_exception
747
 
            raise exc_info[0], exc_info[1], exc_info[2]
748
 
        write_group_tokens, missing_keys = self.insert_result
749
 
        if write_group_tokens or missing_keys:
750
 
            # bzip needed? missing keys should typically be a small set.
751
 
            # Should this be a streaming body response ?
752
 
            missing_keys = sorted(missing_keys)
753
 
            bytes = bencode.bencode((write_group_tokens, missing_keys))
754
 
            self.repository.unlock()
755
 
            return SuccessfulSmartServerResponse(('missing-basis', bytes))
756
 
        else:
757
 
            self.repository.unlock()
758
 
            return SuccessfulSmartServerResponse(('ok', ))
759
 
 
760
 
 
761
 
class SmartServerRepositoryInsertStream_1_19(SmartServerRepositoryInsertStreamLocked):
762
 
    """Insert a record stream from a RemoteSink into a repository.
763
 
 
764
 
    Same as SmartServerRepositoryInsertStreamLocked, except:
765
 
     - the lock token argument is optional
766
 
     - servers that implement this verb accept 'inventory-delta' records in the
767
 
       stream.
768
 
 
769
 
    New in 1.19.
770
 
    """
771
 
 
772
 
    def do_repository_request(self, repository, resume_tokens, lock_token=None):
773
 
        """StreamSink.insert_stream for a remote repository."""
774
 
        SmartServerRepositoryInsertStreamLocked.do_repository_request(
775
 
            self, repository, resume_tokens, lock_token)
776
 
 
777
 
 
778
 
class SmartServerRepositoryInsertStream(SmartServerRepositoryInsertStreamLocked):
779
 
    """Insert a record stream from a RemoteSink into an unlocked repository.
780
 
 
781
 
    This is the same as SmartServerRepositoryInsertStreamLocked, except it
782
 
    takes no lock_tokens; i.e. it works with an unlocked (or lock-free, e.g.
783
 
    like pack format) repository.
784
 
 
785
 
    New in 1.13.
786
 
    """
787
 
 
788
 
    def do_repository_request(self, repository, resume_tokens):
789
 
        """StreamSink.insert_stream for a remote repository."""
790
 
        repository.lock_write()
791
 
        self.do_insert_stream_request(repository, resume_tokens)
792
 
 
793