|
1 # changegroup.py - Mercurial changegroup manipulation functions |
|
2 # |
|
3 # Copyright 2006 Matt Mackall <mpm@selenic.com> |
|
4 # |
|
5 # This software may be used and distributed according to the terms of the |
|
6 # GNU General Public License version 2 or any later version. |
|
7 |
|
8 from i18n import _ |
|
9 import util |
|
10 import struct, os, bz2, zlib, tempfile |
|
11 |
|
12 def getchunk(source): |
|
13 """return the next chunk from changegroup 'source' as a string""" |
|
14 d = source.read(4) |
|
15 if not d: |
|
16 return "" |
|
17 l = struct.unpack(">l", d)[0] |
|
18 if l <= 4: |
|
19 return "" |
|
20 d = source.read(l - 4) |
|
21 if len(d) < l - 4: |
|
22 raise util.Abort(_("premature EOF reading chunk" |
|
23 " (got %d bytes, expected %d)") |
|
24 % (len(d), l - 4)) |
|
25 return d |
|
26 |
|
27 def chunkheader(length): |
|
28 """return a changegroup chunk header (string)""" |
|
29 return struct.pack(">l", length + 4) |
|
30 |
|
31 def closechunk(): |
|
32 """return a changegroup chunk header (string) for a zero-length chunk""" |
|
33 return struct.pack(">l", 0) |
|
34 |
|
35 class nocompress(object): |
|
36 def compress(self, x): |
|
37 return x |
|
38 def flush(self): |
|
39 return "" |
|
40 |
|
41 bundletypes = { |
|
42 "": ("", nocompress), |
|
43 "HG10UN": ("HG10UN", nocompress), |
|
44 "HG10BZ": ("HG10", lambda: bz2.BZ2Compressor()), |
|
45 "HG10GZ": ("HG10GZ", lambda: zlib.compressobj()), |
|
46 } |
|
47 |
|
48 def collector(cl, mmfs, files): |
|
49 # Gather information about changeset nodes going out in a bundle. |
|
50 # We want to gather manifests needed and filelogs affected. |
|
51 def collect(node): |
|
52 c = cl.read(node) |
|
53 files.update(c[3]) |
|
54 mmfs.setdefault(c[0], node) |
|
55 return collect |
|
56 |
|
57 # hgweb uses this list to communicate its preferred type |
|
58 bundlepriority = ['HG10GZ', 'HG10BZ', 'HG10UN'] |
|
59 |
|
60 def writebundle(cg, filename, bundletype): |
|
61 """Write a bundle file and return its filename. |
|
62 |
|
63 Existing files will not be overwritten. |
|
64 If no filename is specified, a temporary file is created. |
|
65 bz2 compression can be turned off. |
|
66 The bundle file will be deleted in case of errors. |
|
67 """ |
|
68 |
|
69 fh = None |
|
70 cleanup = None |
|
71 try: |
|
72 if filename: |
|
73 fh = open(filename, "wb") |
|
74 else: |
|
75 fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg") |
|
76 fh = os.fdopen(fd, "wb") |
|
77 cleanup = filename |
|
78 |
|
79 header, compressor = bundletypes[bundletype] |
|
80 fh.write(header) |
|
81 z = compressor() |
|
82 |
|
83 # parse the changegroup data, otherwise we will block |
|
84 # in case of sshrepo because we don't know the end of the stream |
|
85 |
|
86 # an empty chunkgroup is the end of the changegroup |
|
87 # a changegroup has at least 2 chunkgroups (changelog and manifest). |
|
88 # after that, an empty chunkgroup is the end of the changegroup |
|
89 empty = False |
|
90 count = 0 |
|
91 while not empty or count <= 2: |
|
92 empty = True |
|
93 count += 1 |
|
94 while 1: |
|
95 chunk = getchunk(cg) |
|
96 if not chunk: |
|
97 break |
|
98 empty = False |
|
99 fh.write(z.compress(chunkheader(len(chunk)))) |
|
100 pos = 0 |
|
101 while pos < len(chunk): |
|
102 next = pos + 2**20 |
|
103 fh.write(z.compress(chunk[pos:next])) |
|
104 pos = next |
|
105 fh.write(z.compress(closechunk())) |
|
106 fh.write(z.flush()) |
|
107 cleanup = None |
|
108 return filename |
|
109 finally: |
|
110 if fh is not None: |
|
111 fh.close() |
|
112 if cleanup is not None: |
|
113 os.unlink(cleanup) |
|
114 |
|
115 def decompressor(fh, alg): |
|
116 if alg == 'UN': |
|
117 return fh |
|
118 elif alg == 'GZ': |
|
119 def generator(f): |
|
120 zd = zlib.decompressobj() |
|
121 for chunk in f: |
|
122 yield zd.decompress(chunk) |
|
123 elif alg == 'BZ': |
|
124 def generator(f): |
|
125 zd = bz2.BZ2Decompressor() |
|
126 zd.decompress("BZ") |
|
127 for chunk in util.filechunkiter(f, 4096): |
|
128 yield zd.decompress(chunk) |
|
129 else: |
|
130 raise util.Abort("unknown bundle compression '%s'" % alg) |
|
131 return util.chunkbuffer(generator(fh)) |
|
132 |
|
133 class unbundle10(object): |
|
134 def __init__(self, fh, alg): |
|
135 self._stream = decompressor(fh, alg) |
|
136 self._type = alg |
|
137 self.callback = None |
|
138 def compressed(self): |
|
139 return self._type != 'UN' |
|
140 def read(self, l): |
|
141 return self._stream.read(l) |
|
142 def seek(self, pos): |
|
143 return self._stream.seek(pos) |
|
144 def tell(self): |
|
145 return self._stream.tell() |
|
146 def close(self): |
|
147 return self._stream.close() |
|
148 |
|
149 def chunklength(self): |
|
150 d = self.read(4) |
|
151 if not d: |
|
152 return 0 |
|
153 l = max(0, struct.unpack(">l", d)[0] - 4) |
|
154 if l and self.callback: |
|
155 self.callback() |
|
156 return l |
|
157 |
|
158 def chunk(self): |
|
159 """return the next chunk from changegroup 'source' as a string""" |
|
160 l = self.chunklength() |
|
161 d = self.read(l) |
|
162 if len(d) < l: |
|
163 raise util.Abort(_("premature EOF reading chunk" |
|
164 " (got %d bytes, expected %d)") |
|
165 % (len(d), l)) |
|
166 return d |
|
167 |
|
168 def parsechunk(self): |
|
169 l = self.chunklength() |
|
170 if not l: |
|
171 return {} |
|
172 h = self.read(80) |
|
173 node, p1, p2, cs = struct.unpack("20s20s20s20s", h) |
|
174 data = self.read(l - 80) |
|
175 return dict(node=node, p1=p1, p2=p2, cs=cs, data=data) |
|
176 |
|
177 class headerlessfixup(object): |
|
178 def __init__(self, fh, h): |
|
179 self._h = h |
|
180 self._fh = fh |
|
181 def read(self, n): |
|
182 if self._h: |
|
183 d, self._h = self._h[:n], self._h[n:] |
|
184 if len(d) < n: |
|
185 d += self._fh.read(n - len(d)) |
|
186 return d |
|
187 return self._fh.read(n) |
|
188 |
|
189 def readbundle(fh, fname): |
|
190 header = fh.read(6) |
|
191 |
|
192 if not fname: |
|
193 fname = "stream" |
|
194 if not header.startswith('HG') and header.startswith('\0'): |
|
195 fh = headerlessfixup(fh, header) |
|
196 header = "HG10UN" |
|
197 |
|
198 magic, version, alg = header[0:2], header[2:4], header[4:6] |
|
199 |
|
200 if magic != 'HG': |
|
201 raise util.Abort(_('%s: not a Mercurial bundle') % fname) |
|
202 if version != '10': |
|
203 raise util.Abort(_('%s: unknown bundle version %s') % (fname, version)) |
|
204 return unbundle10(fh, alg) |