~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

  • Committer: Robert Collins
  • Date: 2006-03-28 14:29:13 UTC
  • mto: (1626.2.1 integration)
  • mto: This revision was merged to the branch mainline in revision 1628.
  • Revision ID: robertc@robertcollins.net-20060328142913-ac5afb37075719c6
Convert log to use the new tsort.merge_sort routine.

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2006, 2007, 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 tarfile
24
 
import tempfile
25
 
import threading
26
 
 
27
 
from bzrlib import (
28
 
    bencode,
29
 
    errors,
30
 
    graph,
31
 
    osutils,
32
 
    pack,
33
 
    ui,
34
 
    versionedfile,
35
 
    )
36
 
from bzrlib.bzrdir import BzrDir
37
 
from bzrlib.smart.request import (
38
 
    FailedSmartServerResponse,
39
 
    SmartServerRequest,
40
 
    SuccessfulSmartServerResponse,
41
 
    )
42
 
from bzrlib.repository import _strip_NULL_ghosts, network_format_registry
43
 
from bzrlib import revision as _mod_revision
44
 
from bzrlib.versionedfile import (
45
 
    NetworkRecordStream,
46
 
    record_to_fulltext_bytes,
47
 
    )
48
 
 
49
 
 
50
 
class SmartServerRepositoryRequest(SmartServerRequest):
51
 
    """Common base class for Repository requests."""
52
 
 
53
 
    def do(self, path, *args):
54
 
        """Execute a repository request.
55
 
 
56
 
        All Repository requests take a path to the repository as their first
57
 
        argument.  The repository must be at the exact path given by the
58
 
        client - no searching is done.
59
 
 
60
 
        The actual logic is delegated to self.do_repository_request.
61
 
 
62
 
        :param client_path: The path for the repository as received from the
63
 
            client.
64
 
        :return: A SmartServerResponse from self.do_repository_request().
65
 
        """
66
 
        transport = self.transport_from_client_path(path)
67
 
        bzrdir = BzrDir.open_from_transport(transport)
68
 
        # Save the repository for use with do_body.
69
 
        self._repository = bzrdir.open_repository()
70
 
        return self.do_repository_request(self._repository, *args)
71
 
 
72
 
    def do_repository_request(self, repository, *args):
73
 
        """Override to provide an implementation for a verb."""
74
 
        # No-op for verbs that take bodies (None as a result indicates a body
75
 
        # is expected)
76
 
        return None
77
 
 
78
 
    def recreate_search(self, repository, search_bytes, discard_excess=False):
79
 
        """Recreate a search from its serialised form.
80
 
 
81
 
        :param discard_excess: If True, and the search refers to data we don't
82
 
            have, just silently accept that fact - the verb calling
83
 
            recreate_search trusts that clients will look for missing things
84
 
            they expected and get it from elsewhere.
85
 
        """
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)
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
 
        :param repository: The repository to stream from.
418
 
        :param to_network_name: The network name of the format of the target
419
 
            repository.
420
 
        """
421
 
        self._to_format = network_format_registry.get(to_network_name)
422
 
        if self._should_fake_unknown():
423
 
            return FailedSmartServerResponse(
424
 
                ('UnknownMethod', 'Repository.get_stream'))
425
 
        return None # Signal that we want a body.
426
 
 
427
 
    def _should_fake_unknown(self):
428
 
        """Return True if we should return UnknownMethod to the client.
429
 
        
430
 
        This is a workaround for bugs in pre-1.19 clients that claim to
431
 
        support receiving streams of CHK repositories.  The pre-1.19 client
432
 
        expects inventory records to be serialized in the format defined by
433
 
        to_network_name, but in pre-1.19 (at least) that format definition
434
 
        tries to use the xml5 serializer, which does not correctly handle
435
 
        rich-roots.  After 1.19 the client can also accept inventory-deltas
436
 
        (which avoids this issue), and those clients will use the
437
 
        Repository.get_stream_1.19 verb instead of this one.
438
 
        So: if this repository is CHK, and the to_format doesn't match,
439
 
        we should just fake an UnknownSmartMethod error so that the client
440
 
        will fallback to VFS, rather than sending it a stream we know it
441
 
        cannot handle.
