/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to bzrlib/smart/repository.py

  • Committer: Robert Collins
  • Date: 2010-05-05 00:05:29 UTC
  • mto: This revision was merged to the branch mainline in revision 5206.
  • Revision ID: robertc@robertcollins.net-20100505000529-ltmllyms5watqj5u
Make 'pydoc bzrlib.tests.build_tree_shape' useful.

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/smart/repository.py
14
14
# along with this program; if not, write to the Free Software
15
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
16
16
 
17
 
"""Server-side repository related request implementations."""
18
 
 
19
 
from __future__ import absolute_import
 
17
"""Server-side repository related request implmentations."""
20
18
 
21
19
import bz2
22
 
import itertools
23
20
import os
24
 
try:
25
 
    import queue
26
 
except ImportError:
27
 
    import Queue as queue
 
21
import Queue
28
22
import sys
29
23
import tempfile
30
24
import threading
31
 
import zlib
32
25
 
33
 
from ... import (
 
26
from bzrlib import (
34
27
    bencode,
35
28
    errors,
36
 
    estimate_compressed_size,
 
29
    graph,
37
30
    osutils,
38
 
    trace,
 
31
    pack,
39
32
    ui,
40
 
    )
41
 
from .. import (
42
 
    inventory as _mod_inventory,
43
 
    inventory_delta,
44
 
    pack,
45
 
    vf_search,
46
 
    )
47
 
from ..bzrdir import BzrDir
48
 
from ...sixish import (
49
 
    reraise,
50
 
)
51
 
from .request import (
 
33
    versionedfile,
 
34
    )
 
35
from bzrlib.bzrdir import BzrDir
 
36
from bzrlib.smart.request import (
52
37
    FailedSmartServerResponse,
53
38
    SmartServerRequest,
54
39
    SuccessfulSmartServerResponse,
55
40
    )
56
 
from ...repository import (
57
 
    _strip_NULL_ghosts,
58
 
    network_format_registry,
59
 
    )
60
 
from ... import revision as _mod_revision
61
 
from ..versionedfile import (
62
 
    ChunkedContentFactory,
 
41
from bzrlib.repository import _strip_NULL_ghosts, network_format_registry
 
42
from bzrlib import revision as _mod_revision
 
43
from bzrlib.versionedfile import (
63
44
    NetworkRecordStream,
64
45
    record_to_fulltext_bytes,
65
46
    )
101
82
            recreate_search trusts that clients will look for missing things
102
83
            they expected and get it from elsewhere.
103
84
        """
104
 
        if search_bytes == b'everything':
105
 
            return vf_search.EverythingResult(repository), None
106
 
        lines = search_bytes.split(b'\n')
107
 
        if lines[0] == b'ancestry-of':
 
85
        lines = search_bytes.split('\n')
 
86
        if lines[0] == 'ancestry-of':
108
87
            heads = lines[1:]
109
 
            search_result = vf_search.PendingAncestryResult(heads, repository)
 
88
            search_result = graph.PendingAncestryResult(heads, repository)
110
89
            return search_result, None
111
 
        elif lines[0] == b'search':
 
90
        elif lines[0] == 'search':
112
91
            return self.recreate_search_from_recipe(repository, lines[1:],
113
 
                                                    discard_excess=discard_excess)
 
92
                discard_excess=discard_excess)
114
93
        else:
115
 
            return (None, FailedSmartServerResponse((b'BadSearch',)))
 
94
            return (None, FailedSmartServerResponse(('BadSearch',)))
116
95
 
117
96
    def recreate_search_from_recipe(self, repository, lines,
118
 
                                    discard_excess=False):
 
97
        discard_excess=False):
119
98
        """Recreate a specific revision search (vs a from-tip search).
120
99
 
121
100
        :param discard_excess: If True, and the search refers to data we don't
123
102
            recreate_search trusts that clients will look for missing things
124
103
            they expected and get it from elsewhere.
125
104
        """
126
 
        start_keys = set(lines[0].split(b' '))
127
 
        exclude_keys = set(lines[1].split(b' '))
128
 
        revision_count = int(lines[2].decode('ascii'))
129
 
        with repository.lock_read():
 
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:
130
110
            search = repository.get_graph()._make_breadth_first_searcher(
131
111
                start_keys)
132
112
            while True:
133
113
                try:
134
 
                    next_revs = next(search)
 
114
                    next_revs = search.next()
135
115
                except StopIteration:
136
116
                    break
137
117
                search.stop_searching_any(exclude_keys.intersection(next_revs))
138
 
            (started_keys, excludes, included_keys) = search.get_state()
139
 
            if (not discard_excess and len(included_keys) != revision_count):
 
118
            search_result = search.get_result()
 
119
            if (not discard_excess and
 
120
                search_result.get_recipe()[3] != revision_count):
140
121
                # we got back a different amount of data than expected, this
141
122
                # gets reported as NoSuchRevision, because less revisions
142
123
                # indicates missing revisions, and more should never happen as
143
124
                # the excludes list considers ghosts and ensures that ghost
144
125
                # filling races are not a problem.
145
 
                return (None, FailedSmartServerResponse((b'NoSuchRevision',)))
146
 
            search_result = vf_search.SearchResult(started_keys, excludes,
147
 
                                                   len(included_keys), included_keys)
 
126
                return (None, FailedSmartServerResponse(('NoSuchRevision',)))
148
127
            return (search_result, None)
 
128
        finally:
 
129
            repository.unlock()
149
130
 
150
131
 
151
132
class SmartServerRepositoryReadLocked(SmartServerRepositoryRequest):
153
134
 
154
135
    def do_repository_request(self, repository, *args):
155
136
        """Read lock a repository for do_readlocked_repository_request."""
156
 
        with repository.lock_read():
 
137
        repository.lock_read()
 
138
        try:
157
139
            return self.do_readlocked_repository_request(repository, *args)
158
 
 
159
 
 
160
 
class SmartServerRepositoryBreakLock(SmartServerRepositoryRequest):
161
 
    """Break a repository lock."""
162
 
 
163
 
    def do_repository_request(self, repository):
164
 
        repository.break_lock()
165
 
        return SuccessfulSmartServerResponse((b'ok', ))
166
 
 
167
 
 
168
 
_lsprof_count = 0
 
140
        finally:
 
141
            repository.unlock()
169
142
 
170
143
 
171
144
class SmartServerRepositoryGetParentMap(SmartServerRepositoryRequest):
189
162
        :param revision_ids: The utf8 encoded revision_id to answer for.
190
163
        """
191
164
        self._revision_ids = revision_ids
192
 
        return None  # Signal that we want a body.
 
165
        return None # Signal that we want a body.
193
166
 
194
167
    def do_body(self, body_bytes):
195
168
        """Process the current search state and perform the parent lookup.
200
173
            compressed.
201
174
        """
202
175
        repository = self._repository
203
 
        with repository.lock_read():
 
176
        repository.lock_read()
 
177
        try:
204
178
            return self._do_repository_request(body_bytes)
 
179
        finally:
 
180
            repository.unlock()
205
181
 
206
 
    def _expand_requested_revs(self, repo_graph, revision_ids, client_seen_revs,
207
 
                               include_missing, max_size=65536):
 
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()
208
200
        result = {}
209
201
        queried_revs = set()
210
 
        estimator = estimate_compressed_size.ZLibEstimator(max_size)
 
202
        size_so_far = 0
211
203
        next_revs = revision_ids
212
204
        first_loop_done = False
213
205
        while next_revs:
227
219
                    encoded_id = revision_id
228
220
                else:
229
221
                    missing_rev = True
230
 
                    encoded_id = b"missing:" + revision_id
 
222
                    encoded_id = "missing:" + revision_id
231
223
                    parents = []
232
 
                if (revision_id not in client_seen_revs
233
 
                        and (not missing_rev or include_missing)):
 
224
                if (revision_id not in client_seen_revs and
 
225
                    (not missing_rev or include_missing)):
234
226
                    # Client does not have this revision, give it to it.
235
227
                    # add parents to the result
236
228
                    result[encoded_id] = parents
237
229
                    # Approximate the serialized cost of this revision_id.
238
 
                    line = encoded_id + b' ' + b' '.join(parents) + b'\n'
239
 
                    estimator.add_content(line)
 
230
                    size_so_far += 2 + len(encoded_id) + sum(map(len, parents))
240
231
            # get all the directly asked for parents, and then flesh out to
241
232
            # 64K (compressed) or so. We do one level of depth at a time to
242
233
            # stay in sync with the client. The 250000 magic number is
243
234
            # estimated compression ratio taken from bzr.dev itself.
244
 
            if self.no_extra_results or (first_loop_done and estimator.full()):
245
 
                trace.mutter('size: %d, z_size: %d'
246
 
                             % (estimator._uncompressed_size_added,
247
 
                                estimator._compressed_size_added))
 
235
            if self.no_extra_results or (
 
236
                first_loop_done and size_so_far > 250000):
248
237
                next_revs = set()
249
238
                break
250
239
            # don't query things we've already queried
251
 
            next_revs = next_revs.difference(queried_revs)
 
240
            next_revs.difference_update(queried_revs)
252
241
            first_loop_done = True
253
 
        return result
254
 
 
255
 
    def _do_repository_request(self, body_bytes):
256
 
        repository = self._repository
257
 
        revision_ids = set(self._revision_ids)
258
 
        include_missing = b'include-missing:' in revision_ids
259
 
        if include_missing:
260
 
            revision_ids.remove(b'include-missing:')
261
 
        body_lines = body_bytes.split(b'\n')
262
 
        search_result, error = self.recreate_search_from_recipe(
263
 
            repository, body_lines)
264
 
        if error is not None:
265
 
            return error
266
 
        # TODO might be nice to start up the search again; but thats not
267
 
        # written or tested yet.
268
 
        client_seen_revs = set(search_result.get_keys())
269
 
        # Always include the requested ids.
270
 
        client_seen_revs.difference_update(revision_ids)
271
 
 
272
 
        repo_graph = repository.get_graph()
273
 
        result = self._expand_requested_revs(repo_graph, revision_ids,
274
 
                                             client_seen_revs, include_missing)
275
242
 
276
243
        # sorting trivially puts lexographically similar revision ids together.
277
244
        # Compression FTW.
278
 
        lines = []
279
245
        for revision, parents in sorted(result.items()):
280
 
            lines.append(b' '.join((revision, ) + tuple(parents)))
 
246
            lines.append(' '.join((revision, ) + tuple(parents)))
281
247
 
282
248
        return SuccessfulSmartServerResponse(
283
 
            (b'ok', ), bz2.compress(b'\n'.join(lines)))
 
249
            ('ok', ), bz2.compress('\n'.join(lines)))
284
250
 
285
251
 
286
252
class SmartServerRepositoryGetRevisionGraph(SmartServerRepositoryReadLocked):
305
271
        else:
306
272
            search_ids = repository.all_revision_ids()
307
273
        search = graph._make_breadth_first_searcher(search_ids)
308
 
        transitive_ids = set(itertools.chain.from_iterable(search))
 
274
        transitive_ids = set()
 
275
        map(transitive_ids.update, list(search))
309
276
        parent_map = graph.get_parent_map(transitive_ids)
310
277
        revision_graph = _strip_NULL_ghosts(parent_map)
311
278
        if revision_id and revision_id not in revision_graph:
312
279
            # Note that we return an empty body, rather than omitting the body.
313
280
            # This way the client knows that it can always expect to find a body
314
281
            # in the response for this method, even in the error case.
315
 
            return FailedSmartServerResponse((b'nosuchrevision', revision_id), b'')
 
282
            return FailedSmartServerResponse(('nosuchrevision', revision_id), '')
316
283
 
317
284
        for revision, parents in revision_graph.items():
318
 
            lines.append(b' '.join((revision, ) + tuple(parents)))
 
285
            lines.append(' '.join((revision, ) + tuple(parents)))
319
286
 
320
 
        return SuccessfulSmartServerResponse((b'ok', ), b'\n'.join(lines))
 
287
        return SuccessfulSmartServerResponse(('ok', ), '\n'.join(lines))
321
288
 
322
289
 
323
290
class SmartServerRepositoryGetRevIdForRevno(SmartServerRepositoryReadLocked):
324
291
 
325
292
    def do_readlocked_repository_request(self, repository, revno,
326
 
                                         known_pair):
 
293
            known_pair):
