From bea0c1b1e2e45878557a6317580c8f3a8810e4eb Mon Sep 17 00:00:00 2001
From: Trevor Bergeron <tbergeron@google.com>
Date: Mon, 3 Jun 2024 19:09:49 +0000
Subject: [PATCH] fix: Resolve issue with unset thread-local options

---
 bigframes/_config/__init__.py               | 44 +++++++++++----------
 tests/unit/_config/test_threaded_options.py | 41 +++++++++++++++++++
 2 files changed, 64 insertions(+), 21 deletions(-)
 create mode 100644 tests/unit/_config/test_threaded_options.py

diff --git a/bigframes/_config/__init__.py b/bigframes/_config/__init__.py
index 4729532e98..c9b2a3f95a 100644
--- a/bigframes/_config/__init__.py
+++ b/bigframes/_config/__init__.py
@@ -17,8 +17,12 @@
 DataFrames from this package.
 """
 
+from __future__ import annotations
+
 import copy
+from dataclasses import dataclass, field
 import threading
+from typing import Optional
 
 import bigframes_vendored.pandas._config.config as pandas_config
 
@@ -28,18 +32,27 @@
 import bigframes._config.sampling_options as sampling_options
 
 
+@dataclass
+class ThreadLocalConfig(threading.local):
+    # If unset, global settings will be used
+    bigquery_options: Optional[bigquery_options.BigQueryOptions] = None
+    # Note: use default factory instead of default instance so each thread initializes to default values
+    display_options: display_options.DisplayOptions = field(
+        default_factory=display_options.DisplayOptions
+    )
+    sampling_options: sampling_options.SamplingOptions = field(
+        default_factory=sampling_options.SamplingOptions
+    )
+    compute_options: compute_options.ComputeOptions = field(
+        default_factory=compute_options.ComputeOptions
+    )
+
+
 class Options:
     """Global options affecting BigQuery DataFrames behavior."""
 
     def __init__(self):
-        self._local = threading.local()
-
-        # Initialize these in the property getters to make sure we do have a
-        # separate instance per thread.
-        self._local.bigquery_options = None
-        self._local.display_options = None
-        self._local.sampling_options = None
-        self._local.compute_options = None
+        self._local = ThreadLocalConfig()
 
         # BigQuery options are special because they can only be set once per
         # session, so we need an indicator as to whether we are using the
@@ -61,21 +74,16 @@ def _init_bigquery_thread_local(self):
     @property
     def bigquery(self) -> bigquery_options.BigQueryOptions:
         """Options to use with the BigQuery engine."""
-        if (
-            bigquery_options := getattr(self._local, "bigquery_options", None)
-        ) is not None:
+        if self._local.bigquery_options is not None:
             # The only way we can get here is if someone called
             # _init_bigquery_thread_local.
-            return bigquery_options
+            return self._local.bigquery_options
 
         return self._bigquery_options
 
     @property
     def display(self) -> display_options.DisplayOptions:
         """Options controlling object representation."""
-        if self._local.display_options is None:
-            self._local.display_options = display_options.DisplayOptions()
-
         return self._local.display_options
 
     @property
@@ -88,17 +96,11 @@ def sampling(self) -> sampling_options.SamplingOptions:
         matplotlib plotting). This option can be overriden by
         parameters in specific functions.
         """
-        if self._local.sampling_options is None:
-            self._local.sampling_options = sampling_options.SamplingOptions()
-
         return self._local.sampling_options
 
     @property
     def compute(self) -> compute_options.ComputeOptions:
         """Thread-local options controlling object computation."""
-        if self._local.compute_options is None:
-            self._local.compute_options = compute_options.ComputeOptions()
-
         return self._local.compute_options
 
     @property
diff --git a/tests/unit/_config/test_threaded_options.py b/tests/unit/_config/test_threaded_options.py
new file mode 100644
index 0000000000..7fc97a9f72
--- /dev/null
+++ b/tests/unit/_config/test_threaded_options.py
@@ -0,0 +1,41 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import threading
+
+import bigframes._config
+
+
+def test_mutate_options_threaded():
+    options = bigframes._config.Options()
+    options.display.max_rows = 50
+    result_dict = {"this_before": options.display.max_rows}
+
+    def mutate_options_threaded(options, result_dict):
+        result_dict["other_before"] = options.display.max_rows
+
+        options.display.max_rows = 100
+        result_dict["other_after"] = options.display.max_rows
+
+    thread = threading.Thread(
+        target=(lambda: mutate_options_threaded(options, result_dict))
+    )
+    thread.start()
+    thread.join(1)
+    result_dict["this_after"] = options.display.max_rows
+
+    assert result_dict["this_before"] == 50
+    assert result_dict["this_after"] == 50
+    assert result_dict["other_before"] == 25
+    assert result_dict["other_after"] == 100