feat: get page and chunk size from env vars

Author: jorwoods

tableauserverclient/config.py

@@ -1,13 +1,23 @@
-# 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()

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)

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

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

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

test/test_pager.py

@@ -1,16 +1,27 @@
+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")
 
 GET_XML_PAGE1 = os.path.join(TEST_ASSET_DIR, "workbook_get_page_1.xml")
 GET_XML_PAGE2 = os.path.join(TEST_ASSET_DIR, "workbook_get_page_2.xml")
 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):
@@ -88,3 +99,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) -> None:
+        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) -> None:
+        with set_env(TSC_PAGE_SIZE="1000"):
+            assert config.PAGE_SIZE == 1000
+            loop = self.server.workbooks.all()
+            assert loop.request_options.pagesize == 1000