|
1 | | -from email.message import Message |
2 | 1 | import copy |
3 | 2 | import json |
4 | 3 | import io |
5 | 4 | import os |
6 | 5 |
|
7 | | -from contextlib import closing |
8 | 6 | from pathlib import Path |
9 | 7 | from typing import Literal, TYPE_CHECKING, TypedDict, TypeVar, overload |
10 | 8 | from collections.abc import Iterable, Sequence |
|
18 | 16 | from .schedules_endpoint import AddResponse |
19 | 17 |
|
20 | 18 | from tableauserverclient.server.endpoint.dqw_endpoint import _DataQualityWarningEndpoint |
21 | | -from tableauserverclient.server.endpoint.endpoint import QuerysetEndpoint, api, parameter_added_in |
| 19 | +from tableauserverclient.server.endpoint.endpoint import DownloadableMixin, QuerysetEndpoint, api, parameter_added_in |
22 | 20 | from tableauserverclient.server.endpoint.exceptions import ( |
23 | 21 | DUPLICATE_EXTRACT_JOB_CODE, |
24 | 22 | InternalServerError, |
|
30 | 28 |
|
31 | 29 | from tableauserverclient.config import ALLOWED_FILE_EXTENSIONS, BYTES_PER_MB, config |
32 | 30 | from tableauserverclient.filesys_helpers import ( |
33 | | - make_download_path, |
34 | 31 | get_file_type, |
35 | 32 | get_file_object_size, |
36 | | - to_filename, |
37 | 33 | ) |
38 | 34 | from tableauserverclient.helpers.logging import logger |
39 | 35 | from tableauserverclient.models import ( |
|
46 | 42 | ) |
47 | 43 | from tableauserverclient.server import RequestFactory, RequestOptions |
48 | 44 |
|
49 | | -io_types = (io.BytesIO, io.BufferedReader) |
50 | 45 | io_types_r = (io.BytesIO, io.BufferedReader) |
51 | | -io_types_w = (io.BytesIO, io.BufferedWriter) |
52 | 46 |
|
53 | 47 | FilePath = str | os.PathLike |
54 | 48 | FileObject = io.BufferedReader | io.BytesIO |
|
101 | 95 | _UNSET = object() |
102 | 96 |
|
103 | 97 |
|
104 | | -class Datasources(QuerysetEndpoint[DatasourceItem], TaggingMixin[DatasourceItem]): |
| 98 | +class Datasources(QuerysetEndpoint[DatasourceItem], TaggingMixin[DatasourceItem], DownloadableMixin): |
105 | 99 | def __init__(self, parent_srv: "Server") -> None: |
106 | 100 | super().__init__(parent_srv) |
107 | 101 | self._permissions = _PermissionsEndpoint(parent_srv, lambda: self.baseurl) |
@@ -1061,22 +1055,7 @@ def download_revision( |
1061 | 1055 | if not include_extract: |
1062 | 1056 | url += "?includeExtract=False" |
1063 | 1057 |
|
1064 | | - with closing(self.get_request(url, parameters={"stream": True})) as server_response: |
1065 | | - m = Message() |
1066 | | - m["Content-Disposition"] = server_response.headers["Content-Disposition"] |
1067 | | - filename = m.get_filename(failobj="") |
1068 | | - if isinstance(filepath, io_types_w): |
1069 | | - for chunk in server_response.iter_content(1024): # 1KB |
1070 | | - filepath.write(chunk) |
1071 | | - return_path = filepath |
1072 | | - else: |
1073 | | - filename = to_filename(os.path.basename(filename)) |
1074 | | - download_path = make_download_path(filepath, filename) |
1075 | | - with open(download_path, "wb") as f: |
1076 | | - for chunk in server_response.iter_content(1024): # 1KB |
1077 | | - f.write(chunk) |
1078 | | - return_path = os.path.abspath(download_path) |
1079 | | - |
| 1058 | + return_path = self._download_content(url, filepath) |
1080 | 1059 | logger.info(f"Downloaded datasource revision {revision_number} to {return_path} (ID: {datasource_id})") |
1081 | 1060 | return return_path |
1082 | 1061 |
|
|
0 commit comments