From c7899684653ce1c693304ad9e37c0da7f91b3e56 Mon Sep 17 00:00:00 2001 From: jialuo Date: Wed, 16 Jul 2025 22:44:25 +0000 Subject: [PATCH 1/4] fix: resolve location reset issue in bigquery options --- bigframes/_config/bigquery_options.py | 2 +- tests/unit/_config/test_bigquery_options.py | 28 +++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index 09ffee95d4..648b69dea7 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -171,7 +171,7 @@ def location(self) -> Optional[str]: @location.setter def location(self, value: Optional[str]): - if self._session_started and self._location != value: + if self._session_started and self._location != _get_validated_location(value): raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="location")) self._location = _get_validated_location(value) diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index 686499aa75..5125983aed 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/tests/unit/_config/test_bigquery_options.py @@ -22,6 +22,7 @@ import bigframes import bigframes._config.bigquery_options as bigquery_options import bigframes.exceptions +import bigframes.pandas as bpd @pytest.mark.parametrize( @@ -179,6 +180,33 @@ def set_location_property(): assert possibility in str(w[0].message).replace("\n", "") +def test_location_set_same_location_after_session_starts(): + bpd.options.bigquery.location = "us" + assert bpd.options.bigquery.location == "US" + + # Start a session when creating a dataframe. + bpd.DataFrame({"a": [0]}) + + bpd.options.bigquery.location = "us" + assert bpd.options.bigquery.location == "US" + + +def test_location_set_different_location_after_session_starts(): + bpd.options.bigquery.location = "us" + # Start a session when creating a dataframe. + bpd.DataFrame({"a": [0]}) + + with pytest.raises( + ValueError, + match=( + "Cannot change 'location' once a session has started. Call " + "bigframes.pandas.close_session\\(\\) first, if you are using the " + "bigframes.pandas API." + ), + ): + bpd.options.bigquery.location = "us-central1" + + def test_client_endpoints_override_set_shows_warning(): options = bigquery_options.BigQueryOptions() From 1c847b6745c7e361826e412cdec65309468c9b63 Mon Sep 17 00:00:00 2001 From: jialuo Date: Wed, 16 Jul 2025 23:08:45 +0000 Subject: [PATCH 2/4] fix --- tests/unit/_config/test_bigquery_options.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index 5125983aed..b5fea84a57 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/tests/unit/_config/test_bigquery_options.py @@ -181,6 +181,7 @@ def set_location_property(): def test_location_set_same_location_after_session_starts(): + bpd.close_session() bpd.options.bigquery.location = "us" assert bpd.options.bigquery.location == "US" @@ -192,6 +193,7 @@ def test_location_set_same_location_after_session_starts(): def test_location_set_different_location_after_session_starts(): + bpd.close_session() bpd.options.bigquery.location = "us" # Start a session when creating a dataframe. bpd.DataFrame({"a": [0]}) From d017cee9685fcf491823e81dfd150d0f5b58610b Mon Sep 17 00:00:00 2001 From: jialuo Date: Wed, 16 Jul 2025 23:25:43 +0000 Subject: [PATCH 3/4] fix --- tests/unit/_config/test_bigquery_options.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index b5fea84a57..e8e20e6d78 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/tests/unit/_config/test_bigquery_options.py @@ -185,8 +185,7 @@ def test_location_set_same_location_after_session_starts(): bpd.options.bigquery.location = "us" assert bpd.options.bigquery.location == "US" - # Start a session when creating a dataframe. - bpd.DataFrame({"a": [0]}) + bpd.get_global_session() bpd.options.bigquery.location = "us" assert bpd.options.bigquery.location == "US" @@ -195,8 +194,8 @@ def test_location_set_same_location_after_session_starts(): def test_location_set_different_location_after_session_starts(): bpd.close_session() bpd.options.bigquery.location = "us" - # Start a session when creating a dataframe. - bpd.DataFrame({"a": [0]}) + + bpd.get_global_session() with pytest.raises( ValueError, From f6b737812c2944f824d099e772453b92b99687af Mon Sep 17 00:00:00 2001 From: jialuo Date: Wed, 16 Jul 2025 23:59:33 +0000 Subject: [PATCH 4/4] fix the test --- tests/unit/_config/test_bigquery_options.py | 41 ++++++--------------- 1 file changed, 12 insertions(+), 29 deletions(-) diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index e8e20e6d78..3c80f00a37 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/tests/unit/_config/test_bigquery_options.py @@ -22,7 +22,6 @@ import bigframes import bigframes._config.bigquery_options as bigquery_options import bigframes.exceptions -import bigframes.pandas as bpd @pytest.mark.parametrize( @@ -59,6 +58,18 @@ def test_setter_raises_if_session_started(attribute, original_value, new_value): assert getattr(options, attribute) is not new_value +def test_location_set_us_twice(): + """This test ensures the fix for b/423220936 is working as expected.""" + options = bigquery_options.BigQueryOptions() + setattr(options, "location", "us") + assert getattr(options, "location") == "US" + + options._session_started = True + + setattr(options, "location", "us") + assert getattr(options, "location") == "US" + + @pytest.mark.parametrize( [ "attribute", @@ -180,34 +191,6 @@ def set_location_property(): assert possibility in str(w[0].message).replace("\n", "") -def test_location_set_same_location_after_session_starts(): - bpd.close_session() - bpd.options.bigquery.location = "us" - assert bpd.options.bigquery.location == "US" - - bpd.get_global_session() - - bpd.options.bigquery.location = "us" - assert bpd.options.bigquery.location == "US" - - -def test_location_set_different_location_after_session_starts(): - bpd.close_session() - bpd.options.bigquery.location = "us" - - bpd.get_global_session() - - with pytest.raises( - ValueError, - match=( - "Cannot change 'location' once a session has started. Call " - "bigframes.pandas.close_session\\(\\) first, if you are using the " - "bigframes.pandas API." - ), - ): - bpd.options.bigquery.location = "us-central1" - - def test_client_endpoints_override_set_shows_warning(): options = bigquery_options.BigQueryOptions()