diff --git a/hgitaly/service/repository.py b/hgitaly/service/repository.py index d7138f4d400c1dcd75227c7a50b20345263600dd_aGdpdGFseS9zZXJ2aWNlL3JlcG9zaXRvcnkucHk=..832cbfae47cdbcd8c24b74bb8527596f4bbf5ef3_aGdpdGFseS9zZXJ2aWNlL3JlcG9zaXRvcnkucHk= 100644 --- a/hgitaly/service/repository.py +++ b/hgitaly/service/repository.py @@ -92,7 +92,10 @@ ) from ..servicer import HGitalyServicer -from ..stream import WRITE_BUFFER_SIZE +from ..stream import ( + WRITE_BUFFER_SIZE, + streaming_request_tempfile_extract, +) from ..path import ( InvalidPath, validate_relative_path, @@ -445,4 +448,8 @@ sub-request contains repository+data and subsequent sub-requests contains only data (i.e. the actual bundle sent in parts). """ + def init_repo(req, context): + self.hg_init_repository(req.repository, context) + return self.load_repo(req.repository, context) + try: @@ -448,15 +455,8 @@ try: - with tempfile.NamedTemporaryFile( - mode='wb+', - buffering=WRITE_BUFFER_SIZE) as tmpf: - first_request = True - for req in request: - if first_request: - self.hg_init_repository(req.repository, context) - repo = self.load_repo(req.repository, context) - first_request = False - tmpf.write(req.data) - tmpf.seek(0) + with streaming_request_tempfile_extract( + request, context, + first_request_handler=init_repo + ) as (repo, tmpf): unbundle(repo, tmpf.name) finally: # TODO in case of error, we must remove the created repo, diff --git a/hgitaly/stream.py b/hgitaly/stream.py index d7138f4d400c1dcd75227c7a50b20345263600dd_aGdpdGFseS9zdHJlYW0ucHk=..832cbfae47cdbcd8c24b74bb8527596f4bbf5ef3_aGdpdGFseS9zdHJlYW0ucHk= 100644 --- a/hgitaly/stream.py +++ b/hgitaly/stream.py @@ -21,4 +21,5 @@ # and 10 otherwise. Also, for argument base 0 only, # underscore characters are permitted as defined by the # Go syntax for integer literals. +import contextlib import os @@ -24,4 +25,5 @@ import os +import tempfile def concat_resplit(in_chunks, out_chunks_size): @@ -115,3 +117,33 @@ // the GITALY_STREAMIO_WRITE_BUFFER_SIZE environment variable. var WriteBufferSize = 128 * 1024 """ + + +@contextlib.contextmanager +def streaming_request_tempfile_extract(requests, context, + data_attr='data', + first_request_handler=None): + """Extract data from a streaming request. + + It is a a common Gitaly pattern for streaming requests with binary + data to bear other information in the first request message only. + + This context manager extracts all the binary data in a + :class:`NamedTemporaryFile` and calls `first_request_handler` to + perform actions from the first request message. + + The provided Python context is the (first_result, temporary file) pair, + where `first_result` is the return value of `first_request_handler`. + """ + with tempfile.NamedTemporaryFile(mode='wb+', + buffering=WRITE_BUFFER_SIZE) as tmpf: + first_request = True + first_result = None + for req in requests: + if first_request and first_request_handler is not None: + first_result = first_request_handler(req, context) + first_request = False + tmpf.write(req.data) + tmpf.seek(0) + + yield first_result, tmpf