diff --git a/samples/snippets/conftest.py b/samples/snippets/conftest.py new file mode 100644 index 00000000..59db3c08 --- /dev/null +++ b/samples/snippets/conftest.py @@ -0,0 +1,60 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +import os +import time + +from google.cloud.bigquery_reservation_v1.services import reservation_service +from google.cloud.bigquery_reservation_v1.types import reservation as reservation_types +import pytest + + +@pytest.fixture(scope="session") +def project_id() -> str: + return os.environ["GOOGLE_CLOUD_PROJECT"] + + +@pytest.fixture(scope="session") +def reservation_client() -> reservation_service.ReservationServiceClient: + return reservation_service.ReservationServiceClient() + + +@pytest.fixture(scope="session") +def location() -> str: + return "US" + + +@pytest.fixture(scope="session") +def location_path(project_id: str, location: str) -> str: + return reservation_service.ReservationServiceClient.common_location_path( + project_id, location + ) + + +@pytest.fixture(scope="session", autouse=True) +def capacity_commitment(location_path: str, reservation_client: reservation_service.ReservationServiceClient) -> reservation_types.CapacityCommitment: + # TODO(b/196082966): If custom names or creation date property are added, + # do pre-test cleanup of past commitments. + commitment = reservation_types.CapacityCommitment() + commitment.slot_count = 100 + commitment.plan = reservation_types.CapacityCommitment.CommitmentPlan.FLEX + commitment = reservation_client.create_capacity_commitment(parent=location_path, capacity_commitment=commitment) + yield commitment + # Commitments can only be removed after 1 minute. + now = datetime.datetime.now(datetime.timezone.utc) + delta = commitment.commitment_end_time - now + sleep_seconds = max(0, delta.total_seconds()) + 5 + time.sleep(sleep_seconds) + reservation_client.delete_capacity_commitment(name=commitment.name) diff --git a/samples/snippets/quickstart_test.py b/samples/snippets/quickstart_test.py index 66dd3cd8..3ade3e40 100644 --- a/samples/snippets/quickstart_test.py +++ b/samples/snippets/quickstart_test.py @@ -12,18 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os - import pytest from . import quickstart -@pytest.fixture() -def project_id() -> str: - return os.environ["GOOGLE_CLOUD_PROJECT"] - - def test_quickstart(capsys: pytest.CaptureFixture, project_id: str) -> None: quickstart.main(project_id) out, _ = capsys.readouterr() diff --git a/samples/snippets/requirements-test.txt b/samples/snippets/requirements-test.txt index 95ea1e6a..2ff95fe0 100644 --- a/samples/snippets/requirements-test.txt +++ b/samples/snippets/requirements-test.txt @@ -1 +1,2 @@ pytest==6.2.4 +google-cloud-testutils==1.0.0 diff --git a/samples/snippets/reservation_create.py b/samples/snippets/reservation_create.py new file mode 100644 index 00000000..67955606 --- /dev/null +++ b/samples/snippets/reservation_create.py @@ -0,0 +1,66 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google.cloud.bigquery_reservation_v1.types import reservation as reservation_types + + +def create_reservation( + project_id: str, location: str, reservation_id: str, slot_capacity: str +) -> reservation_types.Reservation: + original_project_id = project_id + original_location = location + original_reservation_id = reservation_id + original_slot_capacity = slot_capacity + + # [START bigqueryreservation_reservation_create] + # TODO(developer): Set project_id to the project ID containing the + # reservation. + project_id = "your-project-id" + + # TODO(developer): Set location to the location of the reservation. + # See: https://cloud.google.com/bigquery/docs/locations for a list of + # available locations. + location = "US" + + # TODO(developer): Set reservation_id to a unique ID of the reservation. + reservation_id = "sample-reservation" + + # TODO(developer): Set slot_capicity to the number of slots in the + # reservation. + slot_capacity = 100 + + # [START_EXCLUDE] + project_id = original_project_id + location = original_location + reservation_id = original_reservation_id + slot_capacity = original_slot_capacity + # [END_EXCLUDE] + + from google.cloud.bigquery_reservation_v1.services import reservation_service + from google.cloud.bigquery_reservation_v1.types import ( + reservation as reservation_types, + ) + + reservation_client = reservation_service.ReservationServiceClient() + + parent = reservation_client.common_location_path(project_id, location) + + reservation = reservation_types.Reservation(slot_capacity=slot_capacity) + reservation = reservation_client.create_reservation( + parent=parent, reservation=reservation, reservation_id=reservation_id, + ) + + print(f"Created reservation: {reservation.name}") + # [END bigqueryreservation_reservation_create] + return reservation diff --git a/samples/snippets/reservation_delete.py b/samples/snippets/reservation_delete.py new file mode 100644 index 00000000..b0537b5e --- /dev/null +++ b/samples/snippets/reservation_delete.py @@ -0,0 +1,49 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def delete_reservation(project_id: str, location: str, reservation_id: str) -> None: + original_project_id = project_id + original_location = location + original_reservation_id = reservation_id + + # [START bigqueryreservation_reservation_delete] + # TODO(developer): Set project_id to the project ID containing the + # reservation. + project_id = "your-project-id" + + # TODO(developer): Set location to the location of the reservation. + # See: https://cloud.google.com/bigquery/docs/locations for a list of + # available locations. + location = "US" + + # TODO(developer): Set reservation_id to a unique ID of the reservation. + reservation_id = "sample-reservation" + + # [START_EXCLUDE] + project_id = original_project_id + location = original_location + reservation_id = original_reservation_id + # [END_EXCLUDE] + + from google.cloud.bigquery_reservation_v1.services import reservation_service + + reservation_client = reservation_service.ReservationServiceClient() + reservation_name = reservation_client.reservation_path( + project_id, location, reservation_id + ) + reservation_client.delete_reservation(name=reservation_name) + + print(f"Deleted reservation: {reservation_name}") + # [END bigqueryreservation_reservation_delete] diff --git a/samples/snippets/reservation_test.py b/samples/snippets/reservation_test.py new file mode 100644 index 00000000..442e2582 --- /dev/null +++ b/samples/snippets/reservation_test.py @@ -0,0 +1,81 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import google.api_core.exceptions +from google.cloud.bigquery_reservation_v1.services import reservation_service +import pytest +import test_utils.prefixer + +from . import reservation_create +from . import reservation_delete +from . import reservation_update + + +# Reservation IDs are limited to 64 characters. +reservation_prefixer = test_utils.prefixer.Prefixer( + "py-bq-r", "snippets", separator="-" +) + + +@pytest.fixture(scope="module", autouse=True) +def cleanup_reservations( + reservation_client: reservation_service.ReservationServiceClient, location_path: str +) -> None: + for reservation in reservation_client.list_reservations(parent=location_path): + reservation_id = reservation.name.split("/")[-1] + if reservation_prefixer.should_cleanup(reservation_id): + reservation_client.delete_reservation(name=reservation.name) + + +@pytest.fixture(scope="session") +def reservation_id( + reservation_client: reservation_service.ReservationServiceClient, + project_id: str, + location: str, +) -> str: + id_ = reservation_prefixer.create_prefix() + yield id_ + + reservation_name = reservation_client.reservation_path(project_id, location, id_) + try: + reservation_client.delete_reservation(name=reservation_name) + except google.api_core.exceptions.NotFound: + pass + + +def test_reservation_samples( + capsys: pytest.CaptureFixture, project_id: str, location: str, reservation_id: str +) -> None: + slot_capacity = 100 + reservation = reservation_create.create_reservation( + project_id, location, reservation_id, slot_capacity + ) + assert reservation.slot_capacity == 100 + assert reservation_id in reservation.name + out, _ = capsys.readouterr() + assert f"Created reservation: {reservation.name}" in out + + slot_capacity = 50 + reservation = reservation_update.update_reservation( + project_id, location, reservation_id, slot_capacity + ) + assert reservation.slot_capacity == 50 + assert reservation_id in reservation.name + out, _ = capsys.readouterr() + assert f"Updated reservation: {reservation.name}" in out + + reservation_delete.delete_reservation(project_id, location, reservation_id) + out, _ = capsys.readouterr() + assert "Deleted reservation" in out + assert reservation_id in out diff --git a/samples/snippets/reservation_update.py b/samples/snippets/reservation_update.py new file mode 100644 index 00000000..93a9cf53 --- /dev/null +++ b/samples/snippets/reservation_update.py @@ -0,0 +1,72 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google.cloud.bigquery_reservation_v1.types import reservation as reservation_types + + +def update_reservation( + project_id: str, location: str, reservation_id: str, slot_capacity: str +) -> reservation_types.Reservation: + original_project_id = project_id + original_location = location + original_reservation_id = reservation_id + original_slot_capacity = slot_capacity + + # [START bigqueryreservation_reservation_update] + # TODO(developer): Set project_id to the project ID containing the + # reservation. + project_id = "your-project-id" + + # TODO(developer): Set location to the location of the reservation. + # See: https://cloud.google.com/bigquery/docs/locations for a list of + # available locations. + location = "US" + + # TODO(developer): Set reservation_id to a unique ID of the reservation. + reservation_id = "sample-reservation" + + # TODO(developer): Set slot_capicity to the new number of slots in the + # reservation. + slot_capacity = 50 + + # [START_EXCLUDE] + project_id = original_project_id + location = original_location + reservation_id = original_reservation_id + slot_capacity = original_slot_capacity + # [END_EXCLUDE] + + from google.cloud.bigquery_reservation_v1.services import reservation_service + from google.cloud.bigquery_reservation_v1.types import ( + reservation as reservation_types, + ) + from google.protobuf import field_mask_pb2 + + reservation_client = reservation_service.ReservationServiceClient() + + reservation_name = reservation_client.reservation_path( + project_id, location, reservation_id + ) + reservation = reservation_types.Reservation( + name=reservation_name, slot_capacity=slot_capacity, + ) + field_mask = field_mask_pb2.FieldMask(paths=["slot_capacity"]) + reservation = reservation_client.update_reservation( + reservation=reservation, update_mask=field_mask + ) + + print(f"Updated reservation: {reservation.name}") + print(f"\tslot_capacity: {reservation.slot_capacity}") + # [END bigqueryreservation_reservation_update] + return reservation