Merge branch 'main' into owl-bot-update-lock-3b3a31be60853477bc39ed8d9bac162cac3ba083724cecaad54eb81d4e4dae9c

Author: parthea

google/cloud/bigquery/_job_helpers.py

@@ -39,7 +39,9 @@
 import functools
 import os
 import uuid
+import textwrap
 from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+import warnings
 
 import google.api_core.exceptions as core_exceptions
 from google.api_core import retry as retries
@@ -198,6 +200,44 @@ def _validate_job_config(request_body: Dict[str, Any], invalid_key: str):
         raise ValueError(f"got unexpected key {repr(invalid_key)} in job_config")
 
 
+def validate_job_retry(job_id: Optional[str], job_retry: Optional[retries.Retry]):
+    """Catch common mistakes, such as setting a job_id and job_retry at the same
+    time.
+    """
+    if job_id is not None and job_retry is not None:
+        # TODO(tswast): To avoid breaking changes but still allow a default
+        # query job retry, we currently only raise if they explicitly set a
+        # job_retry other than the default. In a future version, we may want to
+        # avoid this check for DEFAULT_JOB_RETRY and always raise.
+        if job_retry is not google.cloud.bigquery.retry.DEFAULT_JOB_RETRY:
+            raise TypeError(
+                textwrap.dedent(
+                    """
+                    `job_retry` was provided, but the returned job is
+                    not retryable, because a custom `job_id` was
+                    provided. To customize the job ID and allow for job
+                    retries, set job_id_prefix, instead.
+                    """
+                ).strip()
+            )
+        else:
+            warnings.warn(
+                textwrap.dedent(
+                    """
+                    job_retry must be explicitly set to None if job_id is set.
+                    BigQuery cannot retry a failed job by using the exact
+                    same ID. Setting job_id without explicitly disabling
+                    job_retry will raise an error in the future. To avoid this
+                    warning, either use job_id_prefix instead (preferred) or
+                    set job_retry=None.
+                    """
+                ).strip(),
+                category=FutureWarning,
+                # user code -> client.query / client.query_and_wait -> validate_job_retry
+                stacklevel=3,
+            )
+
+
 def _to_query_request(
     job_config: Optional[job.QueryJobConfig] = None,
     *,
@@ -308,7 +348,7 @@ def query_jobs_query(
     project: str,
     retry: retries.Retry,
     timeout: Optional[float],
-    job_retry: retries.Retry,
+    job_retry: Optional[retries.Retry],
 ) -> job.QueryJob:
     """Initiate a query using jobs.query with jobCreationMode=JOB_CREATION_REQUIRED.
 

google/cloud/bigquery/client.py

@@ -3388,7 +3388,7 @@ def query(
         project: Optional[str] = None,
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
-        job_retry: retries.Retry = DEFAULT_JOB_RETRY,
+        job_retry: Optional[retries.Retry] = DEFAULT_JOB_RETRY,
         api_method: Union[str, enums.QueryApiMethod] = enums.QueryApiMethod.INSERT,
     ) -> job.QueryJob:
         """Run a SQL query.
@@ -3455,18 +3455,9 @@ def query(
                 class, or if both ``job_id`` and non-``None`` non-default
                 ``job_retry`` are provided.
         """
-        job_id_given = job_id is not None
-        if (
-            job_id_given
-            and job_retry is not None
-            and job_retry is not DEFAULT_JOB_RETRY
-        ):
-            raise TypeError(
-                "`job_retry` was provided, but the returned job is"
-                " not retryable, because a custom `job_id` was"
-                " provided."
-            )
+        _job_helpers.validate_job_retry(job_id, job_retry)
 
+        job_id_given = job_id is not None
         if job_id_given and api_method == enums.QueryApiMethod.QUERY:
             raise TypeError(
                 "`job_id` was provided, but the 'QUERY' `api_method` was requested."

google/cloud/bigquery/dataset.py

@@ -298,12 +298,15 @@ def __init__(
         role: Optional[str] = None,
         entity_type: Optional[str] = None,
         entity_id: Optional[Union[Dict[str, Any], str]] = None,
+        **kwargs,
     ):
-        self._properties = {}
+        self._properties: Dict[str, Any] = {}
         if entity_type is not None:
             self._properties[entity_type] = entity_id
         self._properties["role"] = role
-        self._entity_type = entity_type
+        self._entity_type: Optional[str] = entity_type
+        for prop, val in kwargs.items():
+            setattr(self, prop, val)
 
     @property
     def role(self) -> Optional[str]:
@@ -330,6 +333,9 @@ def dataset(self, value):
         if isinstance(value, str):
             value = DatasetReference.from_string(value).to_api_repr()
 
+        if isinstance(value, DatasetReference):
+            value = value.to_api_repr()
+
         if isinstance(value, (Dataset, DatasetListItem)):
             value = value.reference.to_api_repr()
 
@@ -437,15 +443,65 @@ def special_group(self) -> Optional[str]:
     def special_group(self, value):
         self._properties["specialGroup"] = value
 
+    @property
+    def condition(self) -> Optional["Condition"]:
+        """Optional[Condition]: The IAM condition associated with this entry."""
+        value = typing.cast(Dict[str, Any], self._properties.get("condition"))
+        return Condition.from_api_repr(value) if value else None
+
+    @condition.setter
+    def condition(self, value: Union["Condition", dict, None]):
+        """Set the IAM condition for this entry."""
+        if value is None:
+            self._properties["condition"] = None
+        elif isinstance(value, Condition):
+            self._properties["condition"] = value.to_api_repr()
+        elif isinstance(value, dict):
+            self._properties["condition"] = value
+        else:
+            raise TypeError("condition must be a Condition object, dict, or None")
+
     @property
     def entity_type(self) -> Optional[str]:
         """The entity_type of the entry."""
+
+        # The api_repr for an AccessEntry object is expected to be a dict with
+        # only a few keys. Two keys that may be present are role and condition.
+        # Any additional key is going to have one of ~eight different names:
+        #   userByEmail, groupByEmail, domain, dataset, specialGroup, view,
+        #   routine, iamMember
+
+        # if self._entity_type is None, see if it needs setting
+        # i.e. is there a key: value pair that should be associated with
+        # entity_type and entity_id?
+        if self._entity_type is None:
+            resource = self._properties.copy()
+            # we are empyting the dict to get to the last `key: value`` pair
+            # so we don't keep these first entries
+            _ = resource.pop("role", None)
+            _ = resource.pop("condition", None)
+
+            try:
+                # we only need entity_type, because entity_id gets set elsewhere.
+                entity_type, _ = resource.popitem()
+            except KeyError:
+                entity_type = None
+
+            self._entity_type = entity_type
+
         return self._entity_type
 
     @property
     def entity_id(self) -> Optional[Union[Dict[str, Any], str]]:
         """The entity_id of the entry."""
-        return self._properties.get(self._entity_type) if self._entity_type else None
+        if self.entity_type:
+            entity_type = self.entity_type
+        else:
+            return None
+        return typing.cast(
+            Optional[Union[Dict[str, Any], str]],
+            self._properties.get(entity_type, None),
+        )
 
     def __eq__(self, other):
         if not isinstance(other, AccessEntry):
@@ -464,7 +520,16 @@ def _key(self):
         Returns:
             Tuple: The contents of this :class:`~google.cloud.bigquery.dataset.AccessEntry`.
         """
+
         properties = self._properties.copy()
+
+        # Dicts are not hashable.
+        # Convert condition to a hashable datatype(s)
+        condition = properties.get("condition")
+        if isinstance(condition, dict):
+            condition_key = tuple(sorted(condition.items()))
+            properties["condition"] = condition_key
+
         prop_tup = tuple(sorted(properties.items()))
         return (self.role, self._entity_type, self.entity_id, prop_tup)
 
@@ -491,19 +556,11 @@ def from_api_repr(cls, resource: dict) -> "AccessEntry":
         Returns:
             google.cloud.bigquery.dataset.AccessEntry:
                 Access entry parsed from ``resource``.
-
-        Raises:
-            ValueError:
-                If the resource has more keys than ``role`` and one additional
-                key.
         """
-        entry = resource.copy()
-        role = entry.pop("role", None)
-        entity_type, entity_id = entry.popitem()
-        if len(entry) != 0:
-            raise ValueError("Entry has unexpected keys remaining.", entry)
 
-        return cls(role, entity_type, entity_id)
+        access_entry = cls()
+        access_entry._properties = resource.copy()
+        return access_entry
 
 
 class Dataset(object):
@@ -1160,6 +1217,43 @@ def from_api_repr(cls, resource: Dict[str, Any]) -> "Condition":
 
         return cls(
             expression=resource["expression"],
-            title=resource.get("title"),
-            description=resource.get("description"),
+            title=resource.get("title", None),
+            description=resource.get("description", None),
         )
+
+    def __eq__(self, other: object) -> bool:
+        """Check for equality based on expression, title, and description."""
+        if not isinstance(other, Condition):
+            return NotImplemented
+        return self._key() == other._key()
+
+    def _key(self):
+        """A tuple key that uniquely describes this field.
+        Used to compute this instance's hashcode and evaluate equality.
+        Returns:
+            Tuple: The contents of this :class:`~google.cloud.bigquery.dataset.AccessEntry`.
+        """
+
+        properties = self._properties.copy()
+
+        # Dicts are not hashable.
+        # Convert object to a hashable datatype(s)
+        prop_tup = tuple(sorted(properties.items()))
+        return prop_tup
+
+    def __ne__(self, other: object) -> bool:
+        """Check for inequality."""
+        return not self == other
+
+    def __hash__(self) -> int:
+        """Generate a hash based on expression, title, and description."""
+        return hash(self._key())
+
+    def __repr__(self) -> str:
+        """Return a string representation of the Condition object."""
+        parts = [f"expression={self.expression!r}"]
+        if self.title is not None:
+            parts.append(f"title={self.title!r}")
+        if self.description is not None:
+            parts.append(f"description={self.description!r}")
+        return f"Condition({', '.join(parts)})"

google/cloud/bigquery/enums.py

@@ -387,3 +387,19 @@ def _generate_next_value_(name, start, count, last_values):
     ROUNDING_MODE_UNSPECIFIED = enum.auto()
     ROUND_HALF_AWAY_FROM_ZERO = enum.auto()
     ROUND_HALF_EVEN = enum.auto()
+
+
+class BigLakeFileFormat(object):
+    FILE_FORMAT_UNSPECIFIED = "FILE_FORMAT_UNSPECIFIED"
+    """The default unspecified value."""
+
+    PARQUET = "PARQUET"
+    """Apache Parquet format."""
+
+
+class BigLakeTableFormat(object):
+    TABLE_FORMAT_UNSPECIFIED = "TABLE_FORMAT_UNSPECIFIED"
+    """The default unspecified value."""
+
+    ICEBERG = "ICEBERG"
+    """Apache Iceberg format."""