Skip to content

Commit b644721

Browse files
TCatshoekTCatshoek
authored
feat(downloads): allow streaming downloads access to response iterator (#1956)
* feat(downloads): allow streaming downloads access to response iterator Allow access to the underlying response iterator when downloading in streaming mode by specifying `iterator=True`. Update type annotations to support this change. * docs(api-docs): add iterator example to artifact download Document the usage of the `iterator=True` option when downloading artifacts * test(packages): add tests for streaming downloads
1 parent 0f2a602 commit b644721

File tree

13 files changed

+165
-37
lines changed

13 files changed

+165
-37
lines changed

‎docs/gl_objects/pipelines_and_jobs.rst‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,19 @@ You can also directly stream the output into a file, and unzip it afterwards::
274274
subprocess.run(["unzip", "-bo", zipfn])
275275
os.unlink(zipfn)
276276

277+
Or, you can also use the underlying response iterator directly::
278+
279+
artifact_bytes_iterator = build_or_job.artifacts(iterator=True)
280+
281+
This can be used with frameworks that expect an iterator (such as FastAPI/Starlette's
282+
``StreamingResponse``) to forward a download from GitLab without having to download
283+
the entire content server-side first::
284+
285+
@app.get("/download_artifact")
286+
def download_artifact():
287+
artifact_bytes_iterator = build_or_job.artifacts(iterator=True)
288+
return StreamingResponse(artifact_bytes_iterator, media_type="application/zip")
289+
277290
Delete all artifacts of a project that can be deleted::
278291

279292
project.artifacts.delete()

‎gitlab/mixins.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
Any,
2121
Callable,
2222
Dict,
23+
Iterator,
2324
List,
2425
Optional,
2526
Tuple,
@@ -612,16 +613,19 @@ class DownloadMixin(_RestObjectBase):
612613
def download(
613614
self,
614615
streamed: bool = False,
616+
iterator: bool = False,
615617
action: Optional[Callable] = None,
616618
chunk_size: int = 1024,
617619
**kwargs: Any,
618-
) -> Optional[bytes]:
620+
) -> Optional[Union[bytes, Iterator[Any]]]:
619621
"""Download the archive of a resource export.
620622
621623
Args:
622624
streamed: If True the data will be processed by chunks of
623625
`chunk_size` and each chunk is passed to `action` for
624626
treatment
627+
iterator: If True directly return the underlying response
628+
iterator
625629
action: Callable responsible of dealing with chunk of
626630
data
627631
chunk_size: Size of each chunk
@@ -640,7 +644,7 @@ def download(
640644
)
641645
if TYPE_CHECKING:
642646
assert isinstance(result, requests.Response)
643-
return utils.response_content(result, streamed, action, chunk_size)
647+
return utils.response_content(result, streamed, iterator, action, chunk_size)
644648

645649

646650
class SubscribableMixin(_RestObjectBase):

‎gitlab/utils.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import traceback
2020
import urllib.parse
2121
import warnings
22-
from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
22+
from typing import Any, Callable, Dict, Iterator, Optional, Tuple, Type, Union
2323

2424
import requests
2525

@@ -34,9 +34,13 @@ def __call__(self, chunk: Any) -> None:
3434
def response_content(
3535
response: requests.Response,
3636
streamed: bool,
37+
iterator: bool,
3738
action: Optional[Callable],
3839
chunk_size: int,
39-
) -> Optional[bytes]:
40+
) -> Optional[Union[bytes, Iterator[Any]]]:
41+
if iterator:
42+
return response.iter_content(chunk_size=chunk_size)
43+
4044
if streamed is False:
4145
return response.content
4246

‎gitlab/v4/cli.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ def do_project_export_download(self) -> None:
127127
data = export_status.download()
128128
if TYPE_CHECKING:
129129
assert data is not None
130+
assert isinstance(data, bytes)
130131
sys.stdout.buffer.write(data)
131132

132133
except Exception as e: # pragma: no cover, cli.die is unit-tested

‎gitlab/v4/objects/artifacts.py‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
GitLab API:
33
https://docs.gitlab.com/ee/api/job_artifacts.html
44
"""
5-
from typing import Any, Callable, Optional, TYPE_CHECKING
5+
from typing import Any, Callable, Iterator, Optional, TYPE_CHECKING, Union
66

77
import requests
88

