~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

Turn completion assertions into separate methods.

Many common assertions used to be expressed as arguments to the complete
method.  This makes the checks more explicit, and the code easier to read.

Show diffs side-by-side

added added

removed removed

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