googleapis · tswast · Apr 18, 2024 · Apr 12, 2024 · Apr 15, 2024 · Apr 15, 2024
@@ -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

@@ -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."""
@@ -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