Skip to content

Commit a42733e

Browse files
tswasttswast
authored
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 🦕
1 parent 44acfd5 commit a42733e

File tree

7 files changed

+329
-7
lines changed

7 files changed

+329
-7
lines changed

‎bigquery-reservation/snippets/conftest.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import datetime
16+
import os
17+
import time
18+
19+
from google.cloud.bigquery_reservation_v1.services import reservation_service
20+
from google.cloud.bigquery_reservation_v1.types import reservation as reservation_types
21+
import pytest
22+
23+
24+
@pytest.fixture(scope="session")
25+
def project_id() -> str:
26+
return os.environ["GOOGLE_CLOUD_PROJECT"]
27+
28+
29+
@pytest.fixture(scope="session")
30+
def reservation_client() -> reservation_service.ReservationServiceClient:
31+
return reservation_service.ReservationServiceClient()
32+
33+
34+
@pytest.fixture(scope="session")
35+
def location() -> str:
36+
return "US"
37+
38+
39+
@pytest.fixture(scope="session")
40+
def location_path(project_id: str, location: str) -> str:
41+
return reservation_service.ReservationServiceClient.common_location_path(
42+
project_id, location
43+
)
44+
45+
46+
@pytest.fixture(scope="session", autouse=True)
47+
def capacity_commitment(location_path: str, reservation_client: reservation_service.ReservationServiceClient) -> reservation_types.CapacityCommitment:
48+
# TODO(b/196082966): If custom names or creation date property are added,
49+
# do pre-test cleanup of past commitments.
50+
commitment = reservation_types.CapacityCommitment()
51+
commitment.slot_count = 100
52+
commitment.plan = reservation_types.CapacityCommitment.CommitmentPlan.FLEX
53+
commitment = reservation_client.create_capacity_commitment(parent=location_path, capacity_commitment=commitment)
54+
yield commitment
55+
# Commitments can only be removed after 1 minute.
56+
now = datetime.datetime.now(datetime.timezone.utc)
57+
delta = commitment.commitment_end_time - now
58+
sleep_seconds = max(0, delta.total_seconds()) + 5
59+
time.sleep(sleep_seconds)
60+
reservation_client.delete_capacity_commitment(name=commitment.name)

‎bigquery-reservation/snippets/quickstart_test.py‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,18 +12,11 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
import os
16-
1715
import pytest
1816

1917
from . import quickstart
2018

2119

22-
@pytest.fixture()
23-
def project_id() -> str:
24-
return os.environ["GOOGLE_CLOUD_PROJECT"]
25-
26-
2720
def test_quickstart(capsys: pytest.CaptureFixture, project_id: str) -> None:
2821
quickstart.main(project_id)
2922
out, _ = capsys.readouterr()

‎bigquery-reservation/snippets/requirements-test.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
pytest==6.2.4
2+
google-cloud-testutils==1.0.0

‎bigquery-reservation/snippets/reservation_create.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from google.cloud.bigquery_reservation_v1.types import reservation as reservation_types
16+
17+
18+
def create_reservation(
19+
project_id: str, location: str, reservation_id: str, slot_capacity: str
20+
) -> reservation_types.Reservation:
21+
original_project_id = project_id
22+
original_location = location
23+
original_reservation_id = reservation_id
24+
original_slot_capacity = slot_capacity
25+
26+
# [START bigqueryreservation_reservation_create]
27+
# TODO(developer): Set project_id to the project ID containing the
28+
# reservation.
29+
project_id = "your-project-id"
30+
31+
# TODO(developer): Set location to the location of the reservation.
32+
# See: https://cloud.google.com/bigquery/docs/locations for a list of
33+
# available locations.
34+
location = "US"
35+
36+
# TODO(developer): Set reservation_id to a unique ID of the reservation.
37+
reservation_id = "sample-reservation"
38+
39+
# TODO(developer): Set slot_capicity to the number of slots in the
40+
# reservation.
41+
slot_capacity = 100
42+
43+
# [START_EXCLUDE]
44+
project_id = original_project_id
45+
location = original_location
46+
reservation_id = original_reservation_id
47+
slot_capacity = original_slot_capacity
48+
# [END_EXCLUDE]
49+
50+
from google.cloud.bigquery_reservation_v1.services import reservation_service
51+
from google.cloud.bigquery_reservation_v1.types import (
52+
reservation as reservation_types,
53+
)
54+
55+
reservation_client = reservation_service.ReservationServiceClient()
56+
57+
parent = reservation_client.common_location_path(project_id, location)
58+
59+
reservation = reservation_types.Reservation(slot_capacity=slot_capacity)
60+
reservation = reservation_client.create_reservation(
61+
parent=parent, reservation=reservation, reservation_id=reservation_id,
62+
)
63+
64+
print(f"Created reservation: {reservation.name}")
65+
# [END bigqueryreservation_reservation_create]
66+
return reservation

‎bigquery-reservation/snippets/reservation_delete.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def delete_reservation(project_id: str, location: str, reservation_id: str) -> None:
17+
original_project_id = project_id
18+
original_location = location
19+
original_reservation_id = reservation_id
20+
21+
# [START bigqueryreservation_reservation_delete]
22+
# TODO(developer): Set project_id to the project ID containing the
23+
# reservation.
24+
project_id = "your-project-id"
25+
26+
# TODO(developer): Set location to the location of the reservation.
27+
# See: https://cloud.google.com/bigquery/docs/locations for a list of
28+
# available locations.
29+
location = "US"
30+
31+
# TODO(developer): Set reservation_id to a unique ID of the reservation.
32+
reservation_id = "sample-reservation"
33+
34+
# [START_EXCLUDE]
35+
project_id = original_project_id
36+
location = original_location
37+
reservation_id = original_reservation_id
38+
# [END_EXCLUDE]
39+
40+
from google.cloud.bigquery_reservation_v1.services import reservation_service
41+
42+
reservation_client = reservation_service.ReservationServiceClient()
43+
reservation_name = reservation_client.reservation_path(
44+
project_id, location, reservation_id
45+
)
46+
reservation_client.delete_reservation(name=reservation_name)
47+
48+
print(f"Deleted reservation: {reservation_name}")
49+
# [END bigqueryreservation_reservation_delete]

