diff --git a/tableauserverclient/server/endpoint/__init__.py b/tableauserverclient/server/endpoint/__init__.py index d944bc429..0a6eb8640 100644 --- a/tableauserverclient/server/endpoint/__init__.py +++ b/tableauserverclient/server/endpoint/__init__.py @@ -4,7 +4,7 @@ from tableauserverclient.server.endpoint.data_alert_endpoint import DataAlerts from tableauserverclient.server.endpoint.databases_endpoint import Databases from tableauserverclient.server.endpoint.datasources_endpoint import Datasources -from tableauserverclient.server.endpoint.endpoint import Endpoint, QuerysetEndpoint +from tableauserverclient.server.endpoint.endpoint import Endpoint, QuerysetEndpoint, DownloadableMixin from tableauserverclient.server.endpoint.exceptions import ServerResponseError, MissingRequiredFieldError from tableauserverclient.server.endpoint.extensions_endpoint import Extensions from tableauserverclient.server.endpoint.favorites_endpoint import Favorites @@ -40,6 +40,7 @@ "DataAlerts", "Databases", "Datasources", + "DownloadableMixin", "QuerysetEndpoint", "MissingRequiredFieldError", "Endpoint", diff --git a/tableauserverclient/server/endpoint/datasources_endpoint.py b/tableauserverclient/server/endpoint/datasources_endpoint.py index 30c66d7ad..f5212485a 100644 --- a/tableauserverclient/server/endpoint/datasources_endpoint.py +++ b/tableauserverclient/server/endpoint/datasources_endpoint.py @@ -1,10 +1,8 @@ -from email.message import Message import copy import json import io import os -from contextlib import closing from pathlib import Path from typing import Literal, TYPE_CHECKING, TypedDict, TypeVar, overload from collections.abc import Iterable, Sequence @@ -18,7 +16,7 @@ from .schedules_endpoint import AddResponse from tableauserverclient.server.endpoint.dqw_endpoint import _DataQualityWarningEndpoint -from tableauserverclient.server.endpoint.endpoint import QuerysetEndpoint, api, parameter_added_in +from tableauserverclient.server.endpoint.endpoint import DownloadableMixin, QuerysetEndpoint, api, parameter_added_in from tableauserverclient.server.endpoint.exceptions import ( DUPLICATE_EXTRACT_JOB_CODE, InternalServerError, @@ -30,10 +28,8 @@ from tableauserverclient.config import ALLOWED_FILE_EXTENSIONS, BYTES_PER_MB, config from tableauserverclient.filesys_helpers import ( - make_download_path, get_file_type, get_file_object_size, - to_filename, ) from tableauserverclient.helpers.logging import logger from tableauserverclient.models import ( @@ -46,9 +42,7 @@ ) from tableauserverclient.server import RequestFactory, RequestOptions -io_types = (io.BytesIO, io.BufferedReader) io_types_r = (io.BytesIO, io.BufferedReader) -io_types_w = (io.BytesIO, io.BufferedWriter) FilePath = str | os.PathLike FileObject = io.BufferedReader | io.BytesIO @@ -101,7 +95,7 @@ _UNSET = object() -class Datasources(QuerysetEndpoint[DatasourceItem], TaggingMixin[DatasourceItem]): +class Datasources(QuerysetEndpoint[DatasourceItem], TaggingMixin[DatasourceItem], DownloadableMixin): def __init__(self, parent_srv: "Server") -> None: super().__init__(parent_srv) self._permissions = _PermissionsEndpoint(parent_srv, lambda: self.baseurl) @@ -1061,22 +1055,7 @@ def download_revision( if not include_extract: url += "?includeExtract=False" - with closing(self.get_request(url, parameters={"stream": True})) as server_response: - m = Message() - m["Content-Disposition"] = server_response.headers["Content-Disposition"] - filename = m.get_filename(failobj="") - if isinstance(filepath, io_types_w): - for chunk in server_response.iter_content(1024): # 1KB - filepath.write(chunk) - return_path = filepath - else: - filename = to_filename(os.path.basename(filename)) - download_path = make_download_path(filepath, filename) - with open(download_path, "wb") as f: - for chunk in server_response.iter_content(1024): # 1KB - f.write(chunk) - return_path = os.path.abspath(download_path) - + return_path = self._download_content(url, filepath) logger.info(f"Downloaded datasource revision {revision_number} to {return_path} (ID: {datasource_id})") return return_path diff --git a/tableauserverclient/server/endpoint/endpoint.py b/tableauserverclient/server/endpoint/endpoint.py index 149aab5ec..31a0806dc 100644 --- a/tableauserverclient/server/endpoint/endpoint.py +++ b/tableauserverclient/server/endpoint/endpoint.py @@ -1,3 +1,7 @@ +from email.message import Message +import io +import os +from contextlib import closing from typing_extensions import Concatenate, ParamSpec from tableauserverclient import datetime_helpers as datetime @@ -18,6 +22,7 @@ from tableauserverclient.models.pagination_item import PaginationItem from tableauserverclient.server.request_options import RequestOptions +from tableauserverclient.filesys_helpers import to_filename, make_download_path from tableauserverclient.server.endpoint.exceptions import ( FailedSignInError, @@ -323,6 +328,58 @@ def wrapper(self: E, *args: P.args, **kwargs: P.kwargs) -> R: T = TypeVar("T") +_io_types_w = (io.BytesIO, io.BufferedWriter) + +FilePath = str | os.PathLike +FileObjectW = io.BufferedWriter | io.BytesIO +PathOrFileW = FilePath | FileObjectW + + +class DownloadableMixin: + """Mixin for endpoints whose resources can be downloaded as binary files. + + Provides a single private helper that streams a server response to a file + path or writable file object, avoiding copy-paste of the identical streaming + loop in Workbooks, Datasources, and Flows. + """ + + def _download_content( + self, + url: str, + filepath: PathOrFileW | None, + ) -> PathOrFileW: + """Stream content at url to filepath and return the resolved path. + + Parameters + ---------- + url : str + Fully-qualified URL whose response body should be saved. + filepath : PathOrFileW | None + Destination file path or writable file object. When None the file + is saved to the current working directory using the server-supplied + filename from the Content-Disposition header. + + Returns + ------- + PathOrFileW + The absolute file path written, or the caller-supplied file object. + """ + with closing(self.get_request(url, parameters={"stream": True})) as server_response: # type: ignore[attr-defined] + m = Message() + m["Content-Disposition"] = server_response.headers["Content-Disposition"] + filename = m.get_filename(failobj="") + if isinstance(filepath, _io_types_w): + for chunk in server_response.iter_content(1024): # 1KB + filepath.write(chunk) + return filepath + else: + filename = to_filename(os.path.basename(filename)) + download_path = make_download_path(filepath, filename) + with open(download_path, "wb") as f: + for chunk in server_response.iter_content(1024): # 1KB + f.write(chunk) + return os.path.abspath(download_path) + class QuerysetEndpoint(Endpoint, Generic[T]): @api(version="2.0") diff --git a/tableauserverclient/server/endpoint/flows_endpoint.py b/tableauserverclient/server/endpoint/flows_endpoint.py index ea7b59526..8595006b8 100644 --- a/tableauserverclient/server/endpoint/flows_endpoint.py +++ b/tableauserverclient/server/endpoint/flows_endpoint.py @@ -1,15 +1,13 @@ -from email.message import Message import copy import io import logging import os -from contextlib import closing from pathlib import Path from typing import TYPE_CHECKING from collections.abc import Iterable from tableauserverclient.server.endpoint.dqw_endpoint import _DataQualityWarningEndpoint -from tableauserverclient.server.endpoint.endpoint import QuerysetEndpoint, api +from tableauserverclient.server.endpoint.endpoint import DownloadableMixin, QuerysetEndpoint, api from tableauserverclient.server.endpoint.exceptions import ( DUPLICATE_EXTRACT_JOB_CODE, InternalServerError, @@ -21,18 +19,12 @@ from tableauserverclient.models import FlowItem, PaginationItem, ConnectionItem, JobItem from tableauserverclient.server import RequestFactory from tableauserverclient.filesys_helpers import ( - to_filename, - make_download_path, get_file_type, get_file_object_size, ) from tableauserverclient.server.query import QuerySet io_types_r = (io.BytesIO, io.BufferedReader) -io_types_w = (io.BytesIO, io.BufferedWriter) - -io_types_r = (io.BytesIO, io.BufferedReader) -io_types_w = (io.BytesIO, io.BufferedWriter) # The maximum size of a file that can be published in a single request is 64MB FILESIZE_LIMIT = 1024 * 1024 * 64 # 64MB @@ -55,7 +47,7 @@ PathOrFileW = FilePath | FileObjectW -class Flows(QuerysetEndpoint[FlowItem], TaggingMixin[FlowItem]): +class Flows(QuerysetEndpoint[FlowItem], TaggingMixin[FlowItem], DownloadableMixin): def __init__(self, parent_srv): super().__init__(parent_srv) self._resource_tagger = _ResourceTagger(parent_srv) @@ -227,22 +219,7 @@ def download(self, flow_id: str, filepath: PathOrFileW | None = None) -> PathOrF raise ValueError(error) url = f"{self.baseurl}/{flow_id}/content" - with closing(self.get_request(url, parameters={"stream": True})) as server_response: - m = Message() - m["Content-Disposition"] = server_response.headers["Content-Disposition"] - filename = m.get_filename(failobj="") - if isinstance(filepath, io_types_w): - for chunk in server_response.iter_content(1024): # 1KB - filepath.write(chunk) - return_path = filepath - else: - filename = to_filename(os.path.basename(filename)) - download_path = make_download_path(filepath, filename) - with open(download_path, "wb") as f: - for chunk in server_response.iter_content(1024): # 1KB - f.write(chunk) - return_path = os.path.abspath(download_path) - + return_path = self._download_content(url, filepath) logger.info(f"Downloaded flow to {return_path} (ID: {flow_id})") return return_path diff --git a/tableauserverclient/server/endpoint/workbooks_endpoint.py b/tableauserverclient/server/endpoint/workbooks_endpoint.py index dd9722cbe..8310d58b3 100644 --- a/tableauserverclient/server/endpoint/workbooks_endpoint.py +++ b/tableauserverclient/server/endpoint/workbooks_endpoint.py @@ -1,15 +1,13 @@ -from email.message import Message import copy import io import logging import os -from contextlib import closing from pathlib import Path from tableauserverclient.models.permissions_item import PermissionsRule from tableauserverclient.server.query import QuerySet -from tableauserverclient.server.endpoint.endpoint import QuerysetEndpoint, api, parameter_added_in +from tableauserverclient.server.endpoint.endpoint import DownloadableMixin, QuerysetEndpoint, api, parameter_added_in from tableauserverclient.server.endpoint.exceptions import ( DUPLICATE_EXTRACT_JOB_CODE, InternalServerError, @@ -21,8 +19,6 @@ from tableauserverclient.server.endpoint.resource_tagger import TaggingMixin from tableauserverclient.filesys_helpers import ( - to_filename, - make_download_path, get_file_type, get_file_object_size, ) @@ -47,7 +43,6 @@ from tableauserverclient.server.endpoint.schedules_endpoint import AddResponse io_types_r = (io.BytesIO, io.BufferedReader) -io_types_w = (io.BytesIO, io.BufferedWriter) # The maximum size of a file that can be published in a single request is 64MB FILESIZE_LIMIT = 1024 * 1024 * 64 # 64MB @@ -67,7 +62,7 @@ _UNSET = object() -class Workbooks(QuerysetEndpoint[WorkbookItem], TaggingMixin[WorkbookItem]): +class Workbooks(QuerysetEndpoint[WorkbookItem], TaggingMixin[WorkbookItem], DownloadableMixin): def __init__(self, parent_srv: "Server") -> None: super().__init__(parent_srv) self._permissions = _PermissionsEndpoint(parent_srv, lambda: self.baseurl) @@ -1123,22 +1118,7 @@ def download_revision( if not include_extract: url += "?includeExtract=False" - with closing(self.get_request(url, parameters={"stream": True})) as server_response: - m = Message() - m["Content-Disposition"] = server_response.headers["Content-Disposition"] - filename = m.get_filename(failobj="") - if isinstance(filepath, io_types_w): - for chunk in server_response.iter_content(1024): # 1KB - filepath.write(chunk) - return_path = filepath - else: - filename = to_filename(os.path.basename(filename)) - download_path = make_download_path(filepath, filename) - with open(download_path, "wb") as f: - for chunk in server_response.iter_content(1024): # 1KB - f.write(chunk) - return_path = os.path.abspath(download_path) - + return_path = self._download_content(url, filepath) logger.info(f"Downloaded workbook revision {revision_number} to {return_path} (ID: {workbook_id})") return return_path