442
 
        """
443
 
        from_format = self._repository._format
444
 
        to_format = self._to_format
445
 
        if not from_format.supports_chks:
446
 
            # Source not CHK: that's ok
447
 
            return False
448
 
        if (to_format.supports_chks and
449
 
            from_format.repository_class is to_format.repository_class and
450
 
            from_format._serializer == to_format._serializer):
451
 
            # Source is CHK, but target matches: that's ok
452
 
            # (e.g. 2a->2a, or CHK2->2a)
453
 
            return False
454
 
        # Source is CHK, and target is not CHK or incompatible CHK.  We can't
455
 
        # generate a compatible stream.
456
 
        return True
457
 
 
458
 
    def do_body(self, body_bytes):
459
 
        repository = self._repository
460
 
        repository.lock_read()
461
 
        try:
462
 
            search_result, error = self.recreate_search(repository, body_bytes,
463
 
                discard_excess=True)
464
 
            if error is not None:
465
 
                repository.unlock()
466
 
                return error
467
 
            source = repository._get_source(self._to_format)
468
 
            stream = source.get_stream(search_result)
469
 
        except Exception:
470
 
            exc_info = sys.exc_info()
471
 
            try:
472
 
                # On non-error, unlocking is done by the body stream handler.
473
 
                repository.unlock()
474
 
            finally:
475
 
                raise exc_info[0], exc_info[1], exc_info[2]
476
 
        return SuccessfulSmartServerResponse(('ok',),
477
 
            body_stream=self.body_stream(stream, repository))
478
 
 
479
 
    def body_stream(self, stream, repository):
480
 
        byte_stream = _stream_to_byte_stream(stream, repository._format)
481
 
        try:
482
 
            for bytes in byte_stream:
483
 
                yield bytes
484
 
        except errors.RevisionNotPresent, e:
485
 
            # This shouldn't be able to happen, but as we don't buffer
486
 
            # everything it can in theory happen.
487
 
            repository.unlock()
488
 
            yield FailedSmartServerResponse(('NoSuchRevision', e.revision_id))
489
 
        else:
490
 
            repository.unlock()
491
 
 
492
 
 
493
 
class SmartServerRepositoryGetStream_1_19(SmartServerRepositoryGetStream):
494
 
 
495
 
    def _should_fake_unknown(self):
496
 
        """Returns False; we don't need to workaround bugs in 1.19+ clients."""
497
 
        return False
498
 
 
499
 
 
500
 
def _stream_to_byte_stream(stream, src_format):
501
 
    """Convert a record stream to a self delimited byte stream."""
502
 
    pack_writer = pack.ContainerSerialiser()
503
 
    yield pack_writer.begin()
504
 
    yield pack_writer.bytes_record(src_format.network_name(), '')
505
 
    for substream_type, substream in stream:
506
 
        if substream_type == 'inventory-deltas':
507
 
            # This doesn't feel like the ideal place to issue this warning;
508
 
            # however we don't want to do it in the Repository that's
509
 
            # generating the stream, because that might be on the server.
510
 
            # Instead we try to observe it as the stream goes by.
511
 
            ui.ui_factory.warn_cross_format_fetch(src_format,
512
 
                '(remote)')
513
 
        for record in substream:
514
 
            if record.storage_kind in ('chunked', 'fulltext'):
515
 
                serialised = record_to_fulltext_bytes(record)
516
 
            elif record.storage_kind == 'inventory-delta':
517
 
                serialised = record_to_inventory_delta_bytes(record)
518
 
            elif record.storage_kind == 'absent':
519
 
                raise ValueError("Absent factory for %s" % (record.key,))
520
 
            else:
521
 
                serialised = record.get_bytes_as(record.storage_kind)
522
 
            if serialised:
523
 
                # Some streams embed the whole stream into the wire
524
 
                # representation of the first record, which means that
525
 
                # later records have no wire representation: we skip them.
526
 
                yield pack_writer.bytes_record(serialised, [(substream_type,)])
527
 
    yield pack_writer.end()
528
 
 
529
 
 
530
 