327
294
        """Find the revid for a given revno, given a known revno/revid pair.
328
 
 
 
295
        
329
296
        New in 1.17.
330
297
        """
331
298
        try:
332
 
            found_flag, result = repository.get_rev_id_for_revno(
333
 
                revno, known_pair)
334
 
        except errors.NoSuchRevision as err:
335
 
            if err.revision != known_pair[1]:
 
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]:
336
302
                raise AssertionError(
337
303
                    'get_rev_id_for_revno raised RevisionNotPresent for '
338
 
                    'non-initial revision: ' + err.revision)
339
 
            return FailedSmartServerResponse(
340
 
                (b'nosuchrevision', err.revision))
341
 
        except errors.RevnoOutOfBounds as e:
342
 
            return FailedSmartServerResponse(
343
 
                (b'revno-outofbounds', e.revno, e.minimum, e.maximum))
 
304
                    'non-initial revision: ' + err.revision_id)
 
305
            return FailedSmartServerResponse(
 
306
                ('nosuchrevision', err.revision_id))
344
307
        if found_flag:
345
 
            return SuccessfulSmartServerResponse((b'ok', result))
 
308
            return SuccessfulSmartServerResponse(('ok', result))
346
309
        else:
347
310
            earliest_revno, earliest_revid = result
348
311
            return SuccessfulSmartServerResponse(
349
 
                (b'history-incomplete', earliest_revno, earliest_revid))
350
 
 
351
 
 
352
 
class SmartServerRepositoryGetSerializerFormat(SmartServerRepositoryRequest):
353
 
 
354
 
    def do_repository_request(self, repository):
355
 
        """Return the serializer format for this repository.
356
 
 
357
 
        New in 2.5.0.
358
 
 
359
 
        :param repository: The repository to query
360
 
        :return: A smart server response (b'ok', FORMAT)
361
 
        """
