Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: use the project id from BQ dataset instead of the default project id #717

Merged
merged 4 commits into from
Sep 28, 2021
Merged

fix: use the project id from BQ dataset instead of the default project id #717

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
428e683
fix: use the project id from BQ dataset instead of the default projec…
geraint0923 Sep 17, 2021
29d9092
Merge branch 'main' into fix_bq
geraint0923 Sep 22, 2021
7803ab3
Merge branch 'main' into fix_bq
geraint0923 Sep 23, 2021
90c3238
Merge branch 'main' into fix_bq
geraint0923 Sep 28, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions google/cloud/aiplatform/jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -795,17 +795,14 @@ def iter_outputs(
if bq_dataset.startswith("bq://"):
bq_dataset = bq_dataset[5:]

# # Split project ID and BQ dataset ID
_, bq_dataset_id = bq_dataset.split(".", 1)

# Build a BigQuery Client using the same credentials as JobServiceClient
bq_client = bigquery.Client(
project=self.project,
credentials=self.api_client._transport._credentials,
)

row_iterator = bq_client.list_rows(
table=f"{bq_dataset_id}.{bq_table}", max_results=bq_max_results
table=f"{bq_dataset}.{bq_table}", max_results=bq_max_results
)

return row_iterator
Expand Down
5 changes: 3 additions & 2 deletions tests/unit/aiplatform/test_jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,13 +48,14 @@
_TEST_ID = "1028944691210842416"
_TEST_ALT_ID = "8834795523125638878"
_TEST_DISPLAY_NAME = "my_job_1234"
_TEST_BQ_PROJECT_ID = "projectId"
_TEST_BQ_DATASET_ID = "bqDatasetId"
_TEST_BQ_TABLE_NAME = "someBqTable"
_TEST_BQ_JOB_ID = "123459876"
_TEST_BQ_MAX_RESULTS = 100
_TEST_GCS_BUCKET_NAME = "my-bucket"

_TEST_BQ_PATH = f"bq://projectId.{_TEST_BQ_DATASET_ID}"
_TEST_BQ_PATH = f"bq://{_TEST_BQ_PROJECT_ID}.{_TEST_BQ_DATASET_ID}"
_TEST_GCS_BUCKET_PATH = f"gs://{_TEST_GCS_BUCKET_NAME}"
_TEST_GCS_JSONL_SOURCE_URI = f"{_TEST_GCS_BUCKET_PATH}/bp_input_config.jsonl"
_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}"
Expand Down Expand Up @@ -420,7 +421,7 @@ def test_batch_prediction_iter_dirs_bq(self, bq_list_rows_mock):
bp.iter_outputs()

bq_list_rows_mock.assert_called_once_with(
table=f"{_TEST_BQ_DATASET_ID}.{_TEST_BQ_TABLE_NAME}",
table=f"{_TEST_BQ_PROJECT_ID}.{_TEST_BQ_DATASET_ID}.{_TEST_BQ_TABLE_NAME}",
max_results=_TEST_BQ_MAX_RESULTS,
)

Expand Down