diff --git a/tableauserverclient/server/server.py b/tableauserverclient/server/server.py index 6b0338970..6dc95557f 100644 --- a/tableauserverclient/server/server.py +++ b/tableauserverclient/server/server.py @@ -1,8 +1,10 @@ from tableauserverclient.helpers.logging import logger import requests +import threading import urllib3 import ssl +import weakref from defusedxml.ElementTree import fromstring, ParseError from packaging.version import Version @@ -121,6 +123,18 @@ class Server: Notes ----- + A single Server instance may be shared across threads (for example with + ``concurrent.futures.ThreadPoolExecutor``). Authentication state (the auth + token, site and user IDs, and API version) is shared by all threads, while + each thread transparently gets its own ``requests.Session`` for HTTP calls, + since ``requests.Session`` itself is not guaranteed to be thread-safe. Sign + in, and call ``use_server_version()`` if you need it, before spawning + worker threads; signing out invalidates the sessions of all threads. + ``session_factory`` may be called once per thread, so a custom factory + should be safe to call concurrently. Call ``close()`` (or use the Server + as a context manager) after worker threads finish to release the pooled + HTTP connections of every thread's session. + When using Python 3.12 or later with older versions of Tableau Server, you may encounter SSL errors related to weak Diffie-Hellman keys. This is because newer Python versions enforce stronger security requirements. You can temporarily work around this using @@ -153,6 +167,19 @@ def __init__(self, server_address, use_server_version=False, http_options=None, self._server_address: str = server_address self._session_factory = session_factory or requests.session + # Thread-safety machinery. requests.Session is not guaranteed to be + # thread-safe (see psf/requests#2766), so each thread that makes calls + # through this Server instance gets its own Session object, created + # lazily from session_factory. The epoch counter invalidates every + # thread's cached session when auth state is cleared (sign out). + self._auth_lock = threading.Lock() + self._session_epoch = 0 + self._thread_sessions = threading.local() + # Weak references to every session created for any thread, so close() + # can release their pooled connections. Weak so that sessions belonging + # to threads that have exited can still be garbage collected. + self._all_sessions: "weakref.WeakSet[requests.Session]" = weakref.WeakSet() + self.auth = Auth(self) self.views = Views(self) self.users = Users(self) @@ -187,7 +214,6 @@ def __init__(self, server_address, use_server_version=False, http_options=None, self.oidc = OIDC(self) self.extensions = Extensions(self) - self._session = self._session_factory() self._http_options = dict() # must set this before making a server call if http_options: self.add_http_options(http_options) @@ -203,7 +229,7 @@ def validate_connection_settings(self): Endpoint.set_user_agent(params) if not self._server_address.startswith("http://") and not self._server_address.startswith("https://"): self._server_address = "http://" + self._server_address - self._session.prepare_request(requests.Request("GET", url=self._server_address, params=self._http_options)) + self.session.prepare_request(requests.Request("GET", url=self._server_address, params=self._http_options)) except Exception as req_ex: raise ValueError("Server connection settings not valid", req_ex) @@ -226,21 +252,28 @@ def clear_http_options(self): self._http_options = dict() def _clear_auth(self): - self._site_id = None - self._user_id = None - self._auth_token = None - self._site_url = None - self._session = self._session_factory() + with self._auth_lock: + self._site_id = None + self._user_id = None + self._auth_token = None + self._site_url = None + # Invalidate the cached session of every thread so state such as + # cookies does not leak into a later sign in. Sessions are replaced + # lazily on next use rather than closed here, so requests already + # in flight on other threads are not disrupted (this matches the + # previous behavior of re-assigning the shared session). + self._session_epoch += 1 def _set_auth(self, site_id, user_id, auth_token, site_url=None): - self._site_id = site_id - self._user_id = user_id - self._auth_token = auth_token - self._site_url = site_url + with self._auth_lock: + self._site_id = site_id + self._user_id = user_id + self._auth_token = auth_token + self._site_url = site_url def _get_legacy_version(self): # the serverInfo call was introduced in 2.4, earlier than that we have this different call - response = self._session.get(self.server_address + "/auth?format=xml") + response = self.session.get(self.server_address + "/auth?format=xml") try: info_xml = fromstring(response.content) except ParseError as parseError: @@ -329,8 +362,29 @@ def http_options(self): return self._http_options @property - def session(self): - return self._session + def session(self) -> requests.Session: + """ + The requests.Session used for HTTP calls made by the current thread. + + requests.Session is not guaranteed to be thread-safe (see + psf/requests#2766), so each thread that makes calls through this Server + instance transparently gets its own Session object, created from + ``session_factory``. Sessions are cached per thread, so a thread pool + worker reuses its session (and its connection pool) across tasks. + Signing out invalidates the cached sessions of all threads. + """ + local = self._thread_sessions + epoch = self._session_epoch + if getattr(local, "session", None) is None or local.epoch != epoch: + session = self._session_factory() + with self._auth_lock: + self._all_sessions.add(session) + local.session = session + # `epoch` was read before the factory ran: if a sign out happened + # in between, local.epoch is already stale and the session will be + # replaced on the next access. + local.epoch = epoch + return local.session def is_signed_in(self): return self._auth_token is not None @@ -357,3 +411,33 @@ def configure_ssl(self, *, allow_weak_dh=False): # Remove any custom SSL context if we're reverting to default settings if "verify" in self._http_options: del self._http_options["verify"] + + def close(self) -> None: + """ + Release the pooled HTTP connections held by every thread's session. + + Call this when you are done with the server, after any worker threads + have finished their requests. Closing is a transport-level operation: + it does not sign out, so the auth token remains valid on the server + (use ``auth.sign_out()`` for that). The Server object remains usable + after close; any subsequent call transparently creates a new session. + """ + with self._auth_lock: + sessions = list(self._all_sessions) + self._all_sessions.clear() + # Invalidate every thread's cached (now closed) session so later + # use creates a fresh one instead of hitting closed pools. + self._session_epoch += 1 + for session in sessions: + session.close() + + def __enter__(self) -> "Server": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + # Note: does not sign out. server.auth.sign_in() already returns a + # context manager that signs out; nesting the two composes cleanly: + # with TSC.Server(...) as server: + # with server.auth.sign_in(auth): + # ... + self.close() diff --git a/test/test_thread_safety.py b/test/test_thread_safety.py new file mode 100644 index 000000000..6af454945 --- /dev/null +++ b/test/test_thread_safety.py @@ -0,0 +1,250 @@ +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest +import requests +import requests_mock + +import tableauserverclient as TSC + +TEST_ASSET_DIR = Path(__file__).parent / "assets" + +SIGN_IN_XML = TEST_ASSET_DIR / "auth_sign_in.xml" + +NUM_THREADS = 8 +CALLS_PER_THREAD = 5 + + +@pytest.fixture(scope="function") +def server() -> TSC.Server: + return TSC.Server("http://test", False) + + +@pytest.fixture(scope="function") +def signed_in_server(server: TSC.Server) -> TSC.Server: + with open(SIGN_IN_XML, "rb") as f: + response_xml = f.read().decode("utf-8") + with requests_mock.mock() as m: + m.post(server.auth.baseurl + "/signin", text=response_xml) + server.auth.sign_in(TSC.TableauAuth("testuser", "password", site_id="Samples")) + return server + + +def test_each_thread_gets_its_own_session(server: TSC.Server) -> None: + main_session = server.session + assert server.session is main_session # stable within a thread + + # Hold strong references to the session objects (not just id()s, which can + # be reused once a thread's session is garbage collected) and keep all + # threads alive at the same time behind a barrier. + barrier = threading.Barrier(NUM_THREADS) + results: dict[str, requests.Session] = {} + lock = threading.Lock() + + def worker(name: str) -> None: + barrier.wait() + first = server.session + assert server.session is first # stable within the worker thread too + with lock: + results[name] = first + + threads = [threading.Thread(target=worker, args=(f"t{i}",)) for i in range(NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + + sessions = list(results.values()) + [main_session] + # every thread saw a distinct session object + assert len({id(s) for s in sessions}) == NUM_THREADS + 1 + + +def test_session_factory_called_once_per_thread() -> None: + lock = threading.Lock() + created: list[requests.Session] = [] + + def counting_factory() -> requests.Session: + session = requests.Session() + with lock: + created.append(session) + return session + + server = TSC.Server("http://test", False, session_factory=counting_factory) + # constructing the server creates the constructing thread's session + assert len(created) == 1 + + barrier = threading.Barrier(NUM_THREADS) + + def worker() -> None: + barrier.wait() + for _ in range(CALLS_PER_THREAD): + server.session # repeated access must not create new sessions + + threads = [threading.Thread(target=worker) for _ in range(NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + + # one factory call for the main thread + exactly one per worker thread, + # despite CALLS_PER_THREAD accesses in each worker + assert len(created) == NUM_THREADS + 1 + + +def test_sign_out_invalidates_sessions_of_all_threads(signed_in_server: TSC.Server) -> None: + server = signed_in_server + before_main = server.session + + barrier = threading.Barrier(2) + worker_sessions: dict[str, int] = {} + + def worker() -> None: + worker_sessions["before"] = id(server.session) + barrier.wait() # let the main thread sign out + barrier.wait() + worker_sessions["after"] = id(server.session) + + t = threading.Thread(target=worker) + t.start() + barrier.wait() # worker has cached its session + with requests_mock.mock() as m: + m.post(server.auth.baseurl + "/signout", text="") + server.auth.sign_out() + barrier.wait() # worker re-reads its session + t.join() + + # both the main thread and the worker thread got fresh sessions + assert server.session is not before_main + assert worker_sessions["after"] != worker_sessions["before"] + assert not server.is_signed_in() + + +def test_concurrent_api_calls_use_per_thread_sessions(signed_in_server: TSC.Server) -> None: + server = signed_in_server + response_xml = (TEST_ASSET_DIR / "user_get_empty.xml").read_text() + + with requests_mock.mock() as m: + m.get(server.users.baseurl, text=response_xml) + + def worker() -> int: + for _ in range(CALLS_PER_THREAD): + _, pagination_item = server.users.get() + assert pagination_item.total_available == 0 + return id(server.session) + + with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor: + session_ids = list(executor.map(lambda _: worker(), range(NUM_THREADS))) + + # every request was actually made + assert m.call_count == NUM_THREADS * CALLS_PER_THREAD + # every request carried the shared auth token + assert all(r.headers["x-tableau-auth"] == server.auth_token for r in m.request_history) + + # the pool had NUM_THREADS workers; each distinct worker thread used a + # distinct session, and threads reused their session across tasks + assert 1 <= len(set(session_ids)) <= NUM_THREADS + + +def test_auth_state_is_set_atomically(server: TSC.Server) -> None: + """Readers must never observe a half-updated (site_id, user_id, token) triple.""" + stop = threading.Event() + errors: list[Exception] = [] + + def reader() -> None: + while not stop.is_set(): + try: + if server.is_signed_in(): + token = server.auth_token + site = server.site_id + # tokens and site ids are written together; a mismatched + # pair means a reader saw a partial update + assert (token, site) in (("token-a", "site-a"), ("token-b", "site-b")) + except TSC.server.endpoint.exceptions.NotSignedInError: + pass # signed out at the moment of the read; that's fine + except Exception as e: # pragma: no cover - only on failure + errors.append(e) + stop.set() + + readers = [threading.Thread(target=reader) for _ in range(4)] + for t in readers: + t.start() + try: + for _ in range(500): + server._set_auth("site-a", "user-a", "token-a", "url-a") + server._set_auth("site-b", "user-b", "token-b", "url-b") + server._clear_auth() + finally: + stop.set() + for t in readers: + t.join() + + assert errors == [] + + +class _TrackingSession(requests.Session): + def __init__(self) -> None: + super().__init__() + self.closed = False + + def close(self) -> None: + self.closed = True + super().close() + + +def _tracking_server() -> tuple[TSC.Server, list["_TrackingSession"]]: + created: list[_TrackingSession] = [] + lock = threading.Lock() + + def factory() -> requests.Session: + session = _TrackingSession() + with lock: + created.append(session) + return session + + return TSC.Server("http://test", False, session_factory=factory), created + + +def test_close_closes_sessions_of_all_threads() -> None: + server, created = _tracking_server() + barrier = threading.Barrier(NUM_THREADS) + + def worker() -> None: + barrier.wait() + server.session + + threads = [threading.Thread(target=worker) for _ in range(NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(created) == NUM_THREADS + 1 # workers + constructing thread + assert not any(s.closed for s in created) + + server.close() + assert all(s.closed for s in created) + + # the server remains usable after close: a fresh session is created + reopened = server.session + assert isinstance(reopened, _TrackingSession) + assert not reopened.closed + assert reopened not in created[: NUM_THREADS + 1] or len(created) == NUM_THREADS + 2 + + +def test_close_does_not_sign_out(signed_in_server: TSC.Server) -> None: + server = signed_in_server + token = server.auth_token + server.close() + # close() is transport-level only; auth state is untouched + assert server.is_signed_in() + assert server.auth_token == token + + +def test_context_manager_closes_on_exit() -> None: + server, created = _tracking_server() + with server as entered: + assert entered is server + server.session + assert len(created) == 1 + assert all(s.closed for s in created)