362
 
        serializer = repository.get_serializer_format()
363
 
        return SuccessfulSmartServerResponse((b'ok', serializer))
 
312
                ('history-incomplete', earliest_revno, earliest_revid))
364
313
 
365
314
 
366
315
class SmartServerRequestHasRevision(SmartServerRepositoryRequest):
370
319
 
371
320
        :param repository: The repository to query in.
372
321
        :param revision_id: The utf8 encoded revision_id to lookup.
373
 
        :return: A smart server response of ('yes', ) if the revision is
374
 
            present. ('no', ) if it is missing.
 
322
        :return: A smart server response of ('ok', ) if the revision is
 
323
            present.
375
324
        """
376
325
        if repository.has_revision(revision_id):
377
 
            return SuccessfulSmartServerResponse((b'yes', ))
 
326
            return SuccessfulSmartServerResponse(('yes', ))
378
327
        else:
379
 
            return SuccessfulSmartServerResponse((b'no', ))
380
 
 
381
 
 
382
 
class SmartServerRequestHasSignatureForRevisionId(
383
 
        SmartServerRepositoryRequest):
384
 
 
385
 
    def do_repository_request(self, repository, revision_id):
386
 
        """Return ok if a signature is present for a revision.
387
 
 
388
 
        Introduced in bzr 2.5.0.
389
 
 
390
 
        :param repository: The repository to query in.
391
 
        :param revision_id: The utf8 encoded revision_id to lookup.
392
 
        :return: A smart server response of ('yes', ) if a
393
 
            signature for the revision is present,
394
 
            ('no', ) if it is missing.
395
 
        """
396
 
        try:
397
 
            if repository.has_signature_for_revision_id(revision_id):
398
 
                return SuccessfulSmartServerResponse((b'yes', ))
399
 
            else:
400
 
                return SuccessfulSmartServerResponse((b'no', ))
401
 
        except errors.NoSuchRevision:
402
 
            return FailedSmartServerResponse(
403
 
                (b'nosuchrevision', revision_id))
 
328
            return SuccessfulSmartServerResponse(('no', ))
404
329
 
405
330
 
406
331
class SmartServerRepositoryGatherStats(SmartServerRepositoryRequest):
412
337
        :param revid: utf8 encoded rev id or an empty string to indicate None
413
338
        :param committers: 'yes' or 'no'.
414
339
 
415
 
        :return: A SmartServerResponse (b'ok',), a encoded body looking like
 
340
        :return: A SmartServerResponse ('ok',), a encoded body looking like
416
341
              committers: 1
417
342
              firstrev: 1234.230 0
418
343
              latestrev: 345.700 3600
420
345
 
421
346
              But containing only fields returned by the gather_stats() call
422
347
        """
423
 
        if revid == b'':
 
348
        if revid == '':
424
349
            decoded_revision_id = None
425
350
        else:
426
351
            decoded_revision_id = revid
427
 
        if committers == b'yes':
 
352
        if committers == 'yes':
428
353
            decoded_committers = True
429
354
        else:
430
355
            decoded_committers = None
431
 
        try:
432
 
            stats = repository.gather_stats(decoded_revision_id,
433
 
                                            decoded_committers)
434
 
        except errors.NoSuchRevision:
435
 
            return FailedSmartServerResponse((b'nosuchrevision', revid))
436
 
 
437
 
        body = b''
438
 
        if 'committers' in stats:
439
 
            body += b'committers: %d\n' % stats['committers']
440
 
        if 'firstrev' in stats:
441
 
            body += b'firstrev: %.3f %d\n' % stats['firstrev']
442
 
        if 'latestrev' in stats:
443
 
            body += b'latestrev: %.3f %d\n' % stats['latestrev']
444
 
        if 'revisions' in stats:
445
 
            body += b'revisions: %d\n' % stats['revisions']
446
 
        if 'size' in stats:
447
 
            body += b'size: %d\n' % stats['size']
448
 
 
449
 
        return SuccessfulSmartServerResponse((b'ok', ), body)
450
 
 
451
 
 
452
 
class SmartServerRepositoryGetRevisionSignatureText(
453
 
        SmartServerRepositoryRequest):
454
 
    """Return the signature text of a revision.
455
 
 
456
 
    New in 2.5.
457
 
    """
458
 
 
459
 
    def do_repository_request(self, repository, revision_id):
460
 
        """Return the result of repository.get_signature_text().
461
 
 
462
 
        :param repository: The repository to query in.
463
 
        :return: A smart server response of with the signature text as
464
 
            body.
465
 
        """
466
 
        try:
467
 
            text = repository.get_signature_text(revision_id)
468
 
        except errors.NoSuchRevision as err:
469
 
            return FailedSmartServerResponse(
470
 
                (b'nosuchrevision', err.revision))
471
 
        return SuccessfulSmartServerResponse((b'ok', ), text)
 
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)
472
371
 
473
372
 
474
373
class SmartServerRepositoryIsShared(SmartServerRepositoryRequest):
481
380
            shared, and ('no', ) if it is not.
482
381
        """
483
382
        if repository.is_shared():
484
 
            return SuccessfulSmartServerResponse((b'yes', ))
485
 
        else:
486
 
            return SuccessfulSmartServerResponse((b'no', ))
487
 
 
488
 
 
489
 
class SmartServerRepositoryMakeWorkingTrees(SmartServerRepositoryRequest):
490
 
 
491
 
    def do_repository_request(self, repository):
492
 
        """Return the result of repository.make_working_trees().
493
 
 
494
 
        Introduced in bzr 2.5.0.
495
 
 
496
 
        :param repository: The repository to query in.
497
 
        :return: A smart server response of ('yes', ) if the repository uses
498
 
            working trees, and ('no', ) if it is not.
499
 
        """
500
 
        if repository.make_working_trees():
501
 
            return SuccessfulSmartServerResponse((b'yes', ))
502
 
        else:
503
 
            return SuccessfulSmartServerResponse((b'no', ))
 
383
            return SuccessfulSmartServerResponse(('yes', ))
 
384
        else:
 
385
            return SuccessfulSmartServerResponse(('no', ))
504
386
 
505
387
 
506
388
class SmartServerRepositoryLockWrite(SmartServerRepositoryRequest):
507
389
 
508
 
    def do_repository_request(self, repository, token=b''):
 
390
    def do_repository_request(self, repository, token=''):
509
391
        # XXX: this probably should not have a token.
510
 
        if token == b'':
 
392
        if token == '':
511
393
            token = None
512
394
        try:
513
 
            token = repository.lock_write(token=token).repository_token
514
 
        except errors.LockContention as e:
515
 
            return FailedSmartServerResponse((b'LockContention',))
 
395
            token = repository.lock_write(token=token)
 
396
        except errors.LockContention, e:
 
397
            return FailedSmartServerResponse(('LockContention',))
516
398
        except errors.UnlockableTransport:
517
 
            return FailedSmartServerResponse((b'UnlockableTransport',))
518
 
        except errors.LockFailed as e:
519
 
            return FailedSmartServerResponse((b'LockFailed',
520
 
                                              str(e.lock), str(e.why)))
 
399
            return FailedSmartServerResponse(('UnlockableTransport',))
 
400
        except errors.LockFailed, e:
 
401
            return FailedSmartServerResponse(('LockFailed',
 
402
                str(e.lock), str(e.why)))
521
403
        if token is not None:
522
404
            repository.leave_lock_in_place()
523
405
        repository.unlock()
524
406
        if token is None:
525
 
            token = b''
526
 
        return SuccessfulSmartServerResponse((b'ok', token))
 
407
            token = ''
 
408
        return SuccessfulSmartServerResponse(('ok', token))
527
409
 
528
410
 
529
411
class SmartServerRepositoryGetStream(SmartServerRepositoryRequest):
531
413
    def do_repository_request(self, repository, to_network_name):
532
414
        """Get a stream for inserting into a to_format repository.
