~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

  • Committer: John Arbash Meinel
  • Date: 2010-01-13 16:23:07 UTC
  • mto: (4634.119.7 2.0)
  • mto: This revision was merged to the branch mainline in revision 4959.
  • Revision ID: john@arbash-meinel.com-20100113162307-0bs82td16gzih827
Update the MANIFEST.in file.

Show diffs side-by-side

added added

removed removed

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