fix: pass credentials to BQ and GCS clients (#469)

Author: sasha-gitg

google/cloud/aiplatform/datasets/dataset.py

@@ -68,7 +68,7 @@ def __init__(
                 Optional location to retrieve dataset from. If not set, location
                 set in aiplatform.init will be used.
             credentials (auth_credentials.Credentials):
-                Custom credentials to use to upload this model. Overrides
+                Custom credentials to use to retreive this Dataset. Overrides
                 credentials set in aiplatform.init.
         """
 

google/cloud/aiplatform/datasets/tabular_dataset.py

@@ -73,20 +73,28 @@ def column_names(self) -> List[str]:
                 gcs_source_uris.sort()
 
                 # Get the first file in sorted list
-                return TabularDataset._retrieve_gcs_source_columns(
-                    self.project, gcs_source_uris[0]
+                return self._retrieve_gcs_source_columns(
+                    project=self.project,
+                    gcs_csv_file_path=gcs_source_uris[0],
+                    credentials=self.credentials,
                 )
         elif bq_source:
             bq_table_uri = bq_source.get("uri")
             if bq_table_uri:
-                return TabularDataset._retrieve_bq_source_columns(
-                    self.project, bq_table_uri
+                return self._retrieve_bq_source_columns(
+                    project=self.project,
+                    bq_table_uri=bq_table_uri,
+                    credentials=self.credentials,
                 )
 
         raise RuntimeError("No valid CSV or BigQuery datasource found.")
 
     @staticmethod
-    def _retrieve_gcs_source_columns(project: str, gcs_csv_file_path: str) -> List[str]:
+    def _retrieve_gcs_source_columns(
+        project: str,
+        gcs_csv_file_path: str,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> List[str]:
         """Retrieve the columns from a comma-delimited CSV file stored on Google Cloud Storage
 
         Example Usage:
@@ -104,7 +112,8 @@ def _retrieve_gcs_source_columns(project: str, gcs_csv_file_path: str) -> List[s
             gcs_csv_file_path (str):
                 Required. A full path to a CSV files stored on Google Cloud Storage.
                 Must include "gs://" prefix.
-
+            credentials (auth_credentials.Credentials):
+                Credentials to use to with GCS Client.
         Returns:
             List[str]
                 A list of columns names in the CSV file.
@@ -116,7 +125,7 @@ def _retrieve_gcs_source_columns(project: str, gcs_csv_file_path: str) -> List[s
         gcs_bucket, gcs_blob = utils.extract_bucket_and_prefix_from_gcs_path(
             gcs_csv_file_path
         )
-        client = storage.Client(project=project)
+        client = storage.Client(project=project, credentials=credentials)
         bucket = client.bucket(gcs_bucket)
         blob = bucket.blob(gcs_blob)
 
@@ -135,6 +144,7 @@ def _retrieve_gcs_source_columns(project: str, gcs_csv_file_path: str) -> List[s
                 line += blob.download_as_bytes(
                     start=start_index, end=start_index + increment
                 ).decode("utf-8")
+
                 first_new_line_index = line.find("\n")
                 start_index += increment
 
@@ -156,7 +166,11 @@ def _retrieve_gcs_source_columns(project: str, gcs_csv_file_path: str) -> List[s
         return next(csv_reader)
 
     @staticmethod
-    def _retrieve_bq_source_columns(project: str, bq_table_uri: str) -> List[str]:
+    def _retrieve_bq_source_columns(
+        project: str,
+        bq_table_uri: str,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> List[str]:
         """Retrieve the columns from a table on Google BigQuery
 
         Example Usage:
@@ -174,6 +188,8 @@ def _retrieve_bq_source_columns(project: str, bq_table_uri: str) -> List[str]:
             bq_table_uri (str):
                 Required. A URI to a BigQuery table.
                 Can include "bq://" prefix but not required.
+            credentials (auth_credentials.Credentials):
+                Credentials to use with BQ Client.
 
         Returns:
             List[str]
@@ -185,7 +201,7 @@ def _retrieve_bq_source_columns(project: str, bq_table_uri: str) -> List[str]:
         if bq_table_uri.startswith(prefix):
             bq_table_uri = bq_table_uri[len(prefix) :]
 
-        client = bigquery.Client(project=project)
+        client = bigquery.Client(project=project, credentials=credentials)
         table = client.get_table(bq_table_uri)
         schema = table.schema
         return [schema.name for schema in schema]

tests/unit/aiplatform/test_datasets.py

@@ -341,16 +341,30 @@ def list_datasets_mock():
 
 @pytest.fixture
 def gcs_client_download_as_bytes_mock():
-    with patch.object(storage.Blob, "download_as_bytes") as bigquery_blob_mock:
-        bigquery_blob_mock.return_value = b'"column_1","column_2"\n0, 1'
-        yield bigquery_blob_mock
+    with patch.object(storage.Blob, "download_as_bytes") as gcs_blob_mock:
+        gcs_blob_mock.return_value = b'"column_1","column_2"\n0, 1'
+        yield gcs_blob_mock
 
 
 @pytest.fixture
-def bigquery_client_mock():
-    with patch.object(bigquery.Client, "get_table") as bigquery_client_mock:
-        bigquery_client_mock.return_value = bigquery.Table("project.dataset.table")
-        yield bigquery_client_mock
+def gcs_client_mock():
+    with patch.object(storage, "Client") as client_mock:
+        yield client_mock
+
+
+@pytest.fixture
+def bq_client_mock():
+    with patch.object(bigquery, "Client") as client_mock:
+        yield client_mock
+
+
+@pytest.fixture
+def bigquery_client_table_mock():
+    with patch.object(bigquery.Client, "get_table") as bigquery_client_table_mock:
+        bigquery_client_table_mock.return_value = bigquery.Table(
+            "project.dataset.table"
+        )
+        yield bigquery_client_table_mock
 
 
 @pytest.fixture
@@ -995,9 +1009,37 @@ def test_tabular_dataset_column_name_gcs(self):
 
         assert my_dataset.column_names == ["column_1", "column_2"]
 
+    @pytest.mark.usefixtures("get_dataset_tabular_gcs_mock")
+    def test_tabular_dataset_column_name_gcs_with_creds(self, gcs_client_mock):
+        creds = auth_credentials.AnonymousCredentials()
+        my_dataset = datasets.TabularDataset(dataset_name=_TEST_NAME, credentials=creds)
+
+        # we are just testing creds passing
+        # this exception if from the mock not returning
+        # the csv data which is tested above
+        try:
+            my_dataset.column_names
+        except StopIteration:
+            pass
+
+        gcs_client_mock.assert_called_once_with(
+            project=_TEST_PROJECT, credentials=creds
+        )
+
+    @pytest.mark.usefixtures("get_dataset_tabular_bq_mock",)
+    def test_tabular_dataset_column_name_bq_with_creds(self, bq_client_mock):
+        creds = auth_credentials.AnonymousCredentials()
+        my_dataset = datasets.TabularDataset(dataset_name=_TEST_NAME, credentials=creds)
+
+        my_dataset.column_names
+
+        assert bq_client_mock.call_args_list[0] == mock.call(
+            project=_TEST_PROJECT, credentials=creds
+        )
+
     @pytest.mark.usefixtures(
         "get_dataset_tabular_bq_mock",
-        "bigquery_client_mock",
+        "bigquery_client_table_mock",
         "bigquery_table_schema_mock",
     )
     def test_tabular_dataset_column_name_bigquery(self):