Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions tableauserverclient/server/endpoint/custom_views_endpoint.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,13 @@
import io
import logging
import os
from contextlib import closing
from pathlib import Path
from typing import TYPE_CHECKING
from collections.abc import Iterator

from tableauserverclient.config import BYTES_PER_MB, config
from tableauserverclient.filesys_helpers import get_file_object_size
from tableauserverclient.server.endpoint.endpoint import QuerysetEndpoint, api
from tableauserverclient.server.endpoint.endpoint import DownloadableMixin, QuerysetEndpoint, api
from tableauserverclient.server.endpoint.exceptions import MissingRequiredFieldError
from tableauserverclient.models import CustomViewItem, PaginationItem
from tableauserverclient.server import (
Expand Down Expand Up @@ -42,7 +41,7 @@
io_types_w = (io.BufferedWriter, io.BytesIO)


class CustomViews(QuerysetEndpoint[CustomViewItem]):
class CustomViews(QuerysetEndpoint[CustomViewItem], DownloadableMixin):
def __init__(self, parent_srv):
super().__init__(parent_srv)

Expand Down Expand Up @@ -229,9 +228,7 @@ def _get_custom_view_csv(
self, custom_view_item: CustomViewItem, req_options: "CSVRequestOptions | None"
) -> Iterator[bytes]:
url = f"{self.baseurl}/{custom_view_item.id}/data"

with closing(self.get_request(url, request_object=req_options, parameters={"stream": True})) as server_response:
yield from server_response.iter_content(1024)
return self._stream_content(url, req_options)

@api(version="3.18")
def update(self, view_item: CustomViewItem) -> CustomViewItem | None:
Expand Down
36 changes: 30 additions & 6 deletions tableauserverclient/server/endpoint/endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from packaging.version import Version
from functools import wraps
from xml.etree.ElementTree import ParseError
from collections.abc import Iterator
from typing import (
Any,
Callable,
Expand All @@ -20,8 +21,9 @@
)
from typing_extensions import Self

from tableauserverclient.config import BYTES_PER_MB, config
from tableauserverclient.models.pagination_item import PaginationItem
from tableauserverclient.server.request_options import RequestOptions
from tableauserverclient.server.request_options import RequestOptions, RequestOptionsBase
from tableauserverclient.filesys_helpers import to_filename, make_download_path

from tableauserverclient.server.endpoint.exceptions import (
Expand Down Expand Up @@ -338,9 +340,13 @@ def wrapper(self: E, *args: P.args, **kwargs: P.kwargs) -> R:
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.
Provides two private helpers, avoiding copy-paste of the streaming loop
across endpoints:

- _download_content streams a response to a file path or writable object
(used by Workbooks, Datasources, Flows).
- _stream_content yields chunks for callers that want to materialize the
body in-memory or process it lazily (used by Views, CustomViews).
"""

def _download_content(
Expand Down Expand Up @@ -368,18 +374,36 @@ def _download_content(
m = Message()
m["Content-Disposition"] = server_response.headers["Content-Disposition"]
filename = m.get_filename(failobj="")
chunk_size = config.CHUNK_SIZE_MB * BYTES_PER_MB
if isinstance(filepath, _io_types_w):
for chunk in server_response.iter_content(1024): # 1KB
for chunk in server_response.iter_content(chunk_size):
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
for chunk in server_response.iter_content(chunk_size):
f.write(chunk)
return os.path.abspath(download_path)

def _stream_content(
self,
url: str,
request_object: RequestOptionsBase | None = None,
) -> Iterator[bytes]:
"""Stream content at url as chunks.

Suitable for callers that want to materialize the body in-memory
(e.g. b"".join(iterator)) or process it lazily as a stream, rather than
write it to a file. Complements _download_content, which writes to disk.
"""
chunk_size = config.CHUNK_SIZE_MB * BYTES_PER_MB
with closing(
self.get_request(url, request_object=request_object, parameters={"stream": True}) # type: ignore[attr-defined]
) as server_response:
yield from server_response.iter_content(chunk_size)


class QuerysetEndpoint(Endpoint, Generic[T]):
@api(version="2.0")
Expand Down
13 changes: 4 additions & 9 deletions tableauserverclient/server/endpoint/views_endpoint.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
import logging
from contextlib import closing

from tableauserverclient.models.permissions_item import PermissionsRule
from tableauserverclient.server.endpoint.endpoint import QuerysetEndpoint, api
from tableauserverclient.server.endpoint.endpoint import DownloadableMixin, QuerysetEndpoint, api
from tableauserverclient.server.endpoint.exceptions import MissingRequiredFieldError, UnsupportedAttributeError
from tableauserverclient.server.endpoint.permissions_endpoint import _PermissionsEndpoint
from tableauserverclient.server.endpoint.resource_tagger import TaggingMixin
Expand All @@ -25,7 +24,7 @@
)


class Views(QuerysetEndpoint[ViewItem], TaggingMixin[ViewItem]):
class Views(QuerysetEndpoint[ViewItem], TaggingMixin[ViewItem], DownloadableMixin):
"""
The Tableau Server Client provides methods for interacting with view
resources, or endpoints. These methods correspond to the endpoints for views
Expand Down Expand Up @@ -262,9 +261,7 @@ def csv_fetcher():

def _get_view_csv(self, view_item: ViewItem, req_options: "CSVRequestOptions | None") -> Iterator[bytes]:
url = f"{self.baseurl}/{view_item.id}/data"

with closing(self.get_request(url, request_object=req_options, parameters={"stream": True})) as server_response:
yield from server_response.iter_content(1024)
return self._stream_content(url, req_options)

@api(version="3.8")
def populate_excel(self, view_item: ViewItem, req_options: "ExcelRequestOptions | None" = None) -> None:
Expand Down Expand Up @@ -301,9 +298,7 @@ def excel_fetcher():

def _get_view_excel(self, view_item: ViewItem, req_options: "ExcelRequestOptions | None") -> Iterator[bytes]:
url = f"{self.baseurl}/{view_item.id}/crosstab/excel"

with closing(self.get_request(url, request_object=req_options, parameters={"stream": True})) as server_response:
yield from server_response.iter_content(1024)
return self._stream_content(url, req_options)

@api(version="3.2")
def populate_permissions(self, item: ViewItem) -> None:
Expand Down
31 changes: 31 additions & 0 deletions test/test_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,37 @@ def test_populate_csv_default_maxage(server: TSC.Server) -> None:
assert response == csv_file


def test_stream_content_uses_configured_chunk_size(server: TSC.Server, monkeypatch) -> None:
# Regression guard: the shared DownloadableMixin._stream_content path must
# stream with config.CHUNK_SIZE_MB * BYTES_PER_MB, not the pre-fix 1024 bytes.
from tableauserverclient.config import BYTES_PER_MB, config

captured: list[int] = []
real_iter_content = None

def spy_iter_content(self, chunk_size=None, decode_unicode=False):
captured.append(chunk_size)
assert real_iter_content is not None
return real_iter_content(self, chunk_size, decode_unicode)

import requests

real_iter_content = requests.Response.iter_content
monkeypatch.setattr(requests.Response, "iter_content", spy_iter_content)

response = POPULATE_CSV.read_bytes()
with requests_mock.mock() as m:
m.get(server.views.baseurl + "/d79634e1-6063-4ec9-95ff-50acbf609ff5/data", content=response)
single_view = TSC.ViewItem()
single_view._id = "d79634e1-6063-4ec9-95ff-50acbf609ff5"
server.views.populate_csv(single_view)
# Materialize the iterator so iter_content is actually invoked.
b"".join(single_view.csv)

assert captured, "iter_content was never invoked"
assert captured[0] == config.CHUNK_SIZE_MB * BYTES_PER_MB


def test_populate_image_missing_id(server: TSC.Server) -> None:
single_view = TSC.ViewItem()
single_view._id = None
Expand Down
29 changes: 29 additions & 0 deletions test/test_workbook.py
Original file line number Diff line number Diff line change
Expand Up @@ -339,6 +339,35 @@ def test_download_object(server: TSC.Server) -> None:
assert isinstance(file_path, BytesIO)


def test_download_uses_configured_chunk_size(server: TSC.Server, tmp_path: Path, monkeypatch) -> None:
# Regression guard: the shared DownloadableMixin._download_content path must
# stream with config.CHUNK_SIZE_MB * BYTES_PER_MB, not the pre-fix 1024 bytes.
from tableauserverclient.config import BYTES_PER_MB, config

captured: list[int] = []
real_iter_content = None

def spy_iter_content(self, chunk_size=None, decode_unicode=False):
captured.append(chunk_size)
assert real_iter_content is not None
return real_iter_content(self, chunk_size, decode_unicode)

import requests

real_iter_content = requests.Response.iter_content
monkeypatch.setattr(requests.Response, "iter_content", spy_iter_content)

with requests_mock.mock() as m:
m.get(
server.workbooks.baseurl + "/1f951daf-4061-451a-9df1-69a8062664f2/content",
headers={"Content-Disposition": 'name="tableau_workbook"; filename="RESTAPISample.twbx"'},
)
server.workbooks.download("1f951daf-4061-451a-9df1-69a8062664f2", filepath=tmp_path)

assert captured, "iter_content was never invoked"
assert captured[0] == config.CHUNK_SIZE_MB * BYTES_PER_MB


def test_download_sanitizes_name(server: TSC.Server, tmp_path: Path) -> None:
filename = "Name,With,Commas.twbx"
disposition = f'name="tableau_workbook"; filename="{filename}"'
Expand Down
Loading