533
415
 
534
 
        The request body is 'search_bytes', a description of the revisions
535
 
        being requested.
536
 
 
537
 
        In 2.3 this verb added support for search_bytes == 'everything'.  Older
538
 
        implementations will respond with a BadSearch error, and clients should
539
 
        catch this and fallback appropriately.
540
 
 
541
416
        :param repository: The repository to stream from.
542
417
        :param to_network_name: The network name of the format of the target
543
418
            repository.
545
420
        self._to_format = network_format_registry.get(to_network_name)
546
421
        if self._should_fake_unknown():
547
422
            return FailedSmartServerResponse(
548
 
                (b'UnknownMethod', b'Repository.get_stream'))
549
 
        return None  # Signal that we want a body.
 
423
                ('UnknownMethod', 'Repository.get_stream'))
 
424
        return None # Signal that we want a body.
550
425
 
551
426
    def _should_fake_unknown(self):
552
427
        """Return True if we should return UnknownMethod to the client.
553
 
 
 
428
        
554
429
        This is a workaround for bugs in pre-1.19 clients that claim to
555
430
        support receiving streams of CHK repositories.  The pre-1.19 client
556
431
        expects inventory records to be serialized in the format defined by
569
444
        if not from_format.supports_chks:
570
445
            # Source not CHK: that's ok
571
446
            return False
572
 
        if (to_format.supports_chks
573
 
            and from_format.repository_class is to_format.repository_class
574
 
                and from_format._serializer == to_format._serializer):
 
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):
575
450
            # Source is CHK, but target matches: that's ok
576
451
            # (e.g. 2a->2a, or CHK2->2a)
577
452
            return False
584
459
        repository.lock_read()
585
460
        try:
586
461
            search_result, error = self.recreate_search(repository, body_bytes,
587
 
                                                        discard_excess=True)
 
462
                discard_excess=True)
588
463
            if error is not None:
589
464
                repository.unlock()
590
465
                return error
591
466
            source = repository._get_source(self._to_format)
592
467
            stream = source.get_stream(search_result)
593
468
        except Exception:
 
469
            exc_info = sys.exc_info()
594
470
            try:
595
471
                # On non-error, unlocking is done by the body stream handler.
596
472
                repository.unlock()
597
473
            finally:
598
 
                raise
599
 
        return SuccessfulSmartServerResponse((b'ok',),
600
 
                                             body_stream=self.body_stream(stream, repository))
 
474
                raise exc_info[0], exc_info[1], exc_info[2]
 
475
        return SuccessfulSmartServerResponse(('ok',),
 
476
            body_stream=self.body_stream(stream, repository))
601
477
 
602
478
    def body_stream(self, stream, repository):
603
479
        byte_stream = _stream_to_byte_stream(stream, repository._format)
604
480
        try:
605
481
            for bytes in byte_stream:
606
482
                yield bytes
607
 
        except errors.RevisionNotPresent as e:
 
483
        except errors.RevisionNotPresent, e:
608
484
            # This shouldn't be able to happen, but as we don't buffer
609
485
            # everything it can in theory happen.
610
486
            repository.unlock()
611
 
            yield FailedSmartServerResponse((b'NoSuchRevision', e.revision_id))
 
487
            yield FailedSmartServerResponse(('NoSuchRevision', e.revision_id))
612
488
        else:
613
489
            repository.unlock()
614
490
 
615
491
 
616
492
class SmartServerRepositoryGetStream_1_19(SmartServerRepositoryGetStream):
617
 
    """The same as Repository.get_stream, but will return stream CHK formats to
618
 
    clients.
619
 
 
620
 
    See SmartServerRepositoryGetStream._should_fake_unknown.
621
 
 
622
 
    New in 1.19.
623
 
    """
624
493
 
625
494
    def _should_fake_unknown(self):
626
495
        """Returns False; we don't need to workaround bugs in 1.19+ clients."""
631
500
    """Convert a record stream to a self delimited byte stream."""
632
501
    pack_writer = pack.ContainerSerialiser()
633
502
    yield pack_writer.begin()
634
 
    yield pack_writer.bytes_record(src_format.network_name(), b'')
 
503
    yield pack_writer.bytes_record(src_format.network_name(), '')
635
504
    for substream_type, substream in stream:
636
505
        for record in substream:
637
506
            if record.storage_kind in ('chunked', 'fulltext'):
638
507
                serialised = record_to_fulltext_bytes(record)
 
508
            elif record.storage_kind == 'inventory-delta':
 
509
                serialised = record_to_inventory_delta_bytes(record)
639
510
            elif record.storage_kind == 'absent':
640
511
                raise ValueError("Absent factory for %s" % (record.key,))
641
512
            else:
644
515
                # Some streams embed the whole stream into the wire
645
516
                # representation of the first record, which means that
646
517
                # later records have no wire representation: we skip them.
647
 
                yield pack_writer.bytes_record(serialised, [(substream_type.encode('ascii'),)])
 
518
                yield pack_writer.bytes_record(serialised, [(substream_type,)])
648
519
    yield pack_writer.end()
649
520
 
650
521
 
673
544
    :ivar first_bytes: The first bytes to give the next NetworkRecordStream.
