|
9 | 9 | from packaging.version import Version |
10 | 10 | from functools import wraps |
11 | 11 | from xml.etree.ElementTree import ParseError |
| 12 | +from collections.abc import Iterator |
12 | 13 | from typing import ( |
13 | 14 | Any, |
14 | 15 | Callable, |
|
20 | 21 | ) |
21 | 22 | from typing_extensions import Self |
22 | 23 |
|
| 24 | +from tableauserverclient.config import BYTES_PER_MB, config |
23 | 25 | from tableauserverclient.models.pagination_item import PaginationItem |
24 | | -from tableauserverclient.server.request_options import RequestOptions |
| 26 | +from tableauserverclient.server.request_options import RequestOptions, RequestOptionsBase |
25 | 27 | from tableauserverclient.filesys_helpers import to_filename, make_download_path |
26 | 28 |
|
27 | 29 | from tableauserverclient.server.endpoint.exceptions import ( |
@@ -338,9 +340,13 @@ def wrapper(self: E, *args: P.args, **kwargs: P.kwargs) -> R: |
338 | 340 | class DownloadableMixin: |
339 | 341 | """Mixin for endpoints whose resources can be downloaded as binary files. |
340 | 342 |
|
341 | | - Provides a single private helper that streams a server response to a file |
342 | | - path or writable file object, avoiding copy-paste of the identical streaming |
343 | | - loop in Workbooks, Datasources, and Flows. |
| 343 | + Provides two private helpers, avoiding copy-paste of the streaming loop |
| 344 | + across endpoints: |
| 345 | +
|
| 346 | + - _download_content streams a response to a file path or writable object |
| 347 | + (used by Workbooks, Datasources, Flows). |
| 348 | + - _stream_content yields chunks for callers that want to materialize the |
| 349 | + body in-memory or process it lazily (used by Views, CustomViews). |
344 | 350 | """ |
345 | 351 |
|
346 | 352 | def _download_content( |
@@ -368,18 +374,36 @@ def _download_content( |
368 | 374 | m = Message() |
369 | 375 | m["Content-Disposition"] = server_response.headers["Content-Disposition"] |
370 | 376 | filename = m.get_filename(failobj="") |
| 377 | + chunk_size = config.CHUNK_SIZE_MB * BYTES_PER_MB |
371 | 378 | if isinstance(filepath, _io_types_w): |
372 | | - for chunk in server_response.iter_content(1024): # 1KB |
| 379 | + for chunk in server_response.iter_content(chunk_size): |
373 | 380 | filepath.write(chunk) |
374 | 381 | return filepath |
375 | 382 | else: |
376 | 383 | filename = to_filename(os.path.basename(filename)) |
377 | 384 | download_path = make_download_path(filepath, filename) |
378 | 385 | with open(download_path, "wb") as f: |
379 | | - for chunk in server_response.iter_content(1024): # 1KB |
| 386 | + for chunk in server_response.iter_content(chunk_size): |
380 | 387 | f.write(chunk) |
381 | 388 | return os.path.abspath(download_path) |
382 | 389 |
|
| 390 | + def _stream_content( |
| 391 | + self, |
| 392 | + url: str, |
| 393 | + request_object: RequestOptionsBase | None = None, |
| 394 | + ) -> Iterator[bytes]: |
| 395 | + """Stream content at url as chunks. |
| 396 | +
|
| 397 | + Suitable for callers that want to materialize the body in-memory |
| 398 | + (e.g. b"".join(iterator)) or process it lazily as a stream, rather than |
| 399 | + write it to a file. Complements _download_content, which writes to disk. |
| 400 | + """ |
| 401 | + chunk_size = config.CHUNK_SIZE_MB * BYTES_PER_MB |
| 402 | + with closing( |
| 403 | + self.get_request(url, request_object=request_object, parameters={"stream": True}) # type: ignore[attr-defined] |
| 404 | + ) as server_response: |
| 405 | + yield from server_response.iter_content(chunk_size) |
| 406 | + |
383 | 407 |
|
384 | 408 | class QuerysetEndpoint(Endpoint, Generic[T]): |
385 | 409 | @api(version="2.0") |
|
0 commit comments