@@ -40,10 +40,14 @@ def __call__(
4040
),
4141
category=DeprecationWarning,
4242
)
43-
return self.download(
43+
data = self.download(
4444
*args,
4545
**kwargs,
4646
)
47+
if TYPE_CHECKING:
48+
assert data is not None
49+
assert isinstance(data, bytes)
50+
return data
4751

4852
@exc.on_http_error(exc.GitlabDeleteError)
4953
def delete(self, **kwargs: Any) -> None:
@@ -71,10 +75,11 @@ def download(
7175
ref_name: str,
7276
job: str,
7377
streamed: bool = False,
78+
iterator: bool = False,
7479
action: Optional[Callable] = None,
7580
chunk_size: int = 1024,
7681
**kwargs: Any,
77-
) -> Optional[bytes]:
82+
) -> Optional[Union[bytes, Iterator[Any]]]:
7883
"""Get the job artifacts archive from a specific tag or branch.
7984
8085
Args:
@@ -85,6 +90,8 @@ def download(
8590
streamed: If True the data will be processed by chunks of
8691
`chunk_size` and each chunk is passed to `action` for
8792
treatment
93+
iterator: If True directly return the underlying response
94+
iterator
8895
action: Callable responsible of dealing with chunk of
8996
data
9097
chunk_size: Size of each chunk
@@ -103,7 +110,7 @@ def download(
103110
)
104111
if TYPE_CHECKING:
105112
assert isinstance(result, requests.Response)
106-
return utils.response_content(result, streamed, action, chunk_size)
113+
return utils.response_content(result, streamed, iterator, action, chunk_size)
107114

108115
@cli.register_custom_action(
109116
"ProjectArtifactManager", ("ref_name", "artifact_path", "job")
@@ -115,10 +122,11 @@ def raw(
115122
artifact_path: str,
116123
job: str,
117124
streamed: bool = False,
125+
iterator: bool = False,
118126
action: Optional[Callable] = None,
119127
chunk_size: int = 1024,
120128
**kwargs: Any,
121-
) -> Optional[bytes]:
129+
) -> Optional[Union[bytes, Iterator[Any]]]:
122130
"""Download a single artifact file from a specific tag or branch from
123131
within the job's artifacts archive.
124132
@@ -130,6 +138,8 @@ def raw(
130138
streamed: If True the data will be processed by chunks of
131139
`chunk_size` and each chunk is passed to `action` for
132140
treatment
141+
iterator: If True directly return the underlying response
142+
iterator
133143
action: Callable responsible of dealing with chunk of
134144
data
135145
chunk_size: Size of each chunk
@@ -148,4 +158,4 @@ def raw(
148158
)
149159
if TYPE_CHECKING:
150160
assert isinstance(result, requests.Response)
151-
return utils.response_content(result, streamed, action, chunk_size)
161+
return utils.response_content(result, streamed, iterator, action, chunk_size)

‎gitlab/v4/objects/files.py‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,15 @@
11
import base64
2-
from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING
2+
from typing import (
3+
Any,
4+
Callable,
5+
cast,
6+
Dict,
7+
Iterator,
8+
List,
9+
Optional,
10+
TYPE_CHECKING,
11+
Union,
12+
)
313

414
import requests
515

@@ -220,10 +230,11 @@ def raw(
220230
file_path: str,
221231
ref: str,
222232
streamed: bool = False,
233+
iterator: bool = False,
223234
action: Optional[Callable[..., Any]] = None,
224235
chunk_size: int = 1024,
225236
**kwargs: Any,
226-
) -> Optional[bytes]:
237+
) -> Optional[Union[bytes, Iterator[Any]]]:
227238
"""Return the content of a file for a commit.
228239
229240
Args:
@@ -232,6 +243,8 @@ def raw(
232243
streamed: If True the data will be processed by chunks of
233244
`chunk_size` and each chunk is passed to `action` for
234245
treatment
246+
iterator: If True directly return the underlying response
247+
iterator
235248
action: Callable responsible of dealing with chunk of
236249
data
237250
chunk_size: Size of each chunk
@@ -252,7 +265,7 @@ def raw(
252265
)
253266
if TYPE_CHECKING:
254267
assert isinstance(result, requests.Response)
255-
return utils.response_content(result, streamed, action, chunk_size)
268+
return utils.response_content(result, streamed, iterator, action, chunk_size)
256269

257270
@cli.register_custom_action("ProjectFileManager", ("file_path", "ref"))
258271
@exc.on_http_error(exc.GitlabListError)

‎gitlab/v4/objects/jobs.py‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Any, Callable, cast, Dict, Optional, TYPE_CHECKING, Union
1+
from typing import Any, Callable, cast, Dict, Iterator, Optional, TYPE_CHECKING, Union
22

33
import requests
44

@@ -116,16 +116,19 @@ def delete_artifacts(self, **kwargs: Any) -> None:
116116
def artifacts(
117117
self,
118118
streamed: bool = False,
119+
iterator: bool = False,
119120
action: Optional[Callable[..., Any]] = None,
120121
chunk_size: int = 1024,
121122
**kwargs: Any,
122-
) -> Optional[bytes]:
123+
) -> Optional[Union[bytes, Iterator[Any]]]:
123124
"""Get the job artifacts.
124125
125126
Args:
126127
streamed: If True the data will be processed by chunks of
127128
`chunk_size` and each chunk is passed to `action` for
128129
treatment
130+
iterator: If True directly return the underlying response
131+
iterator
129132
action: Callable responsible of dealing with chunk of
130133
data
131134
chunk_size: Size of each chunk
@@ -144,25 +147,28 @@ def artifacts(
144147
)
145148
if TYPE_CHECKING:
146149
assert isinstance(result, requests.Response)
147-
return utils.response_content(result, streamed, action, chunk_size)
150+
return utils.response_content(result, streamed, iterator, action, chunk_size)
148151

