Extract download streaming logic into DownloadableMixin

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",

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
 

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")

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
 

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