docs: samples for managing reservations (#144)

TODO: Testing these samples requires a capacity commitment. I created on manually on my dev project, but I think we'll want to programmatically create one in our CI project(s).

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-reservation/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #56  🦕

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)

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()

samples/snippets/requirements-test.txt

@@ -1 +1,2 @@
 pytest==6.2.4
+google-cloud-testutils==1.0.0

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

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]

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

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