674
545
    """
675
546
 
676
 
    def __init__(self, byte_stream, record_counter):
 
547
    def __init__(self, byte_stream):
677
548
        """Create a _ByteStreamDecoder."""
678
549
        self.stream_decoder = pack.ContainerPushParser()
679
550
        self.current_type = None
680
551
        self.first_bytes = None
681
552
        self.byte_stream = byte_stream
682
 
        self._record_counter = record_counter
683
 
        self.key_count = 0
684
553
 
685
554
    def iter_stream_decoder(self):
686
555
        """Iterate the contents of the pack from stream_decoder."""
711
580
 
712
581
    def record_stream(self):
713
582
        """Yield substream_type, substream from the byte stream."""
714
 
        def wrap_and_count(pb, rc, substream):
715
 
            """Yield records from stream while showing progress."""
716
 
            counter = 0
717
 
            if rc:
718
 
                if self.current_type != 'revisions' and self.key_count != 0:
719
 
                    # As we know the number of revisions now (in self.key_count)
720
 
                    # we can setup and use record_counter (rc).
721
 
                    if not rc.is_initialized():
722
 
                        rc.setup(self.key_count, self.key_count)
723
 
            for record in substream.read():
724
 
                if rc:
725
 
                    if rc.is_initialized() and counter == rc.STEP:
726
 
                        rc.increment(counter)
727
 
                        pb.update('Estimate', rc.current, rc.max)
728
 
                        counter = 0
729
 
                    if self.current_type == 'revisions':
730
 
                        # Total records is proportional to number of revs
731
 
                        # to fetch. With remote, we used self.key_count to
732
 
                        # track the number of revs. Once we have the revs
733
 
                        # counts in self.key_count, the progress bar changes
734
 
                        # from 'Estimating..' to 'Estimate' above.
735
 
                        self.key_count += 1
736
 
                        if counter == rc.STEP:
737
 
                            pb.update('Estimating..', self.key_count)
738
 
                            counter = 0
739
 
                counter += 1
740
 
                yield record
741
 
 
742
583
        self.seed_state()
743
 
        with ui.ui_factory.nested_progress_bar() as pb:
744
 
            rc = self._record_counter
745
 
            try:
746
 
                # Make and consume sub generators, one per substream type:
747
 
                while self.first_bytes is not None:
748
 
                    substream = NetworkRecordStream(
749
 
                        self.iter_substream_bytes())
750
 
                    # after substream is fully consumed, self.current_type is set
751
 
                    # to the next type, and self.first_bytes is set to the matching
752
 
                    # bytes.
753
 
                    yield self.current_type.decode('ascii'), wrap_and_count(pb, rc, substream)
754
 
            finally:
755
 
                if rc:
756
 
                    pb.update('Done', rc.max, rc.max)
 
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()
757
590
 
758
591
    def seed_state(self):
759
592
        """Prepare the _ByteStreamDecoder to decode from the pack stream."""
764
597
        list(self.iter_substream_bytes())
765
598
 
766
599
 
767
 
def _byte_stream_to_stream(byte_stream, record_counter=None):
 
600
def _byte_stream_to_stream(byte_stream):
768
601
    """Convert a byte stream into a format and a stream.
769
602
 
770
603
    :param byte_stream: A bytes iterator, as output by _stream_to_byte_stream.
771
604
    :return: (RepositoryFormat, stream_generator)
772
605
    """
773
 
    decoder = _ByteStreamDecoder(byte_stream, record_counter)
 
606
    decoder = _ByteStreamDecoder(byte_stream)
774
607
    for bytes in byte_stream:
775
608
        decoder.stream_decoder.accept_bytes(bytes)
776
609
        for record in decoder.stream_decoder.read_pending_records(max=1):
784
617
    def do_repository_request(self, repository, token):
785
618
        try:
786
619
            repository.lock_write(token=token)
787
 
        except errors.TokenMismatch as e:
788
 
            return FailedSmartServerResponse((b'TokenMismatch',))
 
620
        except errors.TokenMismatch, e:
 
621
            return FailedSmartServerResponse(('TokenMismatch',))
789
622
        repository.dont_leave_lock_in_place()
790
623
        repository.unlock()
791
 
        return SuccessfulSmartServerResponse((b'ok',))
792
 
 
793
 
 
794
 
class SmartServerRepositoryGetPhysicalLockStatus(SmartServerRepositoryRequest):
795
 
    """Get the physical lock status for a repository.
796
 
 
797
 
    New in 2.5.
798
 
    """
799
 
 
800
 
    def do_repository_request(self, repository):
801
 
        if repository.get_physical_lock_status():
802
 
            return SuccessfulSmartServerResponse((b'yes', ))
803
 
        else:
804
 
            return SuccessfulSmartServerResponse((b'no', ))
 
624
        return SuccessfulSmartServerResponse(('ok',))
805
625
 
806
626
 
807
627
class SmartServerRepositorySetMakeWorkingTrees(SmartServerRepositoryRequest):
808
628
 
809
629
    def do_repository_request(self, repository, str_bool_new_value):
810
 
        if str_bool_new_value == b'True':
 
630
        if str_bool_new_value == 'True':
811
631
            new_value = True
812
632
        else:
813
633
            new_value = False
814
634
        repository.set_make_working_trees(new_value)
815
 
        return SuccessfulSmartServerResponse((b'ok',))
 
635
        return SuccessfulSmartServerResponse(('ok',))
816
636
 
817
637
 
818
638
class SmartServerRepositoryTarball(SmartServerRepositoryRequest):
837
657
 
838
658
    def _copy_to_tempdir(self, from_repo):
839
659
        tmp_dirname = osutils.mkdtemp(prefix='tmpbzrclone')
840
 
        tmp_bzrdir = from_repo.controldir._format.initialize(tmp_dirname)
 
660
        tmp_bzrdir = from_repo.bzrdir._format.initialize(tmp_dirname)
841
661
        tmp_repo = from_repo._format.initialize(tmp_bzrdir)
842
662
        from_repo.copy_content_into(tmp_repo)
843
663
        return tmp_dirname, tmp_repo
844
664
 
845
665
    def _tarfile_response(self, tmp_dirname, compression):
846
 
        with tempfile.NamedTemporaryFile() as temp:
 
666
        temp = tempfile.NamedTemporaryFile()
 
667
        try:
847
668
            self._tarball_of_dir(tmp_dirname, compression, temp.file)
848
669
            # all finished; write the tempfile out to the network
849
670
            temp.seek(0)
850
 
            return SuccessfulSmartServerResponse((b'ok',), temp.read())
 
671
            return SuccessfulSmartServerResponse(('ok',), temp.read())
851
672
            # FIXME: Don't read the whole thing into memory here; rather stream
852
673
            # it out from the file onto the network. mbp 20070411
 
674
        finally:
 
675
            temp.close()
853
676
 
854
677
    def _tarball_of_dir(self, dirname, compression, ofile):
855
678
        import tarfile
856
679
        filename = os.path.basename(ofile.name)
857
 
        with tarfile.open(fileobj=ofile, name=filename,
858
 
                          mode='w|' + compression) as tarball:
 
680
        tarball = tarfile.open(fileobj=ofile, name=filename,
 
681
            mode='w|' + compression)
 
682
        try:
859
683
            # The tarball module only accepts ascii names, and (i guess)
860
684
            # packs them with their 8bit names.  We know all the files
861
685
            # within the repository have ASCII names so the should be safe
865
689
            # override it
866
690
            if not dirname.endswith('.bzr'):
867
691
                raise ValueError(dirname)
868
 
            tarball.add(dirname, '.bzr')  # recursive by default
 
692
            tarball.add(dirname, '.bzr') # recursive by default
 
693
        finally:
 
694
            tarball.close()
869
695
 
870
696
 
871
697
class SmartServerRepositoryInsertStreamLocked(SmartServerRepositoryRequest):
884
710
        self.do_insert_stream_request(repository, resume_tokens)
885
711
 
886
712
    def do_insert_stream_request(self, repository, resume_tokens):
887
 
        tokens = [token.decode('utf-8')
888
 
                  for token in resume_tokens.split(b' ') if token]
 
713
        tokens = [token for token in resume_tokens.split(' ') if token]
889
714
        self.tokens = tokens
890
715
        self.repository = repository
891
 
        self.queue = queue.Queue()
 
716
        self.queue = Queue.Queue()
892
717
        self.insert_thread = threading.Thread(target=self._inserter_thread)
893
718
        self.insert_thread.start()
894
719
 
919
744
        if self.insert_thread is not None:
920
745
            self.insert_thread.join()
921
746
        if not self.insert_ok:
922
 
            try:
923
 
                reraise(*self.insert_exception)
924
 
            finally:
925
 
                del self.insert_exception
 
747
            exc_info = self.insert_exception
 
748
            raise exc_info[0], exc_info[1], exc_info[2]
926
749
        write_group_tokens, missing_keys = self.insert_result
927
750
        if write_group_tokens or missing_keys:
928
751
            # bzip needed? missing keys should typically be a small set.
929
752
            # Should this be a streaming body response ?
930
 
            missing_keys = sorted(
931
 
                [(entry[0].encode('utf-8'),) + entry[1:] for entry in missing_keys])
932
 
            bytes = bencode.bencode((
933
 
                [token.encode('utf-8') for token in write_group_tokens], missing_keys))
 
753
            missing_keys = sorted(missing_keys)
 
754
            bytes = bencode.bencode((write_group_tokens, missing_keys))
934
755
            self.repository.unlock()
935
 
            return SuccessfulSmartServerResponse((b'missing-basis', bytes))
 
756
            return SuccessfulSmartServerResponse(('missing-basis', bytes))
936
757
        else:
937
758
            self.repository.unlock()
938
 
            return SuccessfulSmartServerResponse((b'ok', ))
 
759
            return SuccessfulSmartServerResponse(('ok', ))
939
760
 
940
761
 
941
762
class SmartServerRepositoryInsertStream_1_19(SmartServerRepositoryInsertStreamLocked):
971
792
        self.do_insert_stream_request(repository, resume_tokens)
972
793
 
973
794
 
974
 
class SmartServerRepositoryAddSignatureText(SmartServerRepositoryRequest):
975
 
    """Add a revision signature text.
