feat: warn if location is set to unknown location (#609)

* feat: warn if location is set to unknown location

* tests error message

---------

Co-authored-by: Shobhit Singh <[email protected]>
Co-authored-by: Tim Sweña (Swast) <[email protected]>

Author: 3 people

bigframes/_config/bigquery_options.py

@@ -22,11 +22,33 @@
 import google.api_core.exceptions
 import google.auth.credentials
 
+import bigframes.constants
+import bigframes.exceptions
+
 SESSION_STARTED_MESSAGE = (
     "Cannot change '{attribute}' once a session has started. "
     "Call bigframes.pandas.close_session() first, if you are using the bigframes.pandas API."
 )
 
+UNKNOWN_LOCATION_MESSAGE = "The location '{location}' is set to an unknown value."
+
+
+def _validate_location(value: Optional[str]):
+
+    if value is None:
+        return
+
+    if value not in bigframes.constants.ALL_BIGQUERY_LOCATIONS:
+        warnings.warn(
+            UNKNOWN_LOCATION_MESSAGE.format(location=value),
+            # There are many layers before we get to (possibly) the user's code:
+            # -> bpd.options.bigquery.location = "us-central-1"
+            # -> location.setter
+            # -> _validate_location
+            stacklevel=3,
+            category=bigframes.exceptions.UnknownLocationWarning,
+        )
+
 
 class BigQueryOptions:
     """Encapsulates configuration for working with a session."""
@@ -93,6 +115,7 @@ def location(self) -> Optional[str]:
     def location(self, value: Optional[str]):
         if self._session_started and self._location != value:
             raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="location"))
+        _validate_location(value)
         self._location = value
 
     @property

bigframes/exceptions.py

@@ -0,0 +1,17 @@
+# 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.
+
+
+class UnknownLocationWarning(Warning):
+    """The location is set to an unknown value."""

tests/unit/_config/test_bigquery_options.py

@@ -13,10 +13,13 @@
 # limitations under the License.
 
 import re
+import warnings
 
 import pytest
 
+import bigframes
 import bigframes._config.bigquery_options as bigquery_options
+import bigframes.exceptions
 
 
 @pytest.mark.parametrize(
@@ -78,3 +81,51 @@ def test_setter_if_session_started_but_setting_the_same_value(attribute):
     setattr(options, attribute, original_object)
 
     assert getattr(options, attribute) is original_object
+
+
+@pytest.mark.parametrize(
+    [
+        "valid_location",
+    ],
+    [
+        (None,),
+        ("us-central1",),
+    ],
+)
+def test_location_set_to_valid_no_warning(valid_location):
+    options = bigquery_options.BigQueryOptions()
+    # Ensure that no warnings are emitted.
+    # https://docs.pytest.org/en/7.0.x/how-to/capture-warnings.html#additional-use-cases-of-warnings-in-tests
+    with warnings.catch_warnings():
+        # Turn matching UnknownLocationWarning into exceptions.
+        # https://docs.python.org/3/library/warnings.html#warning-filter
+        warnings.simplefilter(
+            "error", category=bigframes.exceptions.UnknownLocationWarning
+        )
+        options.location = valid_location
+
+
+@pytest.mark.parametrize(
+    [
+        "invalid_location",
+    ],
+    [
+        # Test with common mistakes, see article.
+        # https://en.wikipedia.org/wiki/Edit_distance#Formal_definition_and_properties
+        # Substitution
+        ("us-wist-3",),
+        # Insertion
+        ("us-central-1",),
+        # Deletion
+        ("asia-suth2",),
+    ],
+)
+def test_location_set_to_invalid_warning(invalid_location):
+    options = bigquery_options.BigQueryOptions()
+    with pytest.warns(
+        bigframes.exceptions.UnknownLocationWarning,
+        match=re.escape(
+            f"The location '{invalid_location}' is set to an unknown value."
+        ),
+    ):
+        options.location = invalid_location