‎bigquery-reservation/snippets/reservation_test.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import google.api_core.exceptions
16+
from google.cloud.bigquery_reservation_v1.services import reservation_service
17+
import pytest
18+
import test_utils.prefixer
19+
20+
from . import reservation_create
21+
from . import reservation_delete
22+
from . import reservation_update
23+
24+
25+
# Reservation IDs are limited to 64 characters.
26+
reservation_prefixer = test_utils.prefixer.Prefixer(
27+
"py-bq-r", "snippets", separator="-"
28+
)
29+
30+
31+
@pytest.fixture(scope="module", autouse=True)
32+
def cleanup_reservations(
33+
reservation_client: reservation_service.ReservationServiceClient, location_path: str
34+
) -> None:
35+
for reservation in reservation_client.list_reservations(parent=location_path):
36+
reservation_id = reservation.name.split("/")[-1]
37+
if reservation_prefixer.should_cleanup(reservation_id):
38+
reservation_client.delete_reservation(name=reservation.name)
39+
40+
41+
@pytest.fixture(scope="session")
42+
def reservation_id(
43+
reservation_client: reservation_service.ReservationServiceClient,
44+
project_id: str,
45+
location: str,
46+
) -> str:
47+
id_ = reservation_prefixer.create_prefix()
48+
yield id_
49+
50+
reservation_name = reservation_client.reservation_path(project_id, location, id_)
51+
try:
52+
reservation_client.delete_reservation(name=reservation_name)
53+
except google.api_core.exceptions.NotFound:
54+
pass
55+
56+
57+
def test_reservation_samples(
58+
capsys: pytest.CaptureFixture, project_id: str, location: str, reservation_id: str
59+
) -> None:
60+
slot_capacity = 100
61+
reservation = reservation_create.create_reservation(
62+
project_id, location, reservation_id, slot_capacity
63+
)
64+
assert reservation.slot_capacity == 100
65+
assert reservation_id in reservation.name
66+
out, _ = capsys.readouterr()
67+
assert f"Created reservation: {reservation.name}" in out
68+
69+
slot_capacity = 50
70+
reservation = reservation_update.update_reservation(
71+
project_id, location, reservation_id, slot_capacity
72+
)
73+
assert reservation.slot_capacity == 50
74+
assert reservation_id in reservation.name
75+
out, _ = capsys.readouterr()
76+
assert f"Updated reservation: {reservation.name}" in out
77+
78+
reservation_delete.delete_reservation(project_id, location, reservation_id)
79+
out, _ = capsys.readouterr()
80+
assert "Deleted reservation" in out
81+
assert reservation_id in out

‎bigquery-reservation/snippets/reservation_update.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from google.cloud.bigquery_reservation_v1.types import reservation as reservation_types
16+
17+
18+
def update_reservation(
19+
project_id: str, location: str, reservation_id: str, slot_capacity: str
20+
) -> reservation_types.Reservation:
21+
original_project_id = project_id
22+
original_location = location
23+
original_reservation_id = reservation_id
24+
original_slot_capacity = slot_capacity
25+
26+
# [START bigqueryreservation_reservation_update]
27+
# TODO(developer): Set project_id to the project ID containing the
28+
# reservation.
29+
project_id = "your-project-id"
30+
31+
# TODO(developer): Set location to the location of the reservation.
32+
# See: https://cloud.google.com/bigquery/docs/locations for a list of
33+
# available locations.
34+
location = "US"
35+
36+
# TODO(developer): Set reservation_id to a unique ID of the reservation.
37+
reservation_id = "sample-reservation"
38+
39+
# TODO(developer): Set slot_capicity to the new number of slots in the
40+
# reservation.
41+
slot_capacity = 50
42+
43+
# [START_EXCLUDE]
44+
project_id = original_project_id
45+
location = original_location
46+
reservation_id = original_reservation_id
47+
slot_capacity = original_slot_capacity
48+
# [END_EXCLUDE]
49+
50+
from google.cloud.bigquery_reservation_v1.services import reservation_service
51+
from google.cloud.bigquery_reservation_v1.types import (
52+
reservation as reservation_types,
53+
)
54+
from google.protobuf import field_mask_pb2
55+
56+
reservation_client = reservation_service.ReservationServiceClient()
57+
58+
reservation_name = reservation_client.reservation_path(
59+
project_id, location, reservation_id
60+
)
61+
reservation = reservation_types.Reservation(
62+
name=reservation_name, slot_capacity=slot_capacity,
63+
)
64+
field_mask = field_mask_pb2.FieldMask(paths=["slot_capacity"])
65+
reservation = reservation_client.update_reservation(
66+
reservation=reservation, update_mask=field_mask
67+
)
68+
69+
print(f"Updated reservation: {reservation.name}")
70+
print(f"\tslot_capacity: {reservation.slot_capacity}")
71+
# [END bigqueryreservation_reservation_update]
72+
return reservation

0 commit comments

Comments
 (0)