Make Server thread-safe with per-thread sessions (#1871)

* fix: make Server thread-safe with per-thread sessions

* fix torn auth-state reads exposed by free-threaded Python

* restore Server._session as a backwards-compatible property shim

Author: dxdc

tableauserverclient/server/server.py

@@ -1,8 +1,12 @@
 from tableauserverclient.helpers.logging import logger
 
 import requests
+import threading
 import urllib3
 import ssl
+import weakref
+
+from typing import NamedTuple, Optional
 
 from defusedxml.ElementTree import fromstring, ParseError
 from packaging.version import Version
@@ -60,6 +64,23 @@
 default_server_version = "2.4"  # first version that dropped the legacy auth endpoint
 
 
+class _AuthState(NamedTuple):
+    """
+    Immutable bundle of the auth fields that must change together.
+
+    The whole state is swapped by assigning a new instance, and a single
+    reference assignment or read is atomic on every CPython build,
+    free-threaded (no-GIL) builds included. Readers take one snapshot of
+    ``Server._auth_state`` and destructure it, so they can never observe a
+    token paired with the site or user of a different sign in.
+    """
+
+    site_id: Optional[str] = None
+    user_id: Optional[str] = None
+    auth_token: Optional[str] = None
+    site_url: Optional[str] = None
+
+
 class Server:
     """
     In the Tableau REST API, the server (https://MY-SERVER/) is the base or core
@@ -121,6 +142,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
@@ -141,9 +174,7 @@ class PublishMode:
         Replace = "Replace"
 
     def __init__(self, server_address, use_server_version=False, http_options=None, session_factory=None):
-        self._auth_token = None
-        self._site_id = None
-        self._user_id = None
+        self._auth_state = _AuthState()
         self._ssl_context = None
 
         # TODO: this needs to change to default to https, but without breaking existing code
@@ -153,6 +184,22 @@ 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). The
+        # lock guards the epoch counter and the session WeakSet; auth state
+        # itself is an immutable _AuthState swapped by reference and needs
+        # no lock.
+        self._session_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 +234,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 +249,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 +272,64 @@ 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()
+        # swapping the whole immutable state in one assignment is atomic on
+        # every CPython build, free-threaded (no-GIL) builds included
+        self._auth_state = _AuthState()
+        with self._session_lock:
+            # 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). The lock
+            # guards the epoch increment, which is not atomic without the GIL.
+            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
+        # swapping the whole immutable state in one assignment is atomic on
+        # every CPython build, free-threaded (no-GIL) builds included, so
+        # readers can never observe a partially-updated state
+        self._auth_state = _AuthState(site_id, user_id, auth_token, site_url)
+
+    # Backwards-compatible access to the individual auth fields. Existing code
+    # (including this project's tests) reads and assigns these directly as a
+    # sign-in shortcut. Each assignment swaps in a fully-formed state, so
+    # readers never see torn fields; assigning fields one at a time is not
+    # atomic as a group, though, so concurrent writers should use _set_auth.
+    @property
+    def _auth_token(self):
+        return self._auth_state.auth_token
+
+    @_auth_token.setter
+    def _auth_token(self, value) -> None:
+        self._auth_state = self._auth_state._replace(auth_token=value)
+
+    @property
+    def _site_id(self):
+        return self._auth_state.site_id
+
+    @_site_id.setter
+    def _site_id(self, value) -> None:
+        self._auth_state = self._auth_state._replace(site_id=value)
+
+    @property
+    def _user_id(self):
+        return self._auth_state.user_id
+
+    @_user_id.setter
+    def _user_id(self, value) -> None:
+        self._auth_state = self._auth_state._replace(user_id=value)
+
+    @property
+    def _site_url(self):
+        return self._auth_state.site_url
+
+    @_site_url.setter
+    def _site_url(self, value) -> None:
+        self._auth_state = self._auth_state._replace(site_url=value)
 
     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:
@@ -294,31 +383,37 @@ def namespace(self):
 
     @property
     def auth_token(self):
-        if self._auth_token is None:
+        # read the state once: a second self._auth_state read could observe
+        # a concurrent sign out and return None instead of raising
+        token = self._auth_state.auth_token
+        if token is None:
             error = "Missing authentication token. You must sign in first."
             raise NotSignedInError(error)
-        return self._auth_token
+        return token
 
     @property
     def site_id(self):
-        if self._site_id is None:
+        site_id = self._auth_state.site_id
+        if site_id is None:
             error = "Missing site ID. You must sign in first."
             raise NotSignedInError(error)
-        return self._site_id
+        return site_id
 
     @property
     def site_url(self):
-        if self._site_url is None:
+        site_url = self._auth_state.site_url
+        if site_url is None:
             error = "Missing site URL. You must sign in first."
             raise NotSignedInError(error)
-        return self._site_url
+        return site_url
 
     @property
     def user_id(self):
-        if self._user_id is None:
+        user_id = self._auth_state.user_id
+        if user_id is None:
             error = "Missing user ID. You must sign in first."
             raise NotSignedInError(error)
-        return self._user_id
+        return user_id
 
     @property
     def server_address(self):
@@ -329,11 +424,50 @@ 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._session_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
+
+    # Backwards-compatible access to the pre-thread-safety private attribute.
+    # Reading returns the current thread's session; assigning replaces the
+    # CURRENT thread's session only (other threads keep sessions created by
+    # session_factory), which preserves the common single-threaded pattern of
+    # injecting a prepared session before making calls. The injected session
+    # is registered so close() still reaches it.
+    @property
+    def _session(self) -> requests.Session:
+        return self.session
+
+    @_session.setter
+    def _session(self, value: requests.Session) -> None:
+        local = self._thread_sessions
+        with self._session_lock:
+            self._all_sessions.add(value)
+        local.session = value
+        local.epoch = self._session_epoch
 
     def is_signed_in(self):
-        return self._auth_token is not None
+        return self._auth_state.auth_token is not None
 
     def configure_ssl(self, *, allow_weak_dh=False):
         """Configure SSL/TLS settings for the server connection.
@@ -357,3 +491,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._session_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()