976
 
 
977
 
    New in 2.5.
978
 
    """
979
 
 
980
 
    def do_repository_request(self, repository, lock_token, revision_id,
981
 
                              *write_group_tokens):
982
 
        """Add a revision signature text.
983
 
 
984
 
        :param repository: Repository to operate on
985
 
        :param lock_token: Lock token
986
 
        :param revision_id: Revision for which to add signature
987
 
        :param write_group_tokens: Write group tokens
988
 
        """
989
 
        self._lock_token = lock_token
990
 
        self._revision_id = revision_id
991
 
        self._write_group_tokens = [token.decode(
992
 
            'utf-8') for token in write_group_tokens]
993
 
        return None
994
 
 
995
 
    def do_body(self, body_bytes):
996
 
        """Add a signature text.
997
 
 
998
 
        :param body_bytes: GPG signature text
999
 
        :return: SuccessfulSmartServerResponse with arguments 'ok' and
1000
 
            the list of new write group tokens.
1001
 
        """
1002
 
        with self._repository.lock_write(token=self._lock_token):
1003
 
            self._repository.resume_write_group(self._write_group_tokens)
1004
 
            try:
1005
 
                self._repository.add_signature_text(self._revision_id,
1006
 
                                                    body_bytes)
1007
 
            finally:
1008
 
                new_write_group_tokens = self._repository.suspend_write_group()
1009
 
        return SuccessfulSmartServerResponse(
1010
 
            (b'ok', ) + tuple([token.encode('utf-8') for token in new_write_group_tokens]))
1011
 
 
1012
 
 
1013
 
class SmartServerRepositoryStartWriteGroup(SmartServerRepositoryRequest):
1014
 
    """Start a write group.
1015
 
 
1016
 
    New in 2.5.
1017
 
    """
1018
 
 
1019
 
    def do_repository_request(self, repository, lock_token):
1020
 
        """Start a write group."""
1021
 
        with repository.lock_write(token=lock_token):
1022
 
            repository.start_write_group()
1023
 
            try:
1024
 
                tokens = repository.suspend_write_group()
1025
 
            except errors.UnsuspendableWriteGroup:
1026
 
                return FailedSmartServerResponse((b'UnsuspendableWriteGroup',))
1027
 
        return SuccessfulSmartServerResponse((b'ok', tokens))
1028
 
 
1029
 
 
1030
 
class SmartServerRepositoryCommitWriteGroup(SmartServerRepositoryRequest):
1031
 
    """Commit a write group.
1032
 
 
1033
 
    New in 2.5.
1034
 
    """
1035
 
 
1036
 
    def do_repository_request(self, repository, lock_token,
1037
 
                              write_group_tokens):
1038
 
        """Commit a write group."""
1039
 
        with repository.lock_write(token=lock_token):
1040
 
            try:
1041
 
                repository.resume_write_group(
1042
 
                    [token.decode('utf-8') for token in write_group_tokens])
1043
 
            except errors.UnresumableWriteGroup as e:
1044
 
                return FailedSmartServerResponse(
1045
 
                    (b'UnresumableWriteGroup', [token.encode('utf-8') for token
1046
 
                                                in e.write_groups], e.reason.encode('utf-8')))
1047
 
            try:
1048
 
                repository.commit_write_group()
1049
 
            except:
1050
 
                write_group_tokens = repository.suspend_write_group()
1051
 
                # FIXME JRV 2011-11-19: What if the write_group_tokens
1052
 
                # have changed?
1053
 
                raise
1054
 
        return SuccessfulSmartServerResponse((b'ok', ))
1055
 
 
1056
 
 
1057
 
class SmartServerRepositoryAbortWriteGroup(SmartServerRepositoryRequest):
1058
 
    """Abort a write group.
1059
 
 
1060
 
    New in 2.5.
1061
 
    """
1062
 
 
1063
 
    def do_repository_request(self, repository, lock_token, write_group_tokens):
1064
 
        """Abort a write group."""
1065
 
        with repository.lock_write(token=lock_token):
1066
 
            try:
1067
 
                repository.resume_write_group(
1068
 
                    [token.decode('utf-8') for token in write_group_tokens])
1069
 
            except errors.UnresumableWriteGroup as e:
1070
 
                return FailedSmartServerResponse(
1071
 
                    (b'UnresumableWriteGroup',
1072
 
                        [token.encode('utf-8') for token in e.write_groups],
1073
 
                        e.reason.encode('utf-8')))
1074
 
                repository.abort_write_group()
1075
 
        return SuccessfulSmartServerResponse((b'ok', ))
1076
 
 
1077
 
 
1078
 
class SmartServerRepositoryCheckWriteGroup(SmartServerRepositoryRequest):
1079
 
    """Check that a write group is still valid.
1080
 
 
1081
 
    New in 2.5.
