comparison mercurial/wireproto.py @ 36074:2f7290555c96

wireproto: introduce type for raw byte responses (API) Right now we simply return a str/bytes instance for simple responses. I want all wire protocol response types to be strongly typed. So let's invent and use a type for raw bytes responses. .. api:: Wire protocol command handlers now return a wireprototypes.bytesresponse instead of a raw bytes instance. Protocol handlers will continue handling bytes instances. However, any extensions wrapping wire protocol commands will need to handle the new type. Differential Revision: https://phab.mercurial-scm.org/D2089
author Gregory Szorc <gregory.szorc@gmail.com>
date Wed, 07 Feb 2018 20:27:36 -0800
parents cd6ab329c5c7
children 62bca1c50e96
comparison
equal deleted inserted replaced
36073:cd6ab329c5c7 36074:2f7290555c96
35 ) 35 )
36 36
37 urlerr = util.urlerr 37 urlerr = util.urlerr
38 urlreq = util.urlreq 38 urlreq = util.urlreq
39 39
40 bytesresponse = wireprototypes.bytesresponse
40 ooberror = wireprototypes.ooberror 41 ooberror = wireprototypes.ooberror
41 pushres = wireprototypes.pushres 42 pushres = wireprototypes.pushres
42 pusherr = wireprototypes.pusherr 43 pusherr = wireprototypes.pusherr
43 streamres = wireprototypes.streamres 44 streamres = wireprototypes.streamres
44 streamres_legacy = wireprototypes.streamreslegacy 45 streamres_legacy = wireprototypes.streamreslegacy
694 result = func(repo, proto, *[data[k] for k in keys]) 695 result = func(repo, proto, *[data[k] for k in keys])
695 else: 696 else:
696 result = func(repo, proto) 697 result = func(repo, proto)
697 if isinstance(result, ooberror): 698 if isinstance(result, ooberror):
698 return result 699 return result
700
701 # For now, all batchable commands must return bytesresponse or
702 # raw bytes (for backwards compatibility).
703 assert isinstance(result, (bytesresponse, bytes))
704 if isinstance(result, bytesresponse):
705 result = result.data
699 res.append(escapearg(result)) 706 res.append(escapearg(result))
700 return ';'.join(res) 707
708 return bytesresponse(';'.join(res))
701 709
702 @wireprotocommand('between', 'pairs') 710 @wireprotocommand('between', 'pairs')
703 def between(repo, proto, pairs): 711 def between(repo, proto, pairs):
704 pairs = [decodelist(p, '-') for p in pairs.split(" ")] 712 pairs = [decodelist(p, '-') for p in pairs.split(" ")]
705 r = [] 713 r = []
706 for b in repo.between(pairs): 714 for b in repo.between(pairs):
707 r.append(encodelist(b) + "\n") 715 r.append(encodelist(b) + "\n")
708 return "".join(r) 716
717 return bytesresponse(''.join(r))
709 718
710 @wireprotocommand('branchmap') 719 @wireprotocommand('branchmap')
711 def branchmap(repo, proto): 720 def branchmap(repo, proto):
712 branchmap = repo.branchmap() 721 branchmap = repo.branchmap()
713 heads = [] 722 heads = []
714 for branch, nodes in branchmap.iteritems(): 723 for branch, nodes in branchmap.iteritems():
715 branchname = urlreq.quote(encoding.fromlocal(branch)) 724 branchname = urlreq.quote(encoding.fromlocal(branch))
716 branchnodes = encodelist(nodes) 725 branchnodes = encodelist(nodes)
717 heads.append('%s %s' % (branchname, branchnodes)) 726 heads.append('%s %s' % (branchname, branchnodes))
718 return '\n'.join(heads) 727
728 return bytesresponse('\n'.join(heads))
719 729
720 @wireprotocommand('branches', 'nodes') 730 @wireprotocommand('branches', 'nodes')
721 def branches(repo, proto, nodes): 731 def branches(repo, proto, nodes):
722 nodes = decodelist(nodes) 732 nodes = decodelist(nodes)
723 r = [] 733 r = []
724 for b in repo.branches(nodes): 734 for b in repo.branches(nodes):
725 r.append(encodelist(b) + "\n") 735 r.append(encodelist(b) + "\n")
726 return "".join(r) 736
737 return bytesresponse(''.join(r))
727 738
728 @wireprotocommand('clonebundles', '') 739 @wireprotocommand('clonebundles', '')
729 def clonebundles(repo, proto): 740 def clonebundles(repo, proto):
730 """Server command for returning info for available bundles to seed clones. 741 """Server command for returning info for available bundles to seed clones.
731 742
733 744
734 Extensions may wrap this command to filter or dynamically emit data 745 Extensions may wrap this command to filter or dynamically emit data
735 depending on the request. e.g. you could advertise URLs for the closest 746 depending on the request. e.g. you could advertise URLs for the closest
736 data center given the client's IP address. 747 data center given the client's IP address.
737 """ 748 """
738 return repo.vfs.tryread('clonebundles.manifest') 749 return bytesresponse(repo.vfs.tryread('clonebundles.manifest'))
739 750
740 wireprotocaps = ['lookup', 'changegroupsubset', 'branchmap', 'pushkey', 751 wireprotocaps = ['lookup', 'changegroupsubset', 'branchmap', 'pushkey',
741 'known', 'getbundle', 'unbundlehash', 'batch'] 752 'known', 'getbundle', 'unbundlehash', 'batch']
742 753
743 def _capabilities(repo, proto): 754 def _capabilities(repo, proto):
787 798
788 # If you are writing an extension and consider wrapping this function. Wrap 799 # If you are writing an extension and consider wrapping this function. Wrap
789 # `_capabilities` instead. 800 # `_capabilities` instead.
790 @wireprotocommand('capabilities') 801 @wireprotocommand('capabilities')
791 def capabilities(repo, proto): 802 def capabilities(repo, proto):
792 return ' '.join(_capabilities(repo, proto)) 803 return bytesresponse(' '.join(_capabilities(repo, proto)))
793 804
794 @wireprotocommand('changegroup', 'roots') 805 @wireprotocommand('changegroup', 'roots')
795 def changegroup(repo, proto, roots): 806 def changegroup(repo, proto, roots):
796 nodes = decodelist(roots) 807 nodes = decodelist(roots)
797 outgoing = discovery.outgoing(repo, missingroots=nodes, 808 outgoing = discovery.outgoing(repo, missingroots=nodes,
812 823
813 @wireprotocommand('debugwireargs', 'one two *') 824 @wireprotocommand('debugwireargs', 'one two *')
814 def debugwireargs(repo, proto, one, two, others): 825 def debugwireargs(repo, proto, one, two, others):
815 # only accept optional args from the known set 826 # only accept optional args from the known set
816 opts = options('debugwireargs', ['three', 'four'], others) 827 opts = options('debugwireargs', ['three', 'four'], others)
817 return repo.debugwireargs(one, two, **pycompat.strkwargs(opts)) 828 return bytesresponse(repo.debugwireargs(one, two,
829 **pycompat.strkwargs(opts)))
818 830
819 @wireprotocommand('getbundle', '*') 831 @wireprotocommand('getbundle', '*')
820 def getbundle(repo, proto, others): 832 def getbundle(repo, proto, others):
821 opts = options('getbundle', gboptsmap.keys(), others) 833 opts = options('getbundle', gboptsmap.keys(), others)
822 for k, v in opts.iteritems(): 834 for k, v in opts.iteritems():
883 return streamres(gen=chunks, prefer_uncompressed=not prefercompressed) 895 return streamres(gen=chunks, prefer_uncompressed=not prefercompressed)
884 896
885 @wireprotocommand('heads') 897 @wireprotocommand('heads')
886 def heads(repo, proto): 898 def heads(repo, proto):
887 h = repo.heads() 899 h = repo.heads()
888 return encodelist(h) + "\n" 900 return bytesresponse(encodelist(h) + '\n')
889 901
890 @wireprotocommand('hello') 902 @wireprotocommand('hello')
891 def hello(repo, proto): 903 def hello(repo, proto):
892 '''the hello command returns a set of lines describing various 904 '''the hello command returns a set of lines describing various
893 interesting things about the server, in an RFC822-like format. 905 interesting things about the server, in an RFC822-like format.
894 Currently the only one defined is "capabilities", which 906 Currently the only one defined is "capabilities", which
895 consists of a line in the form: 907 consists of a line in the form:
896 908
897 capabilities: space separated list of tokens 909 capabilities: space separated list of tokens
898 ''' 910 '''
899 return "capabilities: %s\n" % (capabilities(repo, proto)) 911 caps = capabilities(repo, proto).data
912 return bytesresponse('capabilities: %s\n' % caps)
900 913
901 @wireprotocommand('listkeys', 'namespace') 914 @wireprotocommand('listkeys', 'namespace')
902 def listkeys(repo, proto, namespace): 915 def listkeys(repo, proto, namespace):
903 d = repo.listkeys(encoding.tolocal(namespace)).items() 916 d = repo.listkeys(encoding.tolocal(namespace)).items()
904 return pushkeymod.encodekeys(d) 917 return bytesresponse(pushkeymod.encodekeys(d))
905 918
906 @wireprotocommand('lookup', 'key') 919 @wireprotocommand('lookup', 'key')
907 def lookup(repo, proto, key): 920 def lookup(repo, proto, key):
908 try: 921 try:
909 k = encoding.tolocal(key) 922 k = encoding.tolocal(key)
911 r = c.hex() 924 r = c.hex()
912 success = 1 925 success = 1
913 except Exception as inst: 926 except Exception as inst:
914 r = str(inst) 927 r = str(inst)
915 success = 0 928 success = 0
916 return "%d %s\n" % (success, r) 929 return bytesresponse('%d %s\n' % (success, r))
917 930
918 @wireprotocommand('known', 'nodes *') 931 @wireprotocommand('known', 'nodes *')
919 def known(repo, proto, nodes, others): 932 def known(repo, proto, nodes, others):
920 return ''.join(b and "1" or "0" for b in repo.known(decodelist(nodes))) 933 v = ''.join(b and '1' or '0' for b in repo.known(decodelist(nodes)))
934 return bytesresponse(v)
921 935
922 @wireprotocommand('pushkey', 'namespace key old new') 936 @wireprotocommand('pushkey', 'namespace key old new')
923 def pushkey(repo, proto, namespace, key, old, new): 937 def pushkey(repo, proto, namespace, key, old, new):
924 # compatibility with pre-1.8 clients which were accidentally 938 # compatibility with pre-1.8 clients which were accidentally
925 # sending raw binary nodes rather than utf-8-encoded hex 939 # sending raw binary nodes rather than utf-8-encoded hex
936 with proto.mayberedirectstdio() as output: 950 with proto.mayberedirectstdio() as output:
937 r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key), 951 r = repo.pushkey(encoding.tolocal(namespace), encoding.tolocal(key),
938 encoding.tolocal(old), new) or False 952 encoding.tolocal(old), new) or False
939 953
940 output = output.getvalue() if output else '' 954 output = output.getvalue() if output else ''
941 return '%s\n%s' % (int(r), output) 955 return bytesresponse('%s\n%s' % (int(r), output))
942 956
943 @wireprotocommand('stream_out') 957 @wireprotocommand('stream_out')
944 def stream(repo, proto): 958 def stream(repo, proto):
945 '''If the server supports streaming clone, it advertises the "stream" 959 '''If the server supports streaming clone, it advertises the "stream"
946 capability with a value representing the version and flags of the repo 960 capability with a value representing the version and flags of the repo