diff --git a/hgitaly/service/repository_service.py b/hgitaly/service/repository_service.py index 510895fa5f79fd2e69dffb5efb227e4fc2b17c77_aGdpdGFseS9zZXJ2aWNlL3JlcG9zaXRvcnlfc2VydmljZS5weQ==..fe6737d193523df186a66f720e5c0fb06b71203c_aGdpdGFseS9zZXJ2aWNlL3JlcG9zaXRvcnlfc2VydmljZS5weQ== 100644 --- a/hgitaly/service/repository_service.py +++ b/hgitaly/service/repository_service.py @@ -6,4 +6,6 @@ # SPDX-License-Identifier: GPL-2.0-or-later from grpc import StatusCode import logging +import tempfile + from mercurial import ( @@ -9,4 +11,6 @@ from mercurial import ( + archival, + scmutil, error, ) @@ -19,6 +23,8 @@ from ..stub.repository_service_pb2 import ( RepositoryExistsRequest, RepositoryExistsResponse, + GetArchiveRequest, + GetArchiveResponse, HasLocalBranchesRequest, HasLocalBranchesResponse, WriteRefRequest, @@ -26,6 +32,8 @@ ) from ..stub.repository_service_pb2_grpc import RepositoryServiceServicer from ..servicer import HGitalyServicer +from ..stream import WRITE_BUFFER_SIZE + logger = logging.getLogger(__name__) DEFAULT_BRANCH_FILE_NAME = b'default_gitlab_branch' @@ -29,6 +37,12 @@ logger = logging.getLogger(__name__) DEFAULT_BRANCH_FILE_NAME = b'default_gitlab_branch' +ARCHIVE_FORMATS = { + GetArchiveRequest.Format.Value('ZIP'): b'zip', + GetArchiveRequest.Format.Value('TAR'): b'tar', + GetArchiveRequest.Format.Value('TAR_GZ'): b'tgz', + GetArchiveRequest.Format.Value('TAR_BZ2'): b'tbz2', +} class RepositoryServiceServicer(RepositoryServiceServicer, HGitalyServicer): @@ -46,6 +60,39 @@ return RepositoryExistsResponse(exists=exists) + def GetArchive(self, + request: GetArchiveRequest, + context) -> GetArchiveResponse: + logger.debug("GetArchive request=%r", request) + repo = self.load_repo(request.repository, context) + ctx = repo[request.commit_id] + # TODO translate 'path' as an include rule + match = scmutil.match(ctx, [], {}) + + # using an anonymous (not linked) temporary file + # TODO OPTIM check if archive is not by any chance + # using a tempfile already… + with tempfile.TemporaryFile( + mode='wb+', # the default, but let's insist on binary here + buffering=WRITE_BUFFER_SIZE) as tmpf: + archival.archive( + repo, + tmpf, + ctx.node(), + ARCHIVE_FORMATS[request.format], + True, # decode (TODO this is the default but what is this?) + match, + request.prefix.encode(), + subrepos=False # maybe later, check what GitLab's standard is + ) + + tmpf.seek(0) + while True: + data = tmpf.read(WRITE_BUFFER_SIZE) + if not data: + break + yield GetArchiveResponse(data=data) + def HasLocalBranches(self, request: HasLocalBranchesRequest, context) -> HasLocalBranchesResponse: diff --git a/hgitaly/service/tests/test_repository_service.py b/hgitaly/service/tests/test_repository_service.py index 510895fa5f79fd2e69dffb5efb227e4fc2b17c77_aGdpdGFseS9zZXJ2aWNlL3Rlc3RzL3Rlc3RfcmVwb3NpdG9yeV9zZXJ2aWNlLnB5..fe6737d193523df186a66f720e5c0fb06b71203c_aGdpdGFseS9zZXJ2aWNlL3Rlc3RzL3Rlc3RfcmVwb3NpdG9yeV9zZXJ2aWNlLnB5 100644 --- a/hgitaly/service/tests/test_repository_service.py +++ b/hgitaly/service/tests/test_repository_service.py @@ -4,5 +4,7 @@ # GNU General Public License version 2 or any later version. # # SPDX-License-Identifier: GPL-2.0-or-later +from contextlib import contextmanager +from io import BytesIO import grpc import shutil @@ -7,7 +9,9 @@ import grpc import shutil +import tarfile + import pytest from hgitaly.tests.common import make_empty_repo from hgitaly.stub.repository_service_pb2 import ( @@ -9,8 +13,9 @@ import pytest from hgitaly.tests.common import make_empty_repo from hgitaly.stub.repository_service_pb2 import ( + GetArchiveRequest, HasLocalBranchesRequest, RepositoryExistsRequest, WriteRefRequest, @@ -72,3 +77,86 @@ ref=b'HEAD', revision=b'topic/default/wont-last')) assert exc_info.value.code() == grpc.StatusCode.INVALID_ARGUMENT + + +@contextmanager +def get_archive_tarfile(stub, grpc_repo, commit_id): + with BytesIO() as fobj: + for chunk_index, chunk_response in enumerate( + stub.GetArchive(GetArchiveRequest( + repository=grpc_repo, + format=GetArchiveRequest.Format.Value('TAR'), + commit_id=commit_id, + prefix='archive-dir', + ))): + fobj.write(chunk_response.data) + + fobj.seek(0) + with tarfile.open(fileobj=fobj) as tarf: + yield tarf, chunk_index + 1 + + +def test_get_archive(grpc_channel, server_repos_root, tmpdir): + repo_stub = RepositoryServiceStub(grpc_channel) + wrapper, grpc_repo = make_empty_repo(server_repos_root) + + ctx = wrapper.write_commit('foo', content="Foo") + (wrapper.path / 'sub').mkdir() + ctx2 = wrapper.write_commit('sub/bar', content="Bar") + + node_str = ctx.hex().decode() + with get_archive_tarfile(repo_stub, grpc_repo, node_str) as (tarf, _nb): + assert set(tarf.getnames()) == {'archive-dir/.hg_archival.txt', + 'archive-dir/foo'} + + extract_dir = tmpdir.join('extract') + tarf.extractall(path=extract_dir) + + metadata_lines = extract_dir.join('archive-dir', + '.hg_archival.txt').readlines() + + assert 'node: %s\n' % node_str in metadata_lines + assert extract_dir.join('archive-dir', 'foo').read() == "Foo" + + node2_str = ctx2.hex().decode() + with get_archive_tarfile(repo_stub, grpc_repo, node2_str) as (tarf, _nb): + assert set(tarf.getnames()) == {'archive-dir/.hg_archival.txt', + 'archive-dir/foo', + 'archive-dir/sub/bar'} + + extract_dir = tmpdir.join('extract-2') + tarf.extractall(path=extract_dir) + + metadata_lines = extract_dir.join('archive-dir', + '.hg_archival.txt').readlines() + + assert 'node: %s\n' % node2_str in metadata_lines + assert extract_dir.join('archive-dir', 'sub', 'bar').read() == "Bar" + + +def test_get_archive_multiple_chunks(grpc_channel, server_repos_root, + tmpdir, monkeypatch): + + repo_stub = RepositoryServiceStub(grpc_channel) + wrapper, grpc_repo = make_empty_repo(server_repos_root) + + # we can't just override the environment variable: it's read at module + # import time. + large_content = "Foo" * 200000 + ctx = wrapper.write_commit('foo', content=large_content) + node_str = ctx.hex().decode() + with get_archive_tarfile(repo_stub, grpc_repo, node_str) as (tarf, count): + assert count > 1 + assert set(tarf.getnames()) == {'archive-dir/.hg_archival.txt', + 'archive-dir/foo'} + + extract_dir = tmpdir.join('extract') + tarf.extractall(path=extract_dir) + + metadata_lines = extract_dir.join('archive-dir', + '.hg_archival.txt').readlines() + + assert 'node: %s\n' % node_str in metadata_lines + assert extract_dir.join('archive-dir', 'foo').read() == large_content + + del large_content