1082
 
    """
1083
 
 
1084
 
    def do_repository_request(self, repository, lock_token, write_group_tokens):
1085
 
        """Abort a write group."""
1086
 
        with repository.lock_write(token=lock_token):
1087
 
            try:
1088
 
                repository.resume_write_group(
1089
 
                    [token.decode('utf-8') for token in write_group_tokens])
1090
 
            except errors.UnresumableWriteGroup as e:
1091
 
                return FailedSmartServerResponse(
1092
 
                    (b'UnresumableWriteGroup',
1093
 
                        [token.encode('utf-8') for token in e.write_groups],
1094
 
                        e.reason.encode('utf-8')))
1095
 
            else:
1096
 
                repository.suspend_write_group()
1097
 
        return SuccessfulSmartServerResponse((b'ok', ))
1098
 
 
1099
 
 
1100
 
class SmartServerRepositoryAllRevisionIds(SmartServerRepositoryRequest):
1101
 
    """Retrieve all of the revision ids in a repository.
1102
 
 
1103
 
    New in 2.5.
1104
 
    """
1105
 
 
1106
 
    def do_repository_request(self, repository):
1107
 
        revids = repository.all_revision_ids()
1108
 
        return SuccessfulSmartServerResponse((b"ok", ), b"\n".join(revids))
1109
 
 
1110
 
 
1111
 
class SmartServerRepositoryReconcile(SmartServerRepositoryRequest):
1112
 
    """Reconcile a repository.
1113
 
 
1114
 
    New in 2.5.
1115
 
    """
1116
 
 
1117
 
    def do_repository_request(self, repository, lock_token):
1118
 
        try:
1119
 
            repository.lock_write(token=lock_token)
1120
 
        except errors.TokenLockingNotSupported as e:
1121
 
            return FailedSmartServerResponse(
1122
 
                (b'TokenLockingNotSupported', ))
1123
 
        try:
1124
 
            reconciler = repository.reconcile()
1125
 
        finally:
1126
 
            repository.unlock()
1127
 
        body = [
1128
 
            b"garbage_inventories: %d\n" % reconciler.garbage_inventories,
1129
 
            b"inconsistent_parents: %d\n" % reconciler.inconsistent_parents,
1130
 
            ]
1131
 
        return SuccessfulSmartServerResponse((b'ok', ), b"".join(body))
1132
 
 
1133
 
 
1134
 
class SmartServerRepositoryPack(SmartServerRepositoryRequest):
1135
 
    """Pack a repository.
1136
 
 
1137
 
    New in 2.5.
1138
 
    """
1139
 
 
1140
 
    def do_repository_request(self, repository, lock_token, clean_obsolete_packs):
1141
 
        self._repository = repository
1142
 
        self._lock_token = lock_token
1143
 
        if clean_obsolete_packs == b'True':
1144
 
            self._clean_obsolete_packs = True
1145
 
        else:
1146
 
            self._clean_obsolete_packs = False
1147
 
        return None
1148
 
 
1149
 
    def do_body(self, body_bytes):
1150
 
        if body_bytes == "":
1151
 
            hint = None
1152
 
        else:
1153
 
            hint = body_bytes.splitlines()
1154
 
        with self._repository.lock_write(token=self._lock_token):
1155
 
            self._repository.pack(hint, self._clean_obsolete_packs)
1156
 
        return SuccessfulSmartServerResponse((b"ok", ), )
1157
 
 
1158
 
 
1159
 
class SmartServerRepositoryIterFilesBytes(SmartServerRepositoryRequest):
1160
 
    """Iterate over the contents of files.
1161
 
 
1162
 
    The client sends a list of desired files to stream, one
1163
 
    per line, and as tuples of file id and revision, separated by
1164
 
    \0.
1165
 
 
1166
 
    The server replies with a stream. Each entry is preceded by a header,
1167
 
    which can either be:
1168
 
 
1169
 
    * "ok\x00IDX\n" where IDX is the index of the entry in the desired files
1170
 
        list sent by the client. This header is followed by the contents of
1171
 
        the file, bzip2-compressed.
1172
 
    * "absent\x00FILEID\x00REVISION\x00IDX" to indicate a text is missing.
1173
 
        The client can then raise an appropriate RevisionNotPresent error
1174
 
        or check its fallback repositories.
1175
 
 
1176
 
    New in 2.5.
1177
 
    """
1178
 
 
1179
 
    def body_stream(self, repository, desired_files):
1180
 
        with self._repository.lock_read():
1181
 
            text_keys = {}
1182
 
            for i, key in enumerate(desired_files):
1183
 
                text_keys[key] = i
1184
 
            for record in repository.texts.get_record_stream(text_keys,
1185
 
                                                             'unordered', True):
1186
 
                identifier = text_keys[record.key]
1187
 
                if record.storage_kind == 'absent':
1188
 
                    yield b"absent\0%s\0%s\0%d\n" % (record.key[0],
1189
 
                                                     record.key[1], identifier)
1190
 
                    # FIXME: Way to abort early?
1191
 
                    continue
1192
 
                yield b"ok\0%d\n" % identifier
1193
 
                compressor = zlib.compressobj()
1194
 
                for bytes in record.iter_bytes_as('chunked'):
1195
 
                    data = compressor.compress(bytes)
1196
 
                    if data:
1197
 
                        yield data
1198
 
                data = compressor.flush()
1199
 
                if data:
1200
 
                    yield data
1201
 
 
1202
 
    def do_body(self, body_bytes):
1203
 
        desired_files = [
1204
 
            tuple(l.split(b"\0")) for l in body_bytes.splitlines()]
1205
 
        return SuccessfulSmartServerResponse((b'ok', ),
1206
 
                                             body_stream=self.body_stream(self._repository, desired_files))
1207
 
 
1208
 
    def do_repository_request(self, repository):
1209
 
        # Signal that we want a body
1210
 
        return None
1211
 
 
1212
 
 
1213
 
class SmartServerRepositoryIterRevisions(SmartServerRepositoryRequest):
1214
 
    """Stream a list of revisions.
1215
 
 
1216
 
    The client sends a list of newline-separated revision ids in the
1217
 
    body of the request and the server replies with the serializer format,
1218
 
    and a stream of bzip2-compressed revision texts (using the specified
1219
 
    serializer format).
1220
 
 
1221
 
    Any revisions the server does not have are omitted from the stream.
1222
 
 
1223
 
    New in 2.5.
1224
 
    """
1225
 
 
1226
 
    def do_repository_request(self, repository):
1227
 
        self._repository = repository
1228
 
        # Signal there is a body
1229
 
        return None
1230
 
 
1231
 
    def do_body(self, body_bytes):
1232
 
        revision_ids = body_bytes.split(b"\n")
1233
 
        return SuccessfulSmartServerResponse(
1234
 
            (b'ok', self._repository.get_serializer_format()),
1235
 
            body_stream=self.body_stream(self._repository, revision_ids))
1236
 
 
1237
 
    def body_stream(self, repository, revision_ids):
1238
 
        with self._repository.lock_read():
1239
 
            for record in repository.revisions.get_record_stream(
1240
 
                    [(revid,) for revid in revision_ids], 'unordered', True):
1241
 
                if record.storage_kind == 'absent':
1242
 
                    continue
1243
 
                yield zlib.compress(record.get_bytes_as('fulltext'))
1244
 
 
1245
 
 
1246
 
class SmartServerRepositoryGetInventories(SmartServerRepositoryRequest):
1247
 
    """Get the inventory deltas for a set of revision ids.
