~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

  • Committer: Canonical.com Patch Queue Manager
  • Date: 2008-03-16 16:58:03 UTC
  • mfrom: (3224.3.1 news-typo)
  • Revision ID: pqm@pqm.ubuntu.com-20080316165803-tisoc9mpob9z544o
(Matt Nordhoff) Trivial NEWS typo fix

Show diffs side-by-side

added added

removed removed

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