mercurial/wireprotov2server.py

--- a/mercurial/wireprotov2server.py	Tue Oct 16 21:35:33 2018 +0200
+++ b/mercurial/wireprotov2server.py	Tue Oct 16 21:31:21 2018 +0200
@@ -22,6 +22,8 @@
     match as matchmod,
     narrowspec,
     pycompat,
+    streamclone,
+    util,
     wireprotoframing,
     wireprototypes,
 )
@@ -515,6 +517,11 @@
             if meta['validvalues']:
                 args[arg][b'validvalues'] = meta['validvalues']
 
+        # TODO this type of check should be defined in a per-command callback.
+        if (command == b'rawstorefiledata'
+            and not streamclone.allowservergeneration(repo)):
+            continue
+
         caps['commands'][command] = {
             'args': args,
             'permissions': [entry.permission],
@@ -1369,3 +1376,73 @@
                        encoding.tolocal(key),
                        encoding.tolocal(old),
                        encoding.tolocal(new))
+
+
+@wireprotocommand(
+    'rawstorefiledata',
+    args={
+        'files': {
+            'type': 'list',
+            'example': [b'changelog', b'manifestlog'],
+        },
+        'pathfilter': {
+            'type': 'list',
+            'default': lambda: None,
+            'example': {b'include': [b'path:tests']},
+        },
+    },
+    permission='pull')
+def rawstorefiledata(repo, proto, files, pathfilter):
+    if not streamclone.allowservergeneration(repo):
+        raise error.WireprotoCommandError(b'stream clone is disabled')
+
+    # TODO support dynamically advertising what store files "sets" are
+    # available. For now, we support changelog, manifestlog, and files.
+    files = set(files)
+    allowedfiles = {b'changelog', b'manifestlog'}
+
+    unsupported = files - allowedfiles
+    if unsupported:
+        raise error.WireprotoCommandError(b'unknown file type: %s',
+                                          (b', '.join(sorted(unsupported)),))
+
+    with repo.lock():
+        topfiles = list(repo.store.topfiles())
+
+    sendfiles = []
+    totalsize = 0
+
+    # TODO this is a bunch of storage layer interface abstractions because
+    # it assumes revlogs.
+    for name, encodedname, size in topfiles:
+        if b'changelog' in files and name.startswith(b'00changelog'):
+            pass
+        elif b'manifestlog' in files and name.startswith(b'00manifest'):
+            pass
+        else:
+            continue
+
+        sendfiles.append((b'store', name, size))
+        totalsize += size
+
+    yield {
+        b'filecount': len(sendfiles),
+        b'totalsize': totalsize,
+    }
+
+    for location, name, size in sendfiles:
+        yield {
+            b'location': location,
+            b'path': name,
+            b'size': size,
+        }
+
+        # We have to use a closure for this to ensure the context manager is
+        # closed only after sending the final chunk.
+        def getfiledata():
+            with repo.svfs(name, 'rb', auditpath=False) as fh:
+                for chunk in util.filechunkiter(fh, limit=size):
+                    yield chunk
+
+        yield wireprototypes.indefinitebytestringresponse(
+            getfiledata())