def storetobundlestore(orig, repo, op, unbundler): """stores the incoming bundle coming from push command to the bundlestore instead of applying on the revlogs""" repo.ui.status(_(b"storing changesets on the bundlestore\n")) bundler = bundle2.bundle20(repo.ui) # processing each part and storing it in bundler with bundle2.partiterator(repo, op, unbundler) as parts: for part in parts: bundlepart = None if part.type == b'replycaps': # This configures the current operation to allow reply parts. bundle2._processpart(op, part) else: bundlepart = bundle2.bundlepart(part.type, data=part.read()) for key, value in pycompat.iteritems(part.params): bundlepart.addparam(key, value) # Certain parts require a response if part.type in (b'pushkey', b'changegroup'): if op.reply is not None: rpart = op.reply.newpart(b'reply:%s' % part.type) rpart.addparam( b'in-reply-to', b'%d' % part.id, mandatory=False ) rpart.addparam(b'return', b'1', mandatory=False) op.records.add( part.type, { b'return': 1, }, ) if bundlepart: bundler.addpart(bundlepart) # storing the bundle in the bundlestore buf = util.chunkbuffer(bundler.getchunks()) fd, bundlefile = pycompat.mkstemp() try: try: fp = os.fdopen(fd, 'wb') fp.write(buf.read()) finally: fp.close() storebundle(op, {}, bundlefile) finally: try: os.unlink(bundlefile) except Exception: # we would rather see the original exception pass
def processparts(orig, repo, op, unbundler): # make sure we don't wrap processparts in case of `hg unbundle` if op.source == 'unbundle': return orig(repo, op, unbundler) # this server routes each push to bundle store if repo.ui.configbool('infinitepush', 'pushtobundlestore'): return storetobundlestore(orig, repo, op, unbundler) if unbundler.params.get('infinitepush') != 'True': return orig(repo, op, unbundler) handleallparts = repo.ui.configbool('infinitepush', 'storeallparts') bundler = bundle2.bundle20(repo.ui) cgparams = None with bundle2.partiterator(repo, op, unbundler) as parts: for part in parts: bundlepart = None if part.type == 'replycaps': # This configures the current operation to allow reply parts. bundle2._processpart(op, part) elif part.type == bundleparts.scratchbranchparttype: # Scratch branch parts need to be converted to normal # changegroup parts, and the extra parameters stored for later # when we upload to the store. Eventually those parameters will # be put on the actual bundle instead of this part, then we can # send a vanilla changegroup instead of the scratchbranch part. cgversion = part.params.get('cgversion', '01') bundlepart = bundle2.bundlepart('changegroup', data=part.read()) bundlepart.addparam('version', cgversion) cgparams = part.params # If we're not dumping all parts into the new bundle, we need to # alert the future pushkey and phase-heads handler to skip # the part. if not handleallparts: op.records.add(scratchbranchparttype + '_skippushkey', True) op.records.add(scratchbranchparttype + '_skipphaseheads', True) else: if handleallparts: # Ideally we would not process any parts, and instead just # forward them to the bundle for storage, but since this # differs from previous behavior, we need to put it behind a # config flag for incremental rollout. bundlepart = bundle2.bundlepart(part.type, data=part.read()) for key, value in part.params.iteritems(): bundlepart.addparam(key, value) # Certain parts require a response if part.type == 'pushkey': if op.reply is not None: rpart = op.reply.newpart('reply:pushkey') rpart.addparam('in-reply-to', str(part.id), mandatory=False) rpart.addparam('return', '1', mandatory=False) else: bundle2._processpart(op, part) if handleallparts: op.records.add(part.type, { 'return': 1, }) if bundlepart: bundler.addpart(bundlepart) # If commits were sent, store them if cgparams: buf = util.chunkbuffer(bundler.getchunks()) fd, bundlefile = tempfile.mkstemp() try: try: fp = os.fdopen(fd, r'wb') fp.write(buf.read()) finally: fp.close() storebundle(op, cgparams, bundlefile) finally: try: os.unlink(bundlefile) except Exception: # we would rather see the original exception pass