Adds unit tests, updates helpers, revises several client attrs

Author: chalmerlowe

google/cloud/bigquery_v2/services/centralized_service/_helpers.py

@@ -22,3 +22,31 @@ def _drop_self_key(kwargs):
         raise TypeError("kwargs must be a dict.")
     kwargs.pop("self", None)  # Essentially a no-op if 'self' key does not exist
     return kwargs
+
+
+def _make_request(
+    request_class,
+    user_request,
+    identifier_value,
+    identifier_name: str,
+    parser,
+    identifier_required: bool = True,
+):
+    if user_request is not None and identifier_value is not None:
+        raise ValueError(
+            f"Provide either a request object or '{identifier_name}', not both."
+        )
+
+    if user_request is not None:
+        return user_request
+
+    if identifier_required and identifier_value is None:
+        raise ValueError(
+            f"Either a request object or '{identifier_name}' must be provided."
+        )
+
+    if identifier_value is None:
+        request_fields = parser()
+    else:
+        request_fields = parser(identifier_value)
+    return request_class(**request_fields)

google/cloud/bigquery_v2/services/centralized_service/client.py

@@ -16,6 +16,7 @@
 
 import os
 from typing import (
+    Dict,
     Optional,
     Sequence,
     Tuple,
@@ -34,6 +35,7 @@
 # Import types modules (to access *Requests classes)
 from google.cloud.bigquery_v2.types import (
     dataset,
+    dataset_reference,
     job,
     model,
 )
@@ -43,151 +45,182 @@
 from google.api_core import retry as retries
 from google.auth import credentials as auth_credentials
 
-# Create a type alias
+# Create type aliases
 try:
     OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
 except AttributeError:  # pragma: NO COVER
     OptionalRetry = Union[retries.Retry, object, None]  # type: ignore
 
-# TODO: This line is here to simplify prototyping, etc.
-PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
+DatasetIdentifier = Union[str, dataset_reference.DatasetReference]
 
+# TODO: This variable is here to simplify prototyping, etc.
+PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
 DEFAULT_RETRY: OptionalRetry = gapic_v1.method.DEFAULT
 DEFAULT_TIMEOUT: Union[float, object] = gapic_v1.method.DEFAULT
 DEFAULT_METADATA: Sequence[Tuple[str, Union[str, bytes]]] = ()
 
-
 # Create Centralized Client
 class BigQueryClient:
+    """A centralized client for BigQuery API."""
+
     def __init__(
         self,
         *,
         credentials: Optional[auth_credentials.Credentials] = None,
         client_options: Optional[Union[client_options_lib.ClientOptions, dict]] = None,
     ):
-        self._clients = {}
+        """
+        Initializes the BigQueryClient.
+
+        Args:
+            credentials:
+                The credentials to use for authentication. If not provided, the
+                client will attempt to use the default credentials.
+            client_options:
+                A dictionary of client options to pass to the underlying
+                service clients.
+        """
+
+        self._clients: Dict[str, object] = {}
         self._credentials = credentials
         self._client_options = client_options
+        self.project = PROJECT_ID
+
+    # --- HELPER METHODS ---
+    def _parse_dataset_path(self, dataset_path: str) -> Tuple[Optional[str], str]:
+        """
+        Helper to parse project and dataset from a string identifier.
+
+        Args:
+            dataset_path: A string in the format 'project_id.dataset_id' or
+                'dataset_id'.
+
+        Returns:
+            A tuple of (project_id, dataset_id).
+        """
+        if "." in dataset_path:
+            project_id, dataset_id = dataset_path.split(".", 1)
+            return project_id, dataset_id
+        return self.project, dataset_path
+
+    def _parse_dataset_id_to_dict(self, dataset_id: DatasetIdentifier) -> dict:
+        if isinstance(dataset_id, str):
+            project_id, dataset_id_str = self._parse_dataset_path(dataset_id)
+            return {"project_id": project_id, "dataset_id": dataset_id_str}
+        elif isinstance(dataset_id, dataset_reference.DatasetReference):
+            return {
+                "project_id": dataset_id.project_id,
+                "dataset_id": dataset_id.dataset_id,
+            }
+        else:
+            raise TypeError(f"Invalid type for dataset_id: {type(dataset_id)}")
 
-# --- SERVICE CLIENT ATTRIBUTES ---
+    def _parse_project_id_to_dict(self, project_id: Optional[str] = None) -> dict:
+        """Helper to create a request dictionary from a project_id."""
+        final_project_id = project_id or self.project
+        return {"project_id": final_project_id}
 
+    # --- *SERVICECLIENT ATTRIBUTES ---
     @property
     def dataset_service_client(self):
         if "dataset" not in self._clients:
-            from google.cloud.bigquery_v2.services import dataset_service
-
             self._clients["dataset"] = dataset_service.DatasetServiceClient(
                 credentials=self._credentials, client_options=self._client_options
             )
         return self._clients["dataset"]
 
     @dataset_service_client.setter
     def dataset_service_client(self, value):
-        # Check for the methods the centralized client exposes (to allow duck-typing)
-        required_methods = [
-            "get_dataset",
-            "insert_dataset",
-            "patch_dataset",
-            "update_dataset",
-            "delete_dataset",
-            "list_datasets",
-            "undelete_dataset",
-        ]
-        for method in required_methods:
-            if not hasattr(value, method) or not callable(getattr(value, method)):
-                raise AttributeError(
-                    f"Object assigned to dataset_service_client is missing a callable '{method}' method."
-                )
+        if not isinstance(value, dataset_service.DatasetServiceClient):
+            raise TypeError(
+                "Expected an instance of dataset_service.DatasetServiceClient."
+            )
         self._clients["dataset"] = value
 
     @property
     def job_service_client(self):
         if "job" not in self._clients:
-            from google.cloud.bigquery_v2.services import job_service
-
             self._clients["job"] = job_service.JobServiceClient(
                 credentials=self._credentials, client_options=self._client_options
             )
         return self._clients["job"]
 
     @job_service_client.setter
     def job_service_client(self, value):
-        required_methods = [
-            "get_job",
-            "insert_job",
-            "cancel_job",
-            "delete_job",
-            "list_jobs",
-        ]
-        for method in required_methods:
-            if not hasattr(value, method) or not callable(getattr(value, method)):
-                raise AttributeError(
-                    f"Object assigned to job_service_client is missing a callable '{method}' method."
-                )
+        if not isinstance(value, job_service.JobServiceClient):
+            raise TypeError("Expected an instance of job_service.JobServiceClient.")
         self._clients["job"] = value
 
     @property
     def model_service_client(self):
         if "model" not in self._clients:
-            from google.cloud.bigquery_v2.services import model_service
-
             self._clients["model"] = model_service.ModelServiceClient(
                 credentials=self._credentials, client_options=self._client_options
             )
         return self._clients["model"]
 
     @model_service_client.setter
     def model_service_client(self, value):
-        required_methods = [
-            "get_model",
-            "delete_model",
-            "patch_model",
-            "list_models",
-        ]
-        for method in required_methods:
-            if not hasattr(value, method) or not callable(getattr(value, method)):
-                raise AttributeError(
-                    f"Object assigned to model_service_client is missing a callable '{method}' method."
-                )
+        if not isinstance(value, model_service.ModelServiceClient):
+            raise TypeError("Expected an instance of model_service.ModelServiceClient.")
         self._clients["model"] = value
 
-# --- SERVICE CLIENT METHODS ---
-# TODO: refactor the microgenerator template so that everything related to 
-#       a single ServiceClient is kept close togetehr in the same section of this class:
-#       @property
-#       @setter
-#       _method_A()
-#       _method_B()
-#       _method_C()
-#       etc
-
+    # --- *SERVICECLIENT METHODS ---
     def get_dataset(
         self,
-        request: Optional[Union[dataset.GetDatasetRequest, dict]] = None,
+        dataset_id: Optional[DatasetIdentifier] = None,
         *,
+        request: Optional["dataset.GetDatasetRequest"] = None,
         retry: OptionalRetry = DEFAULT_RETRY,
         timeout: Union[float, object] = DEFAULT_TIMEOUT,
         metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
-    ):
+    ) -> "dataset.Dataset":
         """
         TODO: Docstring is purposefully blank. microgenerator will add automatically.
         """