class _ByteStreamDecoder(object):
531
 
    """Helper for _byte_stream_to_stream.
532
 
 
533
 
    The expected usage of this class is via the function _byte_stream_to_stream
534
 
    which creates a _ByteStreamDecoder, pops off the stream format and then
535
 
    yields the output of record_stream(), the main entry point to
536
 
    _ByteStreamDecoder.
537
 
 
538
 
    Broadly this class has to unwrap two layers of iterators:
539
 
    (type, substream)
540
 
    (substream details)
541
 
 
542
 
    This is complicated by wishing to return type, iterator_for_type, but
543
 
    getting the data for iterator_for_type when we find out type: we can't
544
 
    simply pass a generator down to the NetworkRecordStream parser, instead
545
 
    we have a little local state to seed each NetworkRecordStream instance,
546
 
    and gather the type that we'll be yielding.
547
 
 
548
 
    :ivar byte_stream: The byte stream being decoded.
549
 
    :ivar stream_decoder: A pack parser used to decode the bytestream
550
 
    :ivar current_type: The current type, used to join adjacent records of the
551
 
        same type into a single stream.
552
 
    :ivar first_bytes: The first bytes to give the next NetworkRecordStream.
553
 
    """
554
 
 
555
 
    def __init__(self, byte_stream):
556
 
        """Create a _ByteStreamDecoder."""
557
 
        self.stream_decoder = pack.ContainerPushParser()
558
 
        self.current_type = None
559
 
        self.first_bytes = None
560
 
        self.byte_stream = byte_stream
561
 
 
562
 
    def iter_stream_decoder(self):
563
 
        """Iterate the contents of the pack from stream_decoder."""
564
 
        # dequeue pending items
565
 
        for record in self.stream_decoder.read_pending_records():
566
 
            yield record
567
 
        # Pull bytes of the wire, decode them to records, yield those records.
568
 
        for bytes in self.byte_stream:
569
 
            self.stream_decoder.accept_bytes(bytes)
570
 
            for record in self.stream_decoder.read_pending_records():
571
 
                yield record
572
 
 
573
 
    def iter_substream_bytes(self):
574
 
        if self.first_bytes is not None:
575
 
            yield self.first_bytes
576
 
            # If we run out of pack records, single the outer layer to stop.
577
 
            self.first_bytes = None
578
 
        for record in self.iter_pack_records:
579
 
            record_names, record_bytes = record
580
 
            record_name, = record_names
581
 
            substream_type = record_name[0]
582
 
            if substream_type != self.current_type:
583
 
                # end of a substream, seed the next substream.
584
 
                self.current_type = substream_type
585
 
                self.first_bytes = record_bytes
586
 
                return
587
 
            yield record_bytes
588
 
 
589
 
    def record_stream(self):
590
 
        """Yield substream_type, substream from the byte stream."""
591
 
        self.seed_state()
592
 
        # Make and consume sub generators, one per substream type:
593
 
        while self.first_bytes is not None:
594
 
            substream = NetworkRecordStream(self.iter_substream_bytes())
595
 
            # after substream is fully consumed, self.current_type is set to
596
 
            # the next type, and self.first_bytes is set to the matching bytes.
597
 
            yield self.current_type, substream.read()
598
 
 
599
 
    def seed_state(self):
600
 
        """Prepare the _ByteStreamDecoder to decode from the pack stream."""
601
 
        # Set a single generator we can use to get data from the pack stream.
602
 
        self.iter_pack_records = self.iter_stream_decoder()
603
 
        # Seed the very first subiterator with content; after this each one
604
 
        # seeds the next.
605
 
        list(self.iter_substream_bytes())
606
 
 
607
 
 
608
 
def _byte_stream_to_stream(byte_stream):
609
 
    """Convert a byte stream into a format and a stream.
610
 
 
611
 
    :param byte_stream: A bytes iterator, as output by _stream_to_byte_stream.
612
 
    :return: (RepositoryFormat, stream_generator)
613
 
    """
614
 
    decoder = _ByteStreamDecoder(byte_stream)
615
 
    for bytes in byte_stream:
616
 
        decoder.stream_decoder.accept_bytes(bytes)
617
 
        for record in decoder.stream_decoder.read_pending_records(max=1):
618
 
            record_names, src_format_name = record
619
 
            src_format = network_format_registry.get(src_format_name)
