diff --git a/tableauserverclient/config.py b/tableauserverclient/config.py index 1a4a7dc37..63872398f 100644 --- a/tableauserverclient/config.py +++ b/tableauserverclient/config.py @@ -1,13 +1,25 @@ -# TODO: check for env variables, else set default values +import os ALLOWED_FILE_EXTENSIONS = ["tds", "tdsx", "tde", "hyper", "parquet"] BYTES_PER_MB = 1024 * 1024 -# For when a datasource is over 64MB, break it into 5MB(standard chunk size) chunks -CHUNK_SIZE_MB = 5 * 10 # 5MB felt too slow, upped it to 50 - DELAY_SLEEP_SECONDS = 0.1 # The maximum size of a file that can be published in a single request is 64MB FILESIZE_LIMIT_MB = 64 + + +class Config: + # For when a datasource is over 64MB, break it into 5MB(standard chunk size) chunks + @property + def CHUNK_SIZE_MB(self): + return int(os.getenv("TSC_CHUNK_SIZE_MB", 5 * 10)) # 5MB felt too slow, upped it to 50 + + # Default page size + @property + def PAGE_SIZE(self): + return int(os.getenv("TSC_PAGE_SIZE", 100)) + + +config = Config() diff --git a/tableauserverclient/server/endpoint/datasources_endpoint.py b/tableauserverclient/server/endpoint/datasources_endpoint.py index 316f078a2..c795b03a3 100644 --- a/tableauserverclient/server/endpoint/datasources_endpoint.py +++ b/tableauserverclient/server/endpoint/datasources_endpoint.py @@ -21,7 +21,7 @@ from tableauserverclient.server.endpoint.permissions_endpoint import _PermissionsEndpoint from tableauserverclient.server.endpoint.resource_tagger import _ResourceTagger -from tableauserverclient.config import ALLOWED_FILE_EXTENSIONS, FILESIZE_LIMIT_MB, BYTES_PER_MB, CHUNK_SIZE_MB +from tableauserverclient.config import ALLOWED_FILE_EXTENSIONS, FILESIZE_LIMIT_MB, BYTES_PER_MB, config from tableauserverclient.filesys_helpers import ( make_download_path, get_file_type, @@ -272,7 +272,7 @@ def publish( if file_size >= FILESIZE_LIMIT_MB * BYTES_PER_MB: logger.info( "Publishing {} to server with chunking method (datasource over {}MB, chunk size {}MB)".format( - filename, FILESIZE_LIMIT_MB, CHUNK_SIZE_MB + filename, FILESIZE_LIMIT_MB, config.CHUNK_SIZE_MB ) ) upload_session_id = self.parent_srv.fileuploads.upload(file) diff --git a/tableauserverclient/server/endpoint/fileuploads_endpoint.py b/tableauserverclient/server/endpoint/fileuploads_endpoint.py index a0e29e508..0d30797c1 100644 --- a/tableauserverclient/server/endpoint/fileuploads_endpoint.py +++ b/tableauserverclient/server/endpoint/fileuploads_endpoint.py @@ -2,7 +2,7 @@ from tableauserverclient import datetime_helpers as datetime from tableauserverclient.helpers.logging import logger -from tableauserverclient.config import BYTES_PER_MB, CHUNK_SIZE_MB +from tableauserverclient.config import BYTES_PER_MB, config from tableauserverclient.models import FileuploadItem from tableauserverclient.server import RequestFactory @@ -41,7 +41,7 @@ def _read_chunks(self, file): try: while True: - chunked_content = file_content.read(CHUNK_SIZE_MB * BYTES_PER_MB) + chunked_content = file_content.read(config.CHUNK_SIZE_MB * BYTES_PER_MB) if not chunked_content: break yield chunked_content diff --git a/tableauserverclient/server/query.py b/tableauserverclient/server/query.py index 195139269..bbca612e9 100644 --- a/tableauserverclient/server/query.py +++ b/tableauserverclient/server/query.py @@ -1,6 +1,7 @@ from collections.abc import Sized from itertools import count from typing import Iterable, Iterator, List, Optional, Protocol, Tuple, TYPE_CHECKING, TypeVar, overload +from tableauserverclient.config import config from tableauserverclient.models.pagination_item import PaginationItem from tableauserverclient.server.filter import Filter from tableauserverclient.server.request_options import RequestOptions @@ -35,7 +36,7 @@ def to_camel_case(word: str) -> str: class QuerySet(Iterable[T], Sized): def __init__(self, model: "QuerysetEndpoint[T]", page_size: Optional[int] = None) -> None: self.model = model - self.request_options = RequestOptions(pagesize=page_size or 100) + self.request_options = RequestOptions(pagesize=page_size or config.PAGE_SIZE) self._result_cache: List[T] = [] self._pagination_item = PaginationItem() diff --git a/tableauserverclient/server/request_options.py b/tableauserverclient/server/request_options.py index 5cc06bf9d..563018d1c 100644 --- a/tableauserverclient/server/request_options.py +++ b/tableauserverclient/server/request_options.py @@ -2,6 +2,7 @@ from typing_extensions import Self +from tableauserverclient.config import config from tableauserverclient.models.property_decorators import property_is_int import logging @@ -115,9 +116,9 @@ class Direction: Desc = "desc" Asc = "asc" - def __init__(self, pagenumber=1, pagesize=100): + def __init__(self, pagenumber=1, pagesize=None): self.pagenumber = pagenumber - self.pagesize = pagesize + self.pagesize = pagesize or config.PAGE_SIZE self.sort = set() self.filter = set() diff --git a/test/test_fileuploads.py b/test/test_fileuploads.py index cf0861e24..50a5ef48b 100644 --- a/test/test_fileuploads.py +++ b/test/test_fileuploads.py @@ -1,8 +1,11 @@ +import contextlib +import io import os import unittest import requests_mock +from tableauserverclient.config import BYTES_PER_MB, config from tableauserverclient.server import Server from ._utils import asset @@ -11,6 +14,17 @@ FILEUPLOAD_APPEND = os.path.join(TEST_ASSET_DIR, "fileupload_append.xml") +@contextlib.contextmanager +def set_env(**environ): + old_environ = dict(os.environ) + os.environ.update(environ) + try: + yield + finally: + os.environ.clear() + os.environ.update(old_environ) + + class FileuploadsTests(unittest.TestCase): def setUp(self): self.server = Server("http://test", False) @@ -62,3 +76,14 @@ def test_upload_chunks_file_object(self): actual = self.server.fileuploads.upload(file_content) self.assertEqual(upload_id, actual) + + def test_upload_chunks_config(self): + data = io.BytesIO() + data.write(b"1" * (config.CHUNK_SIZE_MB * BYTES_PER_MB + 1)) + data.seek(0) + with set_env(TSC_CHUNK_SIZE_MB="1"): + chunker = self.server.fileuploads._read_chunks(data) + chunk = next(chunker) + assert len(chunk) == config.CHUNK_SIZE_MB * BYTES_PER_MB + data.seek(0) + assert len(chunk) < len(data.read()) diff --git a/test/test_pager.py b/test/test_pager.py index b60559b2b..7659f2725 100644 --- a/test/test_pager.py +++ b/test/test_pager.py @@ -1,9 +1,11 @@ +import contextlib import os import unittest import requests_mock import tableauserverclient as TSC +from tableauserverclient.config import config TEST_ASSET_DIR = os.path.join(os.path.dirname(__file__), "assets") @@ -12,6 +14,17 @@ GET_XML_PAGE3 = os.path.join(TEST_ASSET_DIR, "workbook_get_page_3.xml") +@contextlib.contextmanager +def set_env(**environ): + old_environ = dict(os.environ) + os.environ.update(environ) + try: + yield + finally: + os.environ.clear() + os.environ.update(old_environ) + + class PagerTests(unittest.TestCase): def setUp(self): self.server = TSC.Server("http://test", False) @@ -88,3 +101,15 @@ def test_pager_with_options(self): # Should have the last workbook wb3 = workbooks.pop() self.assertEqual(wb3.name, "Page3Workbook") + + def test_pager_with_env_var(self): + with set_env(TSC_PAGE_SIZE="1000"): + assert config.PAGE_SIZE == 1000 + loop = TSC.Pager(self.server.workbooks) + assert loop._options.pagesize == 1000 + + def test_queryset_with_env_var(self): + with set_env(TSC_PAGE_SIZE="1000"): + assert config.PAGE_SIZE == 1000 + loop = self.server.workbooks.all() + assert loop.request_options.pagesize == 1000