-        kwargs = _helpers._drop_self_key(locals())
-        return self.dataset_service_client.get_dataset(**kwargs)
+        final_request = _helpers._make_request(
+            request_class=dataset.GetDatasetRequest,
+            user_request=request,
+            identifier_value=dataset_id,
+            identifier_name="dataset_id",
+            parser=self._parse_dataset_id_to_dict,
+            identifier_required=True,
+        )
+
+        return self.dataset_service_client.get_dataset(
+            request=final_request,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
 
     def list_datasets(
         self,
-        request: Optional[Union[dataset.ListDatasetsRequest, dict]] = None,
+        project_id: Optional[str] = None,
         *,
+        request: Optional["dataset.ListDatasetsRequest"] = None,
         retry: OptionalRetry = DEFAULT_RETRY,
         timeout: Union[float, object] = DEFAULT_TIMEOUT,
         metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
     ):
         """
         TODO: Docstring is purposefully blank. microgenerator will add automatically.
         """
-        kwargs = _helpers._drop_self_key(locals())
-        return self.dataset_service_client.list_datasets(**kwargs)
+        final_request = _helpers._make_request(
+            request_class=dataset.ListDatasetsRequest,
+            user_request=request,
+            identifier_value=project_id,
+            identifier_name="project_id",
+            parser=self._parse_project_id_to_dict,
+            identifier_required=False,
+        )
+
+        return self.dataset_service_client.list_datasets(
+            request=final_request,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
 
     def list_jobs(
         self,