149152
@cli.register_custom_action("ProjectJob")
150153
@exc.on_http_error(exc.GitlabGetError)
151154
def artifact(
152155
self,
153156
path: str,
154157
streamed: bool = False,
158+
iterator: bool = False,
155159
action: Optional[Callable[..., Any]] = None,
156160
chunk_size: int = 1024,
157161
**kwargs: Any,
158-
) -> Optional[bytes]:
162+
) -> Optional[Union[bytes, Iterator[Any]]]:
159163
"""Get a single artifact file from within the job's artifacts archive.
160164
161165
Args:
162166
path: Path of the artifact
163167
streamed: If True the data will be processed by chunks of
164168
`chunk_size` and each chunk is passed to `action` for
165169
treatment
170+
iterator: If True directly return the underlying response
171+
iterator
166172
action: Callable responsible of dealing with chunk of
167173
data
168174
chunk_size: Size of each chunk
@@ -181,13 +187,14 @@ def artifact(
181187
)
182188
if TYPE_CHECKING:
183189
assert isinstance(result, requests.Response)
184-
return utils.response_content(result, streamed, action, chunk_size)
190+
return utils.response_content(result, streamed, iterator, action, chunk_size)
185191

186192
@cli.register_custom_action("ProjectJob")
187193
@exc.on_http_error(exc.GitlabGetError)
188194
def trace(
189195
self,
190196
streamed: bool = False,
197+
iterator: bool = False,
191198
action: Optional[Callable[..., Any]] = None,
192199
chunk_size: int = 1024,
193200
**kwargs: Any,
@@ -198,6 +205,8 @@ def trace(
198205
streamed: If True the data will be processed by chunks of
199206
`chunk_size` and each chunk is passed to `action` for
200207
treatment
208+
iterator: If True directly return the underlying response
209+
iterator
201210
action: Callable responsible of dealing with chunk of
202211
data
203212
chunk_size: Size of each chunk
@@ -216,7 +225,9 @@ def trace(
216225
)
217226
if TYPE_CHECKING:
218227
assert isinstance(result, requests.Response)
219-
return_value = utils.response_content(result, streamed, action, chunk_size)
228+
return_value = utils.response_content(
229+
result, streamed, iterator, action, chunk_size
230+
)
220231
if TYPE_CHECKING:
221232
assert isinstance(return_value, dict)
222233
return return_value

‎gitlab/v4/objects/packages.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
"""
66

77
from pathlib import Path
8-
from typing import Any, Callable, cast, Optional, TYPE_CHECKING, Union
8+
from typing import Any, Callable, cast, Iterator, Optional, TYPE_CHECKING, Union
99

1010
import requests
1111

@@ -103,10 +103,11 @@ def download(
103103
package_version: str,
104104
file_name: str,
105105
streamed: bool = False,
106+
iterator: bool = False,
106107
action: Optional[Callable] = None,
107108
chunk_size: int = 1024,
108109
**kwargs: Any,
109-
) -> Optional[bytes]:
110+
) -> Optional[Union[bytes, Iterator[Any]]]:
110111
"""Download a generic package.
111112
112113
Args:
@@ -116,6 +117,8 @@ def download(
116117
streamed: If True the data will be processed by chunks of
117118
`chunk_size` and each chunk is passed to `action` for
118119
treatment
120+
iterator: If True directly return the underlying response
121+
iterator
119122
action: Callable responsible of dealing with chunk of
120123
data
121124
chunk_size: Size of each chunk
@@ -132,7 +135,7 @@ def download(
132135
result = self.gitlab.http_get(path, streamed=streamed, raw=True, **kwargs)
133136
if TYPE_CHECKING:
134137
assert isinstance(result, requests.Response)
135-
return utils.response_content(result, streamed, action, chunk_size)
138+
return utils.response_content(result, streamed, iterator, action, chunk_size)
136139

137140

138141
class GroupPackage(RESTObject):

0 commit comments

Comments
 (0)