wireproto.py
833 lines
| 27.7 KiB
| text/x-python
|
PythonLexer
/ mercurial / wireproto.py
|
r11581 | # wireproto.py - generic wire protocol support functions | ||
# | ||||
# Copyright 2005-2010 Matt Mackall <mpm@selenic.com> | ||||
# | ||||
# This software may be used and distributed according to the terms of the | ||||
# GNU General Public License version 2 or any later version. | ||||
|
r11879 | import urllib, tempfile, os, sys | ||
|
r11581 | from i18n import _ | ||
from node import bin, hex | ||||
|
r21069 | import changegroup as changegroupmod, bundle2 | ||
|
r20967 | import peer, error, encoding, util, store, exchange | ||
|
r11581 | |||
|
r20903 | |||
class abstractserverproto(object): | ||||
"""abstract class that summarizes the protocol API | ||||
Used as reference and documentation. | ||||
""" | ||||
def getargs(self, args): | ||||
"""return the value for arguments in <args> | ||||
returns a list of values (same order as <args>)""" | ||||
raise NotImplementedError() | ||||
def getfile(self, fp): | ||||
"""write the whole content of a file into a file like object | ||||
The file is in the form:: | ||||
(<chunk-size>\n<chunk>)+0\n | ||||
chunk size is the ascii version of the int. | ||||
""" | ||||
raise NotImplementedError() | ||||
def redirect(self): | ||||
"""may setup interception for stdout and stderr | ||||
See also the `restore` method.""" | ||||
raise NotImplementedError() | ||||
# If the `redirect` function does install interception, the `restore` | ||||
# function MUST be defined. If interception is not used, this function | ||||
# MUST NOT be defined. | ||||
# | ||||
# left commented here on purpose | ||||
# | ||||
#def restore(self): | ||||
# """reinstall previous stdout and stderr and return intercepted stdout | ||||
# """ | ||||
# raise NotImplementedError() | ||||
def groupchunks(self, cg): | ||||
"""return 4096 chunks from a changegroup object | ||||
Some protocols may have compressed the contents.""" | ||||
raise NotImplementedError() | ||||
|
r14621 | # abstract batching support | ||
class future(object): | ||||
'''placeholder for a value to be set later''' | ||||
def set(self, value): | ||||
|
r14970 | if util.safehasattr(self, 'value'): | ||
|
r14621 | raise error.RepoError("future is already set") | ||
self.value = value | ||||
class batcher(object): | ||||
'''base class for batches of commands submittable in a single request | ||||
|
r16683 | All methods invoked on instances of this class are simply queued and | ||
return a a future for the result. Once you call submit(), all the queued | ||||
calls are performed and the results set in their respective futures. | ||||
|
r14621 | ''' | ||
def __init__(self): | ||||
self.calls = [] | ||||
def __getattr__(self, name): | ||||
def call(*args, **opts): | ||||
resref = future() | ||||
self.calls.append((name, args, opts, resref,)) | ||||
return resref | ||||
return call | ||||
def submit(self): | ||||
pass | ||||
class localbatch(batcher): | ||||
'''performs the queued calls directly''' | ||||
def __init__(self, local): | ||||
batcher.__init__(self) | ||||
self.local = local | ||||
def submit(self): | ||||
for name, args, opts, resref in self.calls: | ||||
resref.set(getattr(self.local, name)(*args, **opts)) | ||||
class remotebatch(batcher): | ||||
'''batches the queued calls; uses as few roundtrips as possible''' | ||||
def __init__(self, remote): | ||||
|
r16683 | '''remote must support _submitbatch(encbatch) and | ||
_submitone(op, encargs)''' | ||||
|
r14621 | batcher.__init__(self) | ||
self.remote = remote | ||||
def submit(self): | ||||
req, rsp = [], [] | ||||
for name, args, opts, resref in self.calls: | ||||
mtd = getattr(self.remote, name) | ||||
|
r14970 | batchablefn = getattr(mtd, 'batchable', None) | ||
if batchablefn is not None: | ||||
batchable = batchablefn(mtd.im_self, *args, **opts) | ||||
|
r14621 | encargsorres, encresref = batchable.next() | ||
if encresref: | ||||
req.append((name, encargsorres,)) | ||||
rsp.append((batchable, encresref, resref,)) | ||||
else: | ||||
resref.set(encargsorres) | ||||
else: | ||||
if req: | ||||
self._submitreq(req, rsp) | ||||
req, rsp = [], [] | ||||
resref.set(mtd(*args, **opts)) | ||||
if req: | ||||
self._submitreq(req, rsp) | ||||
def _submitreq(self, req, rsp): | ||||
encresults = self.remote._submitbatch(req) | ||||
for encres, r in zip(encresults, rsp): | ||||
batchable, encresref, resref = r | ||||
encresref.set(encres) | ||||
resref.set(batchable.next()) | ||||
def batchable(f): | ||||
'''annotation for batchable methods | ||||
Such methods must implement a coroutine as follows: | ||||
@batchable | ||||
def sample(self, one, two=None): | ||||
# Handle locally computable results first: | ||||
if not one: | ||||
yield "a local result", None | ||||
# Build list of encoded arguments suitable for your wire protocol: | ||||
encargs = [('one', encode(one),), ('two', encode(two),)] | ||||
# Create future for injection of encoded result: | ||||
encresref = future() | ||||
# Return encoded arguments and future: | ||||
yield encargs, encresref | ||||
|
r16683 | # Assuming the future to be filled with the result from the batched | ||
# request now. Decode it: | ||||
|
r14621 | yield decode(encresref.value) | ||
|
r16683 | The decorator returns a function which wraps this coroutine as a plain | ||
method, but adds the original method as an attribute called "batchable", | ||||
which is used by remotebatch to split the call into separate encoding and | ||||
decoding phases. | ||||
|
r14621 | ''' | ||
def plain(*args, **opts): | ||||
batchable = f(*args, **opts) | ||||
encargsorres, encresref = batchable.next() | ||||
if not encresref: | ||||
return encargsorres # a local result in this case | ||||
self = args[0] | ||||
encresref.set(self._submitone(f.func_name, encargsorres)) | ||||
return batchable.next() | ||||
setattr(plain, 'batchable', f) | ||||
return plain | ||||
|
r11597 | # list of nodes encoding / decoding | ||
def decodelist(l, sep=' '): | ||||
|
r13722 | if l: | ||
return map(bin, l.split(sep)) | ||||
return [] | ||||
|
r11597 | |||
def encodelist(l, sep=' '): | ||||
return sep.join(map(hex, l)) | ||||
|
r14622 | # batched call argument encoding | ||
def escapearg(plain): | ||||
return (plain | ||||
.replace(':', '::') | ||||
.replace(',', ':,') | ||||
.replace(';', ':;') | ||||
.replace('=', ':=')) | ||||
def unescapearg(escaped): | ||||
return (escaped | ||||
.replace(':=', '=') | ||||
.replace(':;', ';') | ||||
.replace(':,', ',') | ||||
.replace('::', ':')) | ||||
|
r11586 | # client side | ||
|
r17192 | class wirepeer(peer.peerrepository): | ||
|
r14622 | |||
def batch(self): | ||||
return remotebatch(self) | ||||
def _submitbatch(self, req): | ||||
cmds = [] | ||||
for op, argsdict in req: | ||||
args = ','.join('%s=%s' % p for p in argsdict.iteritems()) | ||||
cmds.append('%s %s' % (op, args)) | ||||
rsp = self._call("batch", cmds=';'.join(cmds)) | ||||
return rsp.split(';') | ||||
def _submitone(self, op, args): | ||||
return self._call(op, **args) | ||||
|
r14623 | @batchable | ||
|
r11586 | def lookup(self, key): | ||
self.requirecap('lookup', _('look up remote revision')) | ||||
|
r14623 | f = future() | ||
|
r20671 | yield {'key': encoding.fromlocal(key)}, f | ||
|
r14623 | d = f.value | ||
|
r11586 | success, data = d[:-1].split(" ", 1) | ||
if int(success): | ||||
|
r14623 | yield bin(data) | ||
|
r11586 | self._abort(error.RepoError(data)) | ||
|
r14623 | @batchable | ||
|
r11586 | def heads(self): | ||
|
r14623 | f = future() | ||
yield {}, f | ||||
d = f.value | ||||
|
r11586 | try: | ||
|
r14623 | yield decodelist(d[:-1]) | ||
|
r13726 | except ValueError: | ||
|
r11879 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
|
r11586 | |||
|
r14623 | @batchable | ||
|
r13723 | def known(self, nodes): | ||
|
r14623 | f = future() | ||
|
r20671 | yield {'nodes': encodelist(nodes)}, f | ||
|
r14623 | d = f.value | ||
|
r13723 | try: | ||
|
r14623 | yield [bool(int(f)) for f in d] | ||
|
r13726 | except ValueError: | ||
|
r13723 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
|
r14623 | @batchable | ||
|
r11586 | def branchmap(self): | ||
|
r14623 | f = future() | ||
yield {}, f | ||||
d = f.value | ||||
|
r11586 | try: | ||
branchmap = {} | ||||
for branchpart in d.splitlines(): | ||||
|
r11597 | branchname, branchheads = branchpart.split(' ', 1) | ||
|
r13047 | branchname = encoding.tolocal(urllib.unquote(branchname)) | ||
|
r11597 | branchheads = decodelist(branchheads) | ||
|
r11586 | branchmap[branchname] = branchheads | ||
|
r14623 | yield branchmap | ||
|
r11586 | except TypeError: | ||
self._abort(error.ResponseError(_("unexpected response:"), d)) | ||||
def branches(self, nodes): | ||||
|
r11597 | n = encodelist(nodes) | ||
|
r11586 | d = self._call("branches", nodes=n) | ||
try: | ||||
|
r11597 | br = [tuple(decodelist(b)) for b in d.splitlines()] | ||
|
r11586 | return br | ||
|
r13726 | except ValueError: | ||
|
r11586 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
def between(self, pairs): | ||||
|
r11587 | batch = 8 # avoid giant requests | ||
r = [] | ||||
for i in xrange(0, len(pairs), batch): | ||||
|
r11597 | n = " ".join([encodelist(p, '-') for p in pairs[i:i + batch]]) | ||
|
r11587 | d = self._call("between", pairs=n) | ||
try: | ||||
|
r11597 | r.extend(l and decodelist(l) or [] for l in d.splitlines()) | ||
|
r13726 | except ValueError: | ||
|
r11587 | self._abort(error.ResponseError(_("unexpected response:"), d)) | ||
return r | ||||
|
r11586 | |||
|
r14623 | @batchable | ||
|
r11586 | def pushkey(self, namespace, key, old, new): | ||
if not self.capable('pushkey'): | ||||
|
r14623 | yield False, None | ||
f = future() | ||||
|
r17293 | self.ui.debug('preparing pushkey for "%s:%s"\n' % (namespace, key)) | ||
|
r20671 | yield {'namespace': encoding.fromlocal(namespace), | ||
'key': encoding.fromlocal(key), | ||||
'old': encoding.fromlocal(old), | ||||
'new': encoding.fromlocal(new)}, f | ||||
|
r14623 | d = f.value | ||
|
r15652 | d, output = d.split('\n', 1) | ||
|
r13450 | try: | ||
d = bool(int(d)) | ||||
except ValueError: | ||||
raise error.ResponseError( | ||||
_('push failed (unexpected response):'), d) | ||||
|
r15652 | for l in output.splitlines(True): | ||
self.ui.status(_('remote: '), l) | ||||
|
r14623 | yield d | ||
|
r11586 | |||
|
r14623 | @batchable | ||
|
r11586 | def listkeys(self, namespace): | ||
if not self.capable('pushkey'): | ||||
|
r14623 | yield {}, None | ||
f = future() | ||||
|
r17293 | self.ui.debug('preparing listkeys for "%s"\n' % namespace) | ||
|
r20671 | yield {'namespace': encoding.fromlocal(namespace)}, f | ||
|
r14623 | d = f.value | ||
|
r11586 | r = {} | ||
for l in d.splitlines(): | ||||
k, v = l.split('\t') | ||||
|
r13050 | r[encoding.tolocal(k)] = encoding.tolocal(v) | ||
|
r14623 | yield r | ||
|
r11586 | |||
|
r11588 | def stream_out(self): | ||
return self._callstream('stream_out') | ||||
|
r11591 | def changegroup(self, nodes, kind): | ||
|
r11597 | n = encodelist(nodes) | ||
|
r20905 | f = self._callcompressable("changegroup", roots=n) | ||
return changegroupmod.unbundle10(f, 'UN') | ||||
|
r11591 | |||
def changegroupsubset(self, bases, heads, kind): | ||||
self.requirecap('changegroupsubset', _('look up remote changes')) | ||||
|
r11597 | bases = encodelist(bases) | ||
heads = encodelist(heads) | ||||
|
r20905 | f = self._callcompressable("changegroupsubset", | ||
bases=bases, heads=heads) | ||||
return changegroupmod.unbundle10(f, 'UN') | ||||
|
r11591 | |||
|
r21157 | def getbundle(self, source, heads=None, common=None, bundlecaps=None, | ||
**kwargs): | ||||
|
r13741 | self.requirecap('getbundle', _('look up remote changes')) | ||
opts = {} | ||||
if heads is not None: | ||||
opts['heads'] = encodelist(heads) | ||||
if common is not None: | ||||
opts['common'] = encodelist(common) | ||||
|
r19201 | if bundlecaps is not None: | ||
opts['bundlecaps'] = ','.join(bundlecaps) | ||||
|
r21157 | opts.update(kwargs) | ||
|
r20905 | f = self._callcompressable("getbundle", **opts) | ||
|
r21144 | if bundlecaps is not None and 'HG2X' in bundlecaps: | ||
|
r21069 | return bundle2.unbundle20(self.ui, f) | ||
else: | ||||
return changegroupmod.unbundle10(f, 'UN') | ||||
|
r13741 | |||
|
r11592 | def unbundle(self, cg, heads, source): | ||
'''Send cg (a readable file-like object representing the | ||||
changegroup to push, typically a chunkbuffer object) to the | ||||
|
r21075 | remote server as a bundle. | ||
When pushing a bundle10 stream, return an integer indicating the | ||||
result of the push (see localrepository.addchangegroup()). | ||||
When pushing a bundle20 stream, return a bundle20 stream.''' | ||||
|
r11592 | |||
|
r14419 | if heads != ['force'] and self.capable('unbundlehash'): | ||
|
r13942 | heads = encodelist(['hashed', | ||
util.sha1(''.join(sorted(heads))).digest()]) | ||||
else: | ||||
heads = encodelist(heads) | ||||
|
r21075 | if util.safehasattr(cg, 'deltaheader'): | ||
# this a bundle10, do the old style call sequence | ||||
ret, output = self._callpush("unbundle", cg, heads=heads) | ||||
if ret == "": | ||||
raise error.ResponseError( | ||||
_('push failed:'), output) | ||||
try: | ||||
ret = int(ret) | ||||
except ValueError: | ||||
raise error.ResponseError( | ||||
_('push failed (unexpected response):'), ret) | ||||
|
r11592 | |||
|
r21075 | for l in output.splitlines(True): | ||
self.ui.status(_('remote: '), l) | ||||
else: | ||||
# bundle2 push. Send a stream, fetch a stream. | ||||
stream = self._calltwowaystream('unbundle', cg, heads=heads) | ||||
ret = bundle2.unbundle20(self.ui, stream) | ||||
|
r11592 | return ret | ||
|
r14048 | def debugwireargs(self, one, two, three=None, four=None, five=None): | ||
|
r13720 | # don't pass optional arguments left at their default value | ||
opts = {} | ||||
if three is not None: | ||||
opts['three'] = three | ||||
if four is not None: | ||||
opts['four'] = four | ||||
return self._call('debugwireargs', one=one, two=two, **opts) | ||||
|
r20904 | def _call(self, cmd, **args): | ||
"""execute <cmd> on the server | ||||
The command is expected to return a simple string. | ||||
returns the server reply as a string.""" | ||||
raise NotImplementedError() | ||||
def _callstream(self, cmd, **args): | ||||
"""execute <cmd> on the server | ||||
The command is expected to return a stream. | ||||
returns the server reply as a file like object.""" | ||||
raise NotImplementedError() | ||||
|
r20905 | def _callcompressable(self, cmd, **args): | ||
"""execute <cmd> on the server | ||||
The command is expected to return a stream. | ||||
|
r21024 | The stream may have been compressed in some implementations. This | ||
|
r20905 | function takes care of the decompression. This is the only difference | ||
with _callstream. | ||||
returns the server reply as a file like object. | ||||
""" | ||||
raise NotImplementedError() | ||||
|
r20904 | def _callpush(self, cmd, fp, **args): | ||
"""execute a <cmd> on server | ||||
The command is expected to be related to a push. Push has a special | ||||
return method. | ||||
returns the server reply as a (ret, output) tuple. ret is either | ||||
empty (error) or a stringified int. | ||||
""" | ||||
raise NotImplementedError() | ||||
|
r21072 | def _calltwowaystream(self, cmd, fp, **args): | ||
"""execute <cmd> on server | ||||
The command will send a stream to the server and get a stream in reply. | ||||
""" | ||||
raise NotImplementedError() | ||||
|
r20904 | def _abort(self, exception): | ||
"""clearly abort the wire protocol connection and raise the exception | ||||
""" | ||||
raise NotImplementedError() | ||||
|
r11586 | # server side | ||
|
r20902 | # wire protocol command can either return a string or one of these classes. | ||
|
r11625 | class streamres(object): | ||
|
r20902 | """wireproto reply: binary stream | ||
The call was successful and the result is a stream. | ||||
Iterate on the `self.gen` attribute to retrieve chunks. | ||||
""" | ||||
|
r11625 | def __init__(self, gen): | ||
self.gen = gen | ||||
class pushres(object): | ||||
|
r20902 | """wireproto reply: success with simple integer return | ||
The call was successful and returned an integer contained in `self.res`. | ||||
""" | ||||
|
r11625 | def __init__(self, res): | ||
self.res = res | ||||
|
r12703 | class pusherr(object): | ||
|
r20902 | """wireproto reply: failure | ||
The call failed. The `self.res` attribute contains the error message. | ||||
""" | ||||
|
r12703 | def __init__(self, res): | ||
self.res = res | ||||
|
r15017 | class ooberror(object): | ||
|
r20902 | """wireproto reply: failure of a batch of operation | ||
Something failed during a batch call. The error message is stored in | ||||
`self.message`. | ||||
""" | ||||
|
r15017 | def __init__(self, message): | ||
self.message = message | ||||
|
r11581 | def dispatch(repo, proto, command): | ||
|
r18382 | repo = repo.filtered("served") | ||
|
r11581 | func, spec = commands[command] | ||
args = proto.getargs(spec) | ||||
|
r11625 | return func(repo, proto, *args) | ||
|
r11581 | |||
|
r13721 | def options(cmd, keys, others): | ||
opts = {} | ||||
for k in keys: | ||||
if k in others: | ||||
opts[k] = others[k] | ||||
del others[k] | ||||
if others: | ||||
sys.stderr.write("abort: %s got unexpected arguments %s\n" | ||||
% (cmd, ",".join(others))) | ||||
return opts | ||||
|
r20906 | # list of commands | ||
commands = {} | ||||
def wireprotocommand(name, args=''): | ||||
|
r21024 | """decorator for wire protocol command""" | ||
|
r20906 | def register(func): | ||
commands[name] = (func, args) | ||||
return func | ||||
return register | ||||
|
r20907 | @wireprotocommand('batch', 'cmds *') | ||
|
r14622 | def batch(repo, proto, cmds, others): | ||
|
r18382 | repo = repo.filtered("served") | ||
|
r14622 | res = [] | ||
for pair in cmds.split(';'): | ||||
op, args = pair.split(' ', 1) | ||||
vals = {} | ||||
for a in args.split(','): | ||||
if a: | ||||
n, v = a.split('=') | ||||
vals[n] = unescapearg(v) | ||||
func, spec = commands[op] | ||||
if spec: | ||||
keys = spec.split() | ||||
data = {} | ||||
for k in keys: | ||||
if k == '*': | ||||
star = {} | ||||
for key in vals.keys(): | ||||
if key not in keys: | ||||
star[key] = vals[key] | ||||
data['*'] = star | ||||
else: | ||||
data[k] = vals[k] | ||||
result = func(repo, proto, *[data[k] for k in keys]) | ||||
else: | ||||
result = func(repo, proto) | ||||
|
r15017 | if isinstance(result, ooberror): | ||
return result | ||||
|
r14622 | res.append(escapearg(result)) | ||
return ';'.join(res) | ||||
|
r20908 | @wireprotocommand('between', 'pairs') | ||
|
r11583 | def between(repo, proto, pairs): | ||
|
r11597 | pairs = [decodelist(p, '-') for p in pairs.split(" ")] | ||
|
r11581 | r = [] | ||
for b in repo.between(pairs): | ||||
|
r11597 | r.append(encodelist(b) + "\n") | ||
|
r11581 | return "".join(r) | ||
|
r20909 | @wireprotocommand('branchmap') | ||
|
r11583 | def branchmap(repo, proto): | ||
|
r18281 | branchmap = repo.branchmap() | ||
|
r11581 | heads = [] | ||
for branch, nodes in branchmap.iteritems(): | ||||
|
r13047 | branchname = urllib.quote(encoding.fromlocal(branch)) | ||
|
r11597 | branchnodes = encodelist(nodes) | ||
heads.append('%s %s' % (branchname, branchnodes)) | ||||
|
r11581 | return '\n'.join(heads) | ||
|
r20910 | @wireprotocommand('branches', 'nodes') | ||
|
r11583 | def branches(repo, proto, nodes): | ||
|
r11597 | nodes = decodelist(nodes) | ||
|
r11581 | r = [] | ||
for b in repo.branches(nodes): | ||||
|
r11597 | r.append(encodelist(b) + "\n") | ||
|
r11581 | return "".join(r) | ||
|
r20774 | |||
wireprotocaps = ['lookup', 'changegroupsubset', 'branchmap', 'pushkey', | ||||
'known', 'getbundle', 'unbundlehash', 'batch'] | ||||
|
r20775 | |||
def _capabilities(repo, proto): | ||||
"""return a list of capabilities for a repo | ||||
This function exists to allow extensions to easily wrap capabilities | ||||
computation | ||||
- returns a lists: easy to alter | ||||
- change done here will be propagated to both `capabilities` and `hello` | ||||
|
r21024 | command without any other action needed. | ||
|
r20775 | """ | ||
|
r20774 | # copy to prevent modification of the global list | ||
caps = list(wireprotocaps) | ||||
|
r11627 | if _allowstream(repo.ui): | ||
|
r16361 | if repo.ui.configbool('server', 'preferuncompressed', False): | ||
caps.append('stream-preferred') | ||||
|
r12296 | requiredformats = repo.requirements & repo.supportedformats | ||
# if our local revlogs are just revlogv1, add 'stream' cap | ||||
if not requiredformats - set(('revlogv1',)): | ||||
caps.append('stream') | ||||
# otherwise, add 'streamreqs' detailing our local revlog format | ||||
else: | ||||
caps.append('streamreqs=%s' % ','.join(requiredformats)) | ||||
|
r21148 | if repo.ui.configbool('experimental', 'bundle2-exp', False): | ||
|
r21141 | capsblob = bundle2.encodecaps(repo.bundle2caps) | ||
|
r21145 | caps.append('bundle2-exp=' + urllib.quote(capsblob)) | ||
|
r11594 | caps.append('unbundle=%s' % ','.join(changegroupmod.bundlepriority)) | ||
|
r14093 | caps.append('httpheader=1024') | ||
|
r20775 | return caps | ||
|
r21024 | # If you are writing an extension and consider wrapping this function. Wrap | ||
|
r20775 | # `_capabilities` instead. | ||
|
r20911 | @wireprotocommand('capabilities') | ||
|
r20775 | def capabilities(repo, proto): | ||
return ' '.join(_capabilities(repo, proto)) | ||||
|
r11594 | |||
|
r20912 | @wireprotocommand('changegroup', 'roots') | ||
|
r11584 | def changegroup(repo, proto, roots): | ||
|
r11597 | nodes = decodelist(roots) | ||
|
r20931 | cg = changegroupmod.changegroup(repo, nodes, 'serve') | ||
|
r11625 | return streamres(proto.groupchunks(cg)) | ||
|
r11584 | |||
|
r20913 | @wireprotocommand('changegroupsubset', 'bases heads') | ||
|
r11584 | def changegroupsubset(repo, proto, bases, heads): | ||
|
r11597 | bases = decodelist(bases) | ||
heads = decodelist(heads) | ||||
|
r20927 | cg = changegroupmod.changegroupsubset(repo, bases, heads, 'serve') | ||
|
r11625 | return streamres(proto.groupchunks(cg)) | ||
|
r11584 | |||
|
r20914 | @wireprotocommand('debugwireargs', 'one two *') | ||
|
r13721 | def debugwireargs(repo, proto, one, two, others): | ||
# only accept optional args from the known set | ||||
opts = options('debugwireargs', ['three', 'four'], others) | ||||
return repo.debugwireargs(one, two, **opts) | ||||
|
r13720 | |||
|
r20915 | @wireprotocommand('getbundle', '*') | ||
|
r13741 | def getbundle(repo, proto, others): | ||
|
r19201 | opts = options('getbundle', ['heads', 'common', 'bundlecaps'], others) | ||
|
r13741 | for k, v in opts.iteritems(): | ||
|
r19201 | if k in ('heads', 'common'): | ||
opts[k] = decodelist(v) | ||||
elif k == 'bundlecaps': | ||||
opts[k] = set(v.split(',')) | ||||
|
r21069 | cg = exchange.getbundle(repo, 'serve', **opts) | ||
|
r13741 | return streamres(proto.groupchunks(cg)) | ||
|
r20916 | @wireprotocommand('heads') | ||
|
r11583 | def heads(repo, proto): | ||
|
r18281 | h = repo.heads() | ||
|
r11597 | return encodelist(h) + "\n" | ||
|
r11581 | |||
|
r20917 | @wireprotocommand('hello') | ||
|
r11594 | def hello(repo, proto): | ||
'''the hello command returns a set of lines describing various | ||||
interesting things about the server, in an RFC822-like format. | ||||
Currently the only one defined is "capabilities", which | ||||
consists of a line in the form: | ||||
capabilities: space separated list of tokens | ||||
''' | ||||
return "capabilities: %s\n" % (capabilities(repo, proto)) | ||||
|
r20919 | @wireprotocommand('listkeys', 'namespace') | ||
|
r11583 | def listkeys(repo, proto, namespace): | ||
|
r15217 | d = repo.listkeys(encoding.tolocal(namespace)).items() | ||
|
r13050 | t = '\n'.join(['%s\t%s' % (encoding.fromlocal(k), encoding.fromlocal(v)) | ||
for k, v in d]) | ||||
|
r11581 | return t | ||
|
r20920 | @wireprotocommand('lookup', 'key') | ||
|
r11583 | def lookup(repo, proto, key): | ||
|
r11581 | try: | ||
|
r15925 | k = encoding.tolocal(key) | ||
c = repo[k] | ||||
r = c.hex() | ||||
|
r11581 | success = 1 | ||
except Exception, inst: | ||||
r = str(inst) | ||||
success = 0 | ||||
return "%s %s\n" % (success, r) | ||||
|
r20918 | @wireprotocommand('known', 'nodes *') | ||
|
r14436 | def known(repo, proto, nodes, others): | ||
|
r13723 | return ''.join(b and "1" or "0" for b in repo.known(decodelist(nodes))) | ||
|
r20921 | @wireprotocommand('pushkey', 'namespace key old new') | ||
|
r11583 | def pushkey(repo, proto, namespace, key, old, new): | ||
|
r13050 | # compatibility with pre-1.8 clients which were accidentally | ||
# sending raw binary nodes rather than utf-8-encoded hex | ||||
if len(new) == 20 and new.encode('string-escape') != new: | ||||
# looks like it could be a binary node | ||||
try: | ||||
|
r14064 | new.decode('utf-8') | ||
|
r13050 | new = encoding.tolocal(new) # but cleanly decodes as UTF-8 | ||
except UnicodeDecodeError: | ||||
pass # binary, leave unmodified | ||||
else: | ||||
new = encoding.tolocal(new) # normal path | ||||
|
r17793 | if util.safehasattr(proto, 'restore'): | ||
proto.redirect() | ||||
try: | ||||
r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key), | ||||
encoding.tolocal(old), new) or False | ||||
except util.Abort: | ||||
r = False | ||||
output = proto.restore() | ||||
return '%s\n%s' % (int(r), output) | ||||
|
r15217 | r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key), | ||
encoding.tolocal(old), new) | ||||
|
r11581 | return '%s\n' % int(r) | ||
|
r11627 | def _allowstream(ui): | ||
return ui.configbool('server', 'uncompressed', True, untrusted=True) | ||||
|
r19176 | def _walkstreamfiles(repo): | ||
# this is it's own function so extensions can override it | ||||
return repo.store.walk() | ||||
|
r20922 | @wireprotocommand('stream_out') | ||
|
r11585 | def stream(repo, proto): | ||
|
r11627 | '''If the server supports streaming clone, it advertises the "stream" | ||
capability with a value representing the version and flags of the repo | ||||
it is serving. Client checks to see if it understands the format. | ||||
The format is simple: the server writes out a line with the amount | ||||
|
r17424 | of files, then the total amount of bytes to be transferred (separated | ||
|
r11627 | by a space). Then, for each file, the server first writes the filename | ||
|
r21024 | and file size (separated by the null character), then the file contents. | ||
|
r11627 | ''' | ||
if not _allowstream(repo.ui): | ||||
return '1\n' | ||||
entries = [] | ||||
total_bytes = 0 | ||||
try: | ||||
# get consistent snapshot of repo, lock during scan | ||||
lock = repo.lock() | ||||
try: | ||||
repo.ui.debug('scanning\n') | ||||
|
r19176 | for name, ename, size in _walkstreamfiles(repo): | ||
|
r18381 | if size: | ||
entries.append((name, size)) | ||||
total_bytes += size | ||||
|
r11627 | finally: | ||
lock.release() | ||||
except error.LockError: | ||||
return '2\n' # error: 2 | ||||
def streamer(repo, entries, total): | ||||
'''stream out all metadata files in repository.''' | ||||
yield '0\n' # success | ||||
repo.ui.debug('%d files, %d bytes to transfer\n' % | ||||
(len(entries), total_bytes)) | ||||
yield '%d %d\n' % (len(entries), total_bytes) | ||||
|
r17556 | |||
sopener = repo.sopener | ||||
oldaudit = sopener.mustaudit | ||||
|
r17558 | debugflag = repo.ui.debugflag | ||
|
r17556 | sopener.mustaudit = False | ||
try: | ||||
for name, size in entries: | ||||
|
r17558 | if debugflag: | ||
repo.ui.debug('sending %s (%d bytes)\n' % (name, size)) | ||||
|
r17556 | # partially encode name over the wire for backwards compat | ||
yield '%s\0%d\n' % (store.encodedir(name), size) | ||||
|
r17557 | if size <= 65536: | ||
|
r17567 | fp = sopener(name) | ||
try: | ||||
data = fp.read(size) | ||||
finally: | ||||
fp.close() | ||||
yield data | ||||
|
r17557 | else: | ||
for chunk in util.filechunkiter(sopener(name), limit=size): | ||||
yield chunk | ||||
|
r17603 | # replace with "finally:" when support for python 2.4 has been dropped | ||
except Exception: | ||||
|
r17556 | sopener.mustaudit = oldaudit | ||
|
r17603 | raise | ||
sopener.mustaudit = oldaudit | ||||
|
r11627 | |||
return streamres(streamer(repo, entries, total_bytes)) | ||||
|
r11585 | |||
|
r20923 | @wireprotocommand('unbundle', 'heads') | ||
|
r11593 | def unbundle(repo, proto, heads): | ||
|
r11597 | their_heads = decodelist(heads) | ||
|
r11593 | |||
|
r20967 | try: | ||
proto.redirect() | ||||
|
r11593 | |||
|
r20967 | exchange.check_heads(repo, their_heads, 'preparing changes') | ||
|
r12702 | |||
|
r20967 | # write bundle data to temporary file because it can be big | ||
fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-') | ||||
fp = os.fdopen(fd, 'wb+') | ||||
r = 0 | ||||
|
r11593 | try: | ||
|
r20967 | proto.getfile(fp) | ||
|
r20968 | fp.seek(0) | ||
|
r21064 | gen = exchange.readbundle(repo.ui, fp, None) | ||
|
r20968 | r = exchange.unbundle(repo, gen, their_heads, 'serve', | ||
proto._client()) | ||||
|
r21075 | if util.safehasattr(r, 'addpart'): | ||
# The return looks streameable, we are in the bundle2 case and | ||||
# should return a stream. | ||||
return streamres(r.getchunks()) | ||||
|
r20967 | return pushres(r) | ||
|
r11593 | finally: | ||
|
r20967 | fp.close() | ||
os.unlink(tempname) | ||||
|
r21618 | except error.BundleValueError, exc: | ||
|
r21183 | bundler = bundle2.bundle20(repo.ui) | ||
|
r21600 | bundler.newpart('B2X:ERROR:UNKNOWNPART', [('parttype', str(exc))]) | ||
|
r21183 | return streamres(bundler.getchunks()) | ||
|
r20969 | except util.Abort, inst: | ||
# The old code we moved used sys.stderr directly. | ||||
|
r21024 | # We did not change it to minimise code change. | ||
|
r20969 | # This need to be moved to something proper. | ||
# Feel free to do it. | ||||
|
r21177 | if getattr(inst, 'duringunbundle2', False): | ||
bundler = bundle2.bundle20(repo.ui) | ||||
manargs = [('message', str(inst))] | ||||
advargs = [] | ||||
if inst.hint is not None: | ||||
advargs.append(('hint', inst.hint)) | ||||
bundler.addpart(bundle2.bundlepart('B2X:ERROR:ABORT', | ||||
manargs, advargs)) | ||||
return streamres(bundler.getchunks()) | ||||
else: | ||||
sys.stderr.write("abort: %s\n" % inst) | ||||
return pushres(0) | ||||
|
r21184 | except error.PushRaced, exc: | ||
|
r21186 | if getattr(exc, 'duringunbundle2', False): | ||
bundler = bundle2.bundle20(repo.ui) | ||||
|
r21600 | bundler.newpart('B2X:ERROR:PUSHRACED', [('message', str(exc))]) | ||
|
r21186 | return streamres(bundler.getchunks()) | ||
else: | ||||
return pusherr(str(exc)) | ||||