Merge branch 'main' into pangea-v1alpha

Author: chalmerlowe

google/cloud/bigquery/client.py

@@ -44,6 +44,8 @@
 import uuid
 import warnings
 
+import requests
+
 from google import resumable_media  # type: ignore
 from google.resumable_media.requests import MultipartUpload  # type: ignore
 from google.resumable_media.requests import ResumableUpload
@@ -65,6 +67,7 @@
     DEFAULT_BQSTORAGE_CLIENT_INFO = None  # type: ignore
 
 
+from google.auth.credentials import Credentials
 from google.cloud.bigquery._http import Connection
 from google.cloud.bigquery import _job_helpers
 from google.cloud.bigquery import _pandas_helpers
@@ -126,15 +129,14 @@
     _versions_helpers.PANDAS_VERSIONS.try_import()
 )  # mypy check fails because pandas import is outside module, there are type: ignore comments related to this
 
+
 ResumableTimeoutType = Union[
     None, float, Tuple[float, float]
 ]  # for resumable media methods
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
     # os.PathLike is only subscriptable in Python 3.9+, thus shielding with a condition.
     PathType = Union[str, bytes, os.PathLike[str], os.PathLike[bytes]]
-    import requests  # required by api-core
-
 _DEFAULT_CHUNKSIZE = 100 * 1024 * 1024  # 100 MB
 _MAX_MULTIPART_SIZE = 5 * 1024 * 1024
 _DEFAULT_NUM_RETRIES = 6
@@ -231,30 +233,34 @@ class Client(ClientWithProject):
 
     def __init__(
         self,
-        project=None,
-        credentials=None,
-        _http=None,
-        location=None,
-        default_query_job_config=None,
-        default_load_job_config=None,
-        client_info=None,
-        client_options=None,
+        project: Optional[str] = None,
+        credentials: Optional[Credentials] = None,
+        _http: Optional[requests.Session] = None,
+        location: Optional[str] = None,
+        default_query_job_config: Optional[QueryJobConfig] = None,
+        default_load_job_config: Optional[LoadJobConfig] = None,
+        client_info: Optional[google.api_core.client_info.ClientInfo] = None,
+        client_options: Optional[
+            Union[google.api_core.client_options.ClientOptions, Dict[str, Any]]
+        ] = None,
     ) -> None:
+        if client_options is None:
+            client_options = {}
+        if isinstance(client_options, dict):
+            client_options = google.api_core.client_options.from_dict(client_options)
+        # assert isinstance(client_options, google.api_core.client_options.ClientOptions)
+
         super(Client, self).__init__(
             project=project,
             credentials=credentials,
             client_options=client_options,
             _http=_http,
         )
 
-        kw_args = {"client_info": client_info}
+        kw_args: Dict[str, Any] = {"client_info": client_info}
         bq_host = _get_bigquery_host()
         kw_args["api_endpoint"] = bq_host if bq_host != _DEFAULT_HOST else None
         client_universe = None
-        if client_options is None:
-            client_options = {}
-        if isinstance(client_options, dict):
-            client_options = google.api_core.client_options.from_dict(client_options)
         if client_options.api_endpoint:
             api_endpoint = client_options.api_endpoint
             kw_args["api_endpoint"] = api_endpoint

google/cloud/bigquery/table.py

@@ -410,6 +410,7 @@ class Table(_TableBase):
         "require_partition_filter": "requirePartitionFilter",
         "table_constraints": "tableConstraints",
         "external_catalog_table_options": "externalCatalogTableOptions",
+        "max_staleness": "maxStaleness",
     }
 
     def __init__(self, table_ref, schema=None) -> None:
@@ -1140,6 +1141,40 @@ def __repr__(self):
     def __str__(self):
         return f"{self.project}.{self.dataset_id}.{self.table_id}"
 
+    @property
+    def max_staleness(self):
+        """Union[str, None]: The maximum staleness of data that could be returned when the table is queried.
+
+        Staleness encoded as a string encoding of sql IntervalValue type.
+        This property is optional and defaults to None.
+
+        According to the BigQuery API documentation, maxStaleness specifies the maximum time
+        interval for which stale data can be returned when querying the table.
+        It helps control data freshness in scenarios like metadata-cached external tables.
+
+        Returns:
+            Optional[str]: A string representing the maximum staleness interval
+            (e.g., '1h', '30m', '15s' for hours, minutes, seconds respectively).
+        """
+        return self._properties.get(self._PROPERTY_TO_API_FIELD["max_staleness"])
+
+    @max_staleness.setter
+    def max_staleness(self, value):
+        """Set the maximum staleness for the table.
+
+        Args:
+            value (Optional[str]): A string representing the maximum staleness interval.
+                Must be a valid time interval string.
+                Examples include '1h' (1 hour), '30m' (30 minutes), '15s' (15 seconds).
+
+        Raises:
+            ValueError: If the value is not None and not a string.
+        """
+        if value is not None and not isinstance(value, str):
+            raise ValueError("max_staleness must be a string or None")
+
+        self._properties[self._PROPERTY_TO_API_FIELD["max_staleness"]] = value
+
 
 class TableListItem(_TableBase):
     """A read-only table resource from a list operation.

tests/unit/test_table.py

@@ -1476,6 +1476,49 @@ def test___str__(self):
         table1 = self._make_one(TableReference(dataset, "table1"))
         self.assertEqual(str(table1), "project1.dataset1.table1")
 
+    def test_max_staleness_getter(self):
+        """Test getting max_staleness property."""
+        dataset = DatasetReference("test-project", "test_dataset")
+        table_ref = dataset.table("test_table")
+        table = self._make_one(table_ref)
+        # Initially None
+        self.assertIsNone(table.max_staleness)
+        # Set max_staleness using setter
+        table.max_staleness = "1h"
+        self.assertEqual(table.max_staleness, "1h")
+
+    def test_max_staleness_setter(self):
+        """Test setting max_staleness property."""
+        dataset = DatasetReference("test-project", "test_dataset")
+        table_ref = dataset.table("test_table")
+        table = self._make_one(table_ref)
+        # Set valid max_staleness
+        table.max_staleness = "30m"
+        self.assertEqual(table.max_staleness, "30m")
+        # Set to None
+        table.max_staleness = None
+        self.assertIsNone(table.max_staleness)
+
+    def test_max_staleness_setter_invalid_type(self):
+        """Test setting max_staleness with an invalid type raises ValueError."""
+        dataset = DatasetReference("test-project", "test_dataset")
+        table_ref = dataset.table("test_table")
+        table = self._make_one(table_ref)
+        # Try setting invalid type
+        with self.assertRaises(ValueError):
+            table.max_staleness = 123  # Not a string
+
+    def test_max_staleness_to_api_repr(self):
+        """Test max_staleness is correctly represented in API representation."""
+        dataset = DatasetReference("test-project", "test_dataset")
+        table_ref = dataset.table("test_table")
+        table = self._make_one(table_ref)
+        # Set max_staleness
+        table.max_staleness = "1h"
+        # Convert to API representation
+        resource = table.to_api_repr()
+        self.assertEqual(resource.get("maxStaleness"), "1h")
+
 
 class Test_row_from_mapping(unittest.TestCase, _SchemaBase):
     PROJECT = "prahj-ekt"