620
 
            return src_format, decoder.record_stream()
621
 
 
622
 
 
623
 
class SmartServerRepositoryUnlock(SmartServerRepositoryRequest):
624
 
 
625
 
    def do_repository_request(self, repository, token):
626
 
        try:
627
 
            repository.lock_write(token=token)
628
 
        except errors.TokenMismatch, e:
629
 
            return FailedSmartServerResponse(('TokenMismatch',))
630
 
        repository.dont_leave_lock_in_place()
631
 
        repository.unlock()
632
 
        return SuccessfulSmartServerResponse(('ok',))
633
 
 
634
 
 
635
 
class SmartServerRepositorySetMakeWorkingTrees(SmartServerRepositoryRequest):
636
 
 
637
 
    def do_repository_request(self, repository, str_bool_new_value):
638
 
        if str_bool_new_value == 'True':
639
 
            new_value = True
640
 
        else:
641
 
            new_value = False
642
 
        repository.set_make_working_trees(new_value)
643
 
        return SuccessfulSmartServerResponse(('ok',))
644
 
 
645
 
 
646
 
class SmartServerRepositoryTarball(SmartServerRepositoryRequest):
647
 
    """Get the raw repository files as a tarball.
648
 
 
649
 
    The returned tarball contains a .bzr control directory which in turn
650
 
    contains a repository.
651
 
 
652
 
    This takes one parameter, compression, which currently must be
653
 
    "", "gz", or "bz2".
654
 
 
655
 
    This is used to implement the Repository.copy_content_into operation.
656
 
    """
657
 
 
658
 
    def do_repository_request(self, repository, compression):
659
 
        tmp_dirname, tmp_repo = self._copy_to_tempdir(repository)
660
 
        try:
661
 
            controldir_name = tmp_dirname + '/.bzr'
662
 
            return self._tarfile_response(controldir_name, compression)
663
 
        finally:
664
 
            osutils.rmtree(tmp_dirname)
665
 
 
666
 
    def _copy_to_tempdir(self, from_repo):
667
 
        tmp_dirname = osutils.mkdtemp(prefix='tmpbzrclone')
668
 
        tmp_bzrdir = from_repo.bzrdir._format.initialize(tmp_dirname)
669
 
        tmp_repo = from_repo._format.initialize(tmp_bzrdir)
670
 
        from_repo.copy_content_into(tmp_repo)
671
 
        return tmp_dirname, tmp_repo
672
 
 
673
 
    def _tarfile_response(self, tmp_dirname, compression):
674
 
        temp = tempfile.NamedTemporaryFile()
675
 
        try:
676
 
            self._tarball_of_dir(tmp_dirname, compression, temp.file)
677
 
            # all finished; write the tempfile out to the network
678
 
            temp.seek(0)
679
 
            return SuccessfulSmartServerResponse(('ok',), temp.read())
680
 
            # FIXME: Don't read the whole thing into memory here; rather stream
681
 
            # it out from the file onto the network. mbp 20070411
682
 
        finally:
683
 
            temp.close()
684
 
 
685
 
    def _tarball_of_dir(self, dirname, compression, ofile):
686
 
        filename = os.path.basename(ofile.name)
687
 
        tarball = tarfile.open(fileobj=ofile, name=filename,
688
 
            mode='w|' + compression)
689
 
        try:
690
 
            # The tarball module only accepts ascii names, and (i guess)
691
 
            # packs them with their 8bit names.  We know all the files
692
 
            # within the repository have ASCII names so the should be safe
693
 
            # to pack in.
694
 
            dirname = dirname.encode(sys.getfilesystemencoding())
695
 
            # python's tarball module includes the whole path by default so
696
 
            # override it
697
 
            if not dirname.endswith('.bzr'):
698
 
                raise ValueError(dirname)
699
 
            tarball.add(dirname, '.bzr') # recursive by default
700
 
        finally:
701
 
            tarball.close()
702
 
 
703
 
 
704
 
class SmartServerRepositoryInsertStreamLocked(SmartServerRepositoryRequest):
705
 
    """Insert a record stream from a RemoteSink into a repository.
706
 
 
707
 
    This gets bytes pushed to it by the network infrastructure and turns that
708
 
    into a bytes iterator using a thread. That is then processed by
709
 
    _byte_stream_to_stream.
710
 
 
711
 
    New in 1.14.
712
 
    """