1248
 
 
1249
 
    This accepts a list of revision ids, and then sends a chain
1250
 
    of deltas for the inventories of those revisions. The first
1251
 
    revision will be empty.
1252
 
 
1253
 
    The server writes back zlibbed serialized inventory deltas,
1254
 
    in the ordering specified. The base for each delta is the
1255
 
    inventory generated by the previous delta.
1256
 
 
1257
 
    New in 2.5.
1258
 
    """
1259
 
 
1260
 
    def _inventory_delta_stream(self, repository, ordering, revids):
1261
 
        prev_inv = _mod_inventory.Inventory(root_id=None,
1262
 
                                            revision_id=_mod_revision.NULL_REVISION)
1263
 
        serializer = inventory_delta.InventoryDeltaSerializer(
1264
 
            repository.supports_rich_root(),
1265
 
            repository._format.supports_tree_reference)
1266
 
        with repository.lock_read():
1267
 
            for inv, revid in repository._iter_inventories(revids, ordering):
1268
 
                if inv is None:
1269
 
                    continue
1270
 
                inv_delta = inv._make_delta(prev_inv)
1271
 
                lines = serializer.delta_to_lines(
1272
 
                    prev_inv.revision_id, inv.revision_id, inv_delta)
1273
 
                yield ChunkedContentFactory(
1274
 
                    inv.revision_id, None, None, lines,
1275
 
                    chunks_are_lines=True)
1276
 
                prev_inv = inv
1277
 
 
1278
 
    def body_stream(self, repository, ordering, revids):
1279
 
        substream = self._inventory_delta_stream(repository,
1280
 
                                                 ordering, revids)
1281
 
        return _stream_to_byte_stream([('inventory-deltas', substream)],
1282
 
                                      repository._format)
1283
 
 
1284
 
    def do_body(self, body_bytes):
1285
 
        return SuccessfulSmartServerResponse((b'ok', ),
1286
 
                                             body_stream=self.body_stream(self._repository, self._ordering,
1287
 
                                                                          body_bytes.splitlines()))
1288
 
 
1289
 
    def do_repository_request(self, repository, ordering):
1290
 
        ordering = ordering.decode('ascii')
1291
 
        if ordering == 'unordered':
1292
 
            # inventory deltas for a topologically sorted stream
1293
 
            # are likely to be smaller
1294
 
            ordering = 'topological'
1295
 
        self._ordering = ordering
1296
 
        # Signal that we want a body
1297
 
        return None
1298
 
 
1299
 
 
1300
 
class SmartServerRepositoryGetStreamForMissingKeys(SmartServerRepositoryRequest):
1301
 
 
1302
 
    def do_repository_request(self, repository, to_network_name):
1303
 
        """Get a stream for missing keys.
1304
 
 
1305
 
        :param repository: The repository to stream from.
1306
 
        :param to_network_name: The network name of the format of the target
1307
 
            repository.
1308
 
        """
1309
 
        try:
1310
 
            self._to_format = network_format_registry.get(to_network_name)
1311
 
        except KeyError:
1312
 
            return FailedSmartServerResponse(
1313
 
                (b'UnknownFormat', b'repository', to_network_name))
1314
 
        return None  # Signal that we want a body.
1315
 
 
1316
 
    def do_body(self, body_bytes):
1317
 
        repository = self._repository
1318
 
        repository.lock_read()
1319
 
        try:
1320
 
            source = repository._get_source(self._to_format)
1321
 
            keys = []
1322
 
            for entry in body_bytes.split(b'\n'):
1323
 
                (kind, revid) = entry.split(b'\t')
1324
 
                keys.append((kind.decode('utf-8'), revid))
1325
 
            stream = source.get_stream_for_missing_keys(keys)
1326
 
        except Exception:
1327
 
            try:
1328
 
                # On non-error, unlocking is done by the body stream handler.
1329
 
                repository.unlock()
1330
 
            finally:
1331
 
                raise
1332
 
        return SuccessfulSmartServerResponse((b'ok',),
1333
 
                                             body_stream=self.body_stream(stream, repository))
1334
 
 
1335
 
    def body_stream(self, stream, repository):
1336
 
        byte_stream = _stream_to_byte_stream(stream, repository._format)
1337
 
        try:
1338
 
            for bytes in byte_stream:
1339
 
                yield bytes
1340
 
        except errors.RevisionNotPresent as e:
1341
 
            # This shouldn't be able to happen, but as we don't buffer
1342
 
            # everything it can in theory happen.
1343
 
            repository.unlock()
1344
 
            yield FailedSmartServerResponse((b'NoSuchRevision', e.revision_id))
1345
 
        else:
1346
 
            repository.unlock()
1347
 
 
1348
 
 
1349
 
class SmartServerRepositoryRevisionArchive(SmartServerRepositoryRequest):
1350
 
 
1351
 
    def do_repository_request(self, repository, revision_id, format, name,
1352
 
                              root, subdir=None, force_mtime=None):
1353
 
        """Stream an archive file for a specific revision.
1354
 
        :param repository: The repository to stream from.
1355
 
        :param revision_id: Revision for which to export the tree
1356
 
        :param format: Format (tar, tgz, tbz2, etc)
1357
 
        :param name: Target file name
1358
 
        :param root: Name of root directory (or '')
1359
 
        :param subdir: Subdirectory to export, if not the root
1360
 
        """
1361
 
        tree = repository.revision_tree(revision_id)
1362
 
        if subdir is not None:
1363
 
            subdir = subdir.decode('utf-8')
1364
 
        if root is not None:
1365
 
            root = root.decode('utf-8')
1366
 
        name = name.decode('utf-8')
1367
 
        return SuccessfulSmartServerResponse((b'ok',),
1368
 
                                             body_stream=self.body_stream(
1369
 
            tree, format.decode(
1370
 
                'utf-8'), os.path.basename(name), root, subdir,
1371
 
            force_mtime))
1372
 
 
1373
 
    def body_stream(self, tree, format, name, root, subdir=None, force_mtime=None):
1374
 
        with tree.lock_read():
1375
 
            return tree.archive(format, name, root, subdir, force_mtime)
1376
 
 
1377
 
 
1378
 
class SmartServerRepositoryAnnotateFileRevision(SmartServerRepositoryRequest):
1379
 
 
1380
 
    def do_repository_request(self, repository, revision_id, tree_path,
1381
 
                              file_id=None, default_revision=None):
1382
 
        """Stream an archive file for a specific revision.
1383
 
 
1384
 
        :param repository: The repository to stream from.
1385
 
        :param revision_id: Revision for which to export the tree
1386
 
        :param tree_path: The path inside the tree
1387
 
        :param file_id: Optional file_id for the file
1388
 
        :param default_revision: Default revision
1389
 
        """
1390
 
        tree = repository.revision_tree(revision_id)
1391
 
        with tree.lock_read():
1392
 
            body = bencode.bencode(list(tree.annotate_iter(
1393
 
                tree_path.decode('utf-8'), default_revision)))
1394
 
            return SuccessfulSmartServerResponse((b'ok',), body=body)