713
 
 
714
 
    def do_repository_request(self, repository, resume_tokens, lock_token):
715
 
        """StreamSink.insert_stream for a remote repository."""
716
 
        repository.lock_write(token=lock_token)
717
 
        self.do_insert_stream_request(repository, resume_tokens)
718
 
 
719
 
    def do_insert_stream_request(self, repository, resume_tokens):
720
 
        tokens = [token for token in resume_tokens.split(' ') if token]
721
 
        self.tokens = tokens
722
 
        self.repository = repository
723
 
        self.queue = Queue.Queue()
724
 
        self.insert_thread = threading.Thread(target=self._inserter_thread)
725
 
        self.insert_thread.start()
726
 
 
727
 
    def do_chunk(self, body_stream_chunk):
728
 
        self.queue.put(body_stream_chunk)
729
 
 
730
 
    def _inserter_thread(self):
731
 
        try:
732
 
            src_format, stream = _byte_stream_to_stream(
733
 
                self.blocking_byte_stream())
734
 
            self.insert_result = self.repository._get_sink().insert_stream(
735
 
                stream, src_format, self.tokens)
736
 
            self.insert_ok = True
737
 
        except:
738
 
            self.insert_exception = sys.exc_info()
739
 
            self.insert_ok = False
740
 
 
741
 
    def blocking_byte_stream(self):
742
 
        while True:
743
 
            bytes = self.queue.get()
744
 
            if bytes is StopIteration:
745
 
                return
746
 
            else:
747
 
                yield bytes
748
 
 
749
 
    def do_end(self):
750
 
        self.queue.put(StopIteration)
751
 
        if self.insert_thread is not None:
752
 
            self.insert_thread.join()
753
 
        if not self.insert_ok:
754
 
            exc_info = self.insert_exception
755
 
            raise exc_info[0], exc_info[1], exc_info[2]
756
 
        write_group_tokens, missing_keys = self.insert_result
757
 
        if write_group_tokens or missing_keys:
758
 
            # bzip needed? missing keys should typically be a small set.
759
 
            # Should this be a streaming body response ?
760
 
            missing_keys = sorted(missing_keys)
761
 
            bytes = bencode.bencode((write_group_tokens, missing_keys))
762
 
            self.repository.unlock()
763
 
            return SuccessfulSmartServerResponse(('missing-basis', bytes))
764
 
        else:
765
 
            self.repository.unlock()
766
 
            return SuccessfulSmartServerResponse(('ok', ))
767
 
 
768
 
 
769
 
class SmartServerRepositoryInsertStream_1_19(SmartServerRepositoryInsertStreamLocked):
770
 
    """Insert a record stream from a RemoteSink into a repository.
771
 
 
772
 
    Same as SmartServerRepositoryInsertStreamLocked, except:
773
 
     - the lock token argument is optional
774
 
     - servers that implement this verb accept 'inventory-delta' records in the
775
 
       stream.
776
 
 
777
 
    New in 1.19.
778
 
    """
779
 
 
780
 
    def do_repository_request(self, repository, resume_tokens, lock_token=None):
781
 
        """StreamSink.insert_stream for a remote repository."""
782
 
        SmartServerRepositoryInsertStreamLocked.do_repository_request(
783
 
            self, repository, resume_tokens, lock_token)
784
 
 
785
 
 
786
 
class SmartServerRepositoryInsertStream(SmartServerRepositoryInsertStreamLocked):
787
 
    """Insert a record stream from a RemoteSink into an unlocked repository.
788
 
 
789
 
    This is the same as SmartServerRepositoryInsertStreamLocked, except it
790
 
    takes no lock_tokens; i.e. it works with an unlocked (or lock-free, e.g.
791
 
    like pack format) repository.
792
 
 
793
 
    New in 1.13.
794
 
    """
795
 
 
796
 
    def do_repository_request(self, repository, resume_tokens):
797
 
        """StreamSink.insert_stream for a remote repository."""
798
 
        repository.lock_write()
799
 
        self.do_insert_stream_request(repository, resume_tokens)
800
 
 
801