1196 lines
36 KiB
Python
1196 lines
36 KiB
Python
#!/usr/bin/env python
|
|
"""The BigQuery CLI reservation client library."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
from typing import Any, Dict, NamedTuple, Optional, Set, Tuple
|
|
|
|
from googleapiclient import discovery
|
|
|
|
from clients import utils as bq_client_utils
|
|
from frontend import utils as frontend_utils
|
|
from utils import bq_error
|
|
from utils import bq_id_utils
|
|
|
|
|
|
def GetBodyForCreateReservation(
|
|
api_version: str, # pylint: disable=unused-argument Cleanup is error prone.
|
|
slots: int,
|
|
ignore_idle_slots: bool,
|
|
edition,
|
|
target_job_concurrency: Optional[int],
|
|
multi_region_auxiliary: Optional[bool],
|
|
autoscale_max_slots: Optional[int] = None,
|
|
max_slots: Optional[int] = None,
|
|
scaling_mode: Optional[str] = None,
|
|
reservation_group_name: Optional[str] = None,
|
|
scheduling_policy_concurrency: Optional[int] = None,
|
|
scheduling_policy_max_slots: Optional[int] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Return the request body for CreateReservation.
|
|
|
|
Arguments:
|
|
api_version: The api version to make the request against.
|
|
slots: Number of slots allocated to this reservation subtree.
|
|
ignore_idle_slots: Specifies whether queries should ignore idle slots from
|
|
other reservations.
|
|
edition: The edition for this reservation.
|
|
target_job_concurrency: Job concurrency target.
|
|
multi_region_auxiliary: Whether this reservation is for the auxiliary
|
|
region.
|
|
autoscale_max_slots: Number of slots to be scaled when needed.
|
|
max_slots: The overall max slots for the reservation.
|
|
scaling_mode: The scaling mode for the reservation.
|
|
reservation_group_name: The reservation group name which the reservation
|
|
belongs to.
|
|
scheduling_policy_concurrency: Soft cap on project concurrency within the
|
|
reservation.
|
|
scheduling_policy_max_slots: Soft cap on project max slots within the
|
|
reservation.
|
|
|
|
Returns:
|
|
Reservation object that was created.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if requirements for parameters are not met.
|
|
"""
|
|
reservation = {}
|
|
reservation['slot_capacity'] = slots
|
|
reservation['ignore_idle_slots'] = ignore_idle_slots
|
|
if multi_region_auxiliary is not None:
|
|
reservation['multi_region_auxiliary'] = multi_region_auxiliary
|
|
if target_job_concurrency is not None:
|
|
reservation['concurrency'] = target_job_concurrency
|
|
|
|
if autoscale_max_slots is not None:
|
|
reservation['autoscale'] = {}
|
|
reservation['autoscale']['max_slots'] = autoscale_max_slots
|
|
|
|
if edition is not None:
|
|
reservation['edition'] = edition
|
|
|
|
if frontend_utils.ValidateAtMostOneSelected(max_slots, autoscale_max_slots):
|
|
raise bq_error.BigqueryError(
|
|
'max_slots and autoscale_max_slots cannot be set at the same time.'
|
|
)
|
|
# make sure max_slots and scaling_mode are set at the same time.
|
|
if (max_slots is not None and scaling_mode is None) or (
|
|
max_slots is None and scaling_mode is not None
|
|
):
|
|
raise bq_error.BigqueryError(
|
|
'max_slots and scaling_mode must be set at the same time.'
|
|
)
|
|
|
|
if max_slots is not None:
|
|
reservation['max_slots'] = max_slots
|
|
if scaling_mode is not None:
|
|
reservation['scaling_mode'] = scaling_mode
|
|
|
|
if reservation_group_name is not None:
|
|
reservation['reservation_group'] = reservation_group_name
|
|
|
|
if scheduling_policy_concurrency is not None:
|
|
reservation['scheduling_policy'] = {}
|
|
reservation['scheduling_policy'][
|
|
'concurrency'
|
|
] = scheduling_policy_concurrency
|
|
if scheduling_policy_max_slots is not None:
|
|
if 'scheduling_policy' not in reservation:
|
|
reservation['scheduling_policy'] = {}
|
|
reservation['scheduling_policy']['max_slots'] = scheduling_policy_max_slots
|
|
|
|
return reservation
|
|
|
|
|
|
def GetBodyForCreateReservationAssignment(
|
|
reference: str, # pylint: disable=unused-argument Cleanup is error prone.
|
|
job_type: Optional[str],
|
|
priority: Optional[str],
|
|
assignee_type: str,
|
|
assignee_id: str,
|
|
scheduling_policy_max_slots: Optional[int] = None,
|
|
scheduling_policy_concurrency: Optional[int] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Creates a reservation assignment for a given project/folder/organization.
|
|
|
|
Arguments:
|
|
reference: Reference to the project reservation is assigned. Location must
|
|
be the same location as the reservation.
|
|
job_type: Type of jobs for this assignment.
|
|
priority: Default job priority for this assignment.
|
|
assignee_type: Type of assignees for the reservation assignment.
|
|
assignee_id: Project/folder/organization ID, to which the reservation is
|
|
assigned.
|
|
scheduling_policy_max_slots: Max slots for the scheduling policy.
|
|
scheduling_policy_concurrency: Concurrency for the scheduling policy.
|
|
|
|
Returns:
|
|
ReservationAssignment object that was created.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if assignment cannot be created.
|
|
"""
|
|
reservation_assignment = {}
|
|
if not job_type:
|
|
raise_job_type_unspecified_error = True
|
|
if (
|
|
scheduling_policy_max_slots is not None
|
|
or scheduling_policy_concurrency is not None
|
|
):
|
|
raise_job_type_unspecified_error = False
|
|
if raise_job_type_unspecified_error:
|
|
raise bq_error.BigqueryError('job_type not specified.')
|
|
else:
|
|
reservation_assignment['job_type'] = job_type
|
|
if priority:
|
|
reservation_assignment['priority'] = priority
|
|
if not assignee_type:
|
|
raise bq_error.BigqueryError('assignee_type not specified.')
|
|
if not assignee_id:
|
|
raise bq_error.BigqueryError('assignee_id not specified.')
|
|
# assignee_type is singular, that's why we need additional 's' inside
|
|
# format string for assignee below.
|
|
reservation_assignment['assignee'] = '%ss/%s' % (
|
|
assignee_type.lower(),
|
|
assignee_id,
|
|
)
|
|
if scheduling_policy_max_slots is not None:
|
|
if 'scheduling_policy' not in reservation_assignment:
|
|
reservation_assignment['scheduling_policy'] = {}
|
|
reservation_assignment['scheduling_policy'][
|
|
'max_slots'
|
|
] = scheduling_policy_max_slots
|
|
if scheduling_policy_concurrency is not None:
|
|
if 'scheduling_policy' not in reservation_assignment:
|
|
reservation_assignment['scheduling_policy'] = {}
|
|
reservation_assignment['scheduling_policy'][
|
|
'concurrency'
|
|
] = scheduling_policy_concurrency
|
|
|
|
return reservation_assignment
|
|
|
|
|
|
def CreateReservation(
|
|
client,
|
|
api_version: str,
|
|
reference,
|
|
slots: int,
|
|
ignore_idle_slots: bool,
|
|
edition,
|
|
target_job_concurrency: Optional[int],
|
|
multi_region_auxiliary: Optional[bool],
|
|
autoscale_max_slots: Optional[int] = None,
|
|
max_slots: Optional[int] = None,
|
|
scaling_mode: Optional[str] = None,
|
|
reservation_group_name: Optional[str] = None,
|
|
scheduling_policy_concurrency: Optional[int] = None,
|
|
scheduling_policy_max_slots: Optional[int] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Create a reservation with the given reservation reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
api_version: The api version to make the request against.
|
|
reference: Reservation to create.
|
|
slots: Number of slots allocated to this reservation subtree.
|
|
ignore_idle_slots: Specifies whether queries should ignore idle slots from
|
|
other reservations.
|
|
edition: The edition for this reservation.
|
|
target_job_concurrency: Job concurrency target.
|
|
multi_region_auxiliary: Whether this reservation is for the auxiliary
|
|
region.
|
|
autoscale_max_slots: Number of slots to be scaled when needed.
|
|
max_slots: The overall max slots for the reservation.
|
|
scaling_mode: The scaling mode for the reservation.
|
|
reservation_group_name: The reservation group name which the reservation
|
|
belongs to.
|
|
|
|
Returns:
|
|
Reservation object that was created.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if autoscale_max_slots is used with other
|
|
version.
|
|
"""
|
|
reservation = GetBodyForCreateReservation(
|
|
api_version,
|
|
slots,
|
|
ignore_idle_slots,
|
|
edition,
|
|
target_job_concurrency,
|
|
multi_region_auxiliary,
|
|
autoscale_max_slots,
|
|
max_slots,
|
|
scaling_mode,
|
|
reservation_group_name,
|
|
scheduling_policy_concurrency,
|
|
scheduling_policy_max_slots,
|
|
)
|
|
parent = 'projects/%s/locations/%s' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.create(
|
|
parent=parent, body=reservation, reservationId=reference.reservationId
|
|
)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def ListReservations(
|
|
client: ...,
|
|
reference: ...,
|
|
page_size: int,
|
|
page_token: Optional[str],
|
|
reservation_group: Optional[str] = None,
|
|
) -> Any:
|
|
"""List reservations in the project and location for the given reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reservation reference containing project and location.
|
|
page_size: Number of results to show.
|
|
page_token: Token to retrieve the next page of results.
|
|
reservation_group: When specified, only list reservations in the given
|
|
reservation group.
|
|
|
|
Returns:
|
|
Reservation object that was created.
|
|
"""
|
|
parent = 'projects/%s/locations/%s' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
if reservation_group is not None:
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.list(
|
|
parent=parent,
|
|
pageSize=page_size,
|
|
pageToken=page_token,
|
|
filter='reservation_group=%s' % reservation_group,
|
|
)
|
|
.execute()
|
|
)
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.list(parent=parent, pageSize=page_size, pageToken=page_token)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def ListBiReservations(client, reference):
|
|
"""List BI reservations in the project and location for the given reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reservation reference containing project and location.
|
|
|
|
Returns:
|
|
List of BI reservations in the given project/location.
|
|
"""
|
|
parent = 'projects/%s/locations/%s/biReservation' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
response = (
|
|
client.projects().locations().getBiReservation(name=parent).execute()
|
|
)
|
|
return response
|
|
|
|
|
|
def GetReservation(client, reference):
|
|
"""Gets a reservation with the given reservation reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reservation to get.
|
|
|
|
Returns:
|
|
Reservation object corresponding to the given id.
|
|
"""
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.get(name=reference.path())
|
|
.execute()
|
|
)
|
|
|
|
|
|
def DeleteReservation(
|
|
client,
|
|
reference: ...
|
|
):
|
|
"""Deletes a reservation with the given reservation reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reservation to delete.
|
|
"""
|
|
client.projects().locations().reservations().delete(
|
|
name=reference.path()
|
|
).execute()
|
|
|
|
|
|
def UpdateBiReservation(client, reference, reservation_size: str):
|
|
"""Updates a BI reservation with the given reservation reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reservation to update.
|
|
reservation_size: size of reservation in GBs. It may only contain digits,
|
|
optionally followed by 'G', 'g', 'GB, 'gb', 'gB', or 'Gb'.
|
|
|
|
Returns:
|
|
Reservation object that was updated.
|
|
Raises:
|
|
ValueError: if reservation_size is malformed.
|
|
"""
|
|
|
|
if (
|
|
reservation_size.upper().endswith('GB')
|
|
and reservation_size[:-2].isdigit()
|
|
):
|
|
reservation_digits = reservation_size[:-2]
|
|
elif (
|
|
reservation_size.upper().endswith('G') and reservation_size[:-1].isdigit()
|
|
):
|
|
reservation_digits = reservation_size[:-1]
|
|
elif reservation_size.isdigit():
|
|
reservation_digits = reservation_size
|
|
else:
|
|
raise ValueError("""Invalid reservation size. The unit for BI reservations
|
|
is GB. The specified reservation size may only contain digits, optionally
|
|
followed by G, g, GB, gb, gB, or Gb.""")
|
|
|
|
reservation_size = int(reservation_digits) * 1024 * 1024 * 1024
|
|
|
|
bi_reservation = {}
|
|
update_mask = ''
|
|
bi_reservation['size'] = reservation_size
|
|
update_mask += 'size,'
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.updateBiReservation(
|
|
name=reference.path(), updateMask=update_mask, body=bi_reservation
|
|
)
|
|
.execute()
|
|
)
|
|
|
|
|
|
|
|
|
|
def GetParamsForUpdateReservation(
|
|
client: ...,
|
|
reference: ...,
|
|
api_version: str, # pylint: disable=unused-argument Cleanup is error prone.
|
|
slots: int,
|
|
ignore_idle_slots: bool,
|
|
target_job_concurrency: Optional[int],
|
|
autoscale_max_slots: Optional[int] = None,
|
|
max_slots: Optional[int] = None,
|
|
scaling_mode: Optional[str] = None,
|
|
labels_to_set: Optional[Dict[str, str]] = None,
|
|
label_keys_to_remove: Optional[Set[str]] = None,
|
|
reservation_group_name: Optional[str] = None,
|
|
scheduling_policy_concurrency: Optional[int] = None,
|
|
scheduling_policy_max_slots: Optional[int] = None,
|
|
) -> Tuple[Dict[str, Any], str]:
|
|
"""Return the request body and update mask for UpdateReservation.
|
|
|
|
Arguments:
|
|
api_version: The api version to make the request against.
|
|
slots: Number of slots allocated to this reservation subtree.
|
|
ignore_idle_slots: Specifies whether queries should ignore idle slots from
|
|
other reservations.
|
|
target_job_concurrency: Job concurrency target.
|
|
autoscale_max_slots: Number of slots to be scaled when needed.
|
|
max_slots: The overall max slots for the reservation.
|
|
scaling_mode: The scaling mode for the reservation.
|
|
reservation_group_name: The reservation group name that the reservation
|
|
belongs to.
|
|
scheduling_policy_concurrency: Soft cap on project concurrency within the
|
|
reservation.
|
|
scheduling_policy_max_slots: Soft cap on project max slots within the
|
|
reservation.
|
|
|
|
Returns:
|
|
Reservation object that was updated.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if parameters are incompatible.
|
|
"""
|
|
reservation = {}
|
|
update_mask = ''
|
|
if slots is not None:
|
|
reservation['slot_capacity'] = slots
|
|
update_mask += 'slot_capacity,'
|
|
|
|
if ignore_idle_slots is not None:
|
|
reservation['ignore_idle_slots'] = ignore_idle_slots
|
|
update_mask += 'ignore_idle_slots,'
|
|
|
|
if target_job_concurrency is not None:
|
|
reservation['concurrency'] = target_job_concurrency
|
|
update_mask += 'concurrency,'
|
|
|
|
if autoscale_max_slots is not None:
|
|
if autoscale_max_slots != 0:
|
|
reservation['autoscale'] = {}
|
|
reservation['autoscale']['max_slots'] = autoscale_max_slots
|
|
update_mask += 'autoscale.max_slots,'
|
|
else:
|
|
# Disable autoscale.
|
|
update_mask += 'autoscale,'
|
|
|
|
if label_keys_to_remove is not None or labels_to_set is not None:
|
|
lookup_reservation = GetReservation(client, reference)
|
|
if 'labels' in lookup_reservation:
|
|
reservation['labels'] = lookup_reservation['labels']
|
|
else:
|
|
reservation['labels'] = {}
|
|
update_mask += 'labels,'
|
|
|
|
if label_keys_to_remove is not None:
|
|
for key in label_keys_to_remove:
|
|
if key in reservation['labels']:
|
|
del reservation['labels'][key]
|
|
|
|
if labels_to_set is not None:
|
|
for key, value in labels_to_set.items():
|
|
reservation['labels'][key] = value
|
|
|
|
if frontend_utils.ValidateAtMostOneSelectedAllowsDefault(
|
|
max_slots, autoscale_max_slots
|
|
):
|
|
raise bq_error.BigqueryError(
|
|
'max_slots and autoscale_max_slots cannot be set at the same time.'
|
|
)
|
|
|
|
# Not we don't need to perform the following check for creation,
|
|
# because there we check the co-existence of max_slots and scaling_mode, so
|
|
# we just need to make sure max_slots and autoscale_max_slots doesn't occur
|
|
# at the same time.
|
|
if (scaling_mode is not None and autoscale_max_slots is not None) and (
|
|
scaling_mode != 'SCALING_MODE_UNSPECIFIED' and autoscale_max_slots != 0
|
|
):
|
|
raise bq_error.BigqueryError(
|
|
'scaling_mode and autoscale_max_slots cannot be set at the same time.'
|
|
)
|
|
|
|
# Note: for update we don't require max_slots and scaling_mode being changed
|
|
# set at the same time.
|
|
# Although backend should make sure them to work together.
|
|
|
|
if max_slots is not None:
|
|
reservation['max_slots'] = max_slots
|
|
update_mask += 'max_slots,'
|
|
if scaling_mode is not None:
|
|
reservation['scaling_mode'] = scaling_mode
|
|
update_mask += 'scaling_mode,'
|
|
|
|
if reservation_group_name is not None:
|
|
reservation['reservation_group'] = reservation_group_name
|
|
update_mask += 'reservation_group,'
|
|
|
|
if scheduling_policy_concurrency is not None:
|
|
reservation['scheduling_policy'] = {}
|
|
reservation['scheduling_policy'][
|
|
'concurrency'
|
|
] = scheduling_policy_concurrency
|
|
update_mask += 'scheduling_policy.concurrency,'
|
|
|
|
if scheduling_policy_max_slots is not None:
|
|
if 'scheduling_policy' not in reservation:
|
|
reservation['scheduling_policy'] = {}
|
|
reservation['scheduling_policy']['max_slots'] = scheduling_policy_max_slots
|
|
update_mask += 'scheduling_policy.max_slots,'
|
|
|
|
return reservation, update_mask
|
|
|
|
|
|
def UpdateReservation(
|
|
client,
|
|
api_version: str,
|
|
reference,
|
|
slots,
|
|
ignore_idle_slots,
|
|
target_job_concurrency: Optional[int],
|
|
autoscale_max_slots: Optional[int] = None,
|
|
max_slots: Optional[int] = None,
|
|
scaling_mode: Optional[str] = None,
|
|
labels_to_set: Optional[Dict[str, str]] = None,
|
|
label_keys_to_remove: Optional[Set[str]] = None,
|
|
reservation_group_name: Optional[str] = None,
|
|
scheduling_policy_concurrency: Optional[int] = None,
|
|
scheduling_policy_max_slots: Optional[int] = None,
|
|
):
|
|
"""Updates a reservation with the given reservation reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
api_version: The api version to make the request against.
|
|
reference: Reservation to update.
|
|
slots: Number of slots allocated to this reservation subtree.
|
|
ignore_idle_slots: Specifies whether queries should ignore idle slots from
|
|
other reservations.
|
|
target_job_concurrency: Job concurrency target.
|
|
autoscale_max_slots: Number of slots to be scaled when needed.
|
|
max_slots: The overall max slots for the reservation.
|
|
scaling_mode: The scaling mode for the reservation.
|
|
reservation_group_name: The reservation group name that the reservation
|
|
belongs to.
|
|
scheduling_policy_concurrency: Soft cap on project concurrency within the
|
|
reservation.
|
|
scheduling_policy_max_slots: Soft cap on project max slots within the
|
|
reservation.
|
|
|
|
Returns:
|
|
Reservation object that was updated.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if autoscale_max_slots is used with other
|
|
version.
|
|
"""
|
|
reservation, update_mask = GetParamsForUpdateReservation(
|
|
client,
|
|
reference,
|
|
api_version,
|
|
slots,
|
|
ignore_idle_slots,
|
|
target_job_concurrency,
|
|
autoscale_max_slots,
|
|
max_slots,
|
|
scaling_mode,
|
|
labels_to_set,
|
|
label_keys_to_remove,
|
|
reservation_group_name,
|
|
scheduling_policy_concurrency,
|
|
scheduling_policy_max_slots,
|
|
)
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.patch(name=reference.path(), updateMask=update_mask, body=reservation)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def CreateCapacityCommitment(
|
|
client,
|
|
reference,
|
|
edition,
|
|
slots: int,
|
|
plan: str,
|
|
renewal_plan: str,
|
|
multi_region_auxiliary: Optional[bool],
|
|
) -> Dict[str, Any]:
|
|
"""Create a capacity commitment.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Project to create a capacity commitment within.
|
|
edition: The edition for this capacity commitment.
|
|
slots: Number of slots in this commitment.
|
|
plan: Commitment plan for this capacity commitment.
|
|
renewal_plan: Renewal plan for this capacity commitment.
|
|
multi_region_auxiliary: Whether this commitment is for the auxiliary region.
|
|
|
|
Returns:
|
|
Capacity commitment object that was created.
|
|
"""
|
|
capacity_commitment = {}
|
|
capacity_commitment['slot_count'] = slots
|
|
capacity_commitment['plan'] = plan
|
|
capacity_commitment['renewal_plan'] = renewal_plan
|
|
if multi_region_auxiliary is not None:
|
|
capacity_commitment['multi_region_auxiliary'] = multi_region_auxiliary
|
|
if edition is not None:
|
|
capacity_commitment['edition'] = edition
|
|
parent = 'projects/%s/locations/%s' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
capacity_commitment_id = None
|
|
if reference.capacityCommitmentId and reference.capacityCommitmentId != ' ':
|
|
capacity_commitment_id = reference.capacityCommitmentId
|
|
request = (
|
|
client.projects()
|
|
.locations()
|
|
.capacityCommitments()
|
|
.create(
|
|
parent=parent,
|
|
body=capacity_commitment,
|
|
capacityCommitmentId=capacity_commitment_id,
|
|
)
|
|
)
|
|
return request.execute()
|
|
|
|
|
|
def ListCapacityCommitments(client, reference, page_size, page_token):
|
|
"""Lists capacity commitments for given project and location.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reference to the project and location.
|
|
page_size: Number of results to show.
|
|
page_token: Token to retrieve the next page of results.
|
|
|
|
Returns:
|
|
list of CapacityCommitments objects.
|
|
"""
|
|
parent = 'projects/%s/locations/%s' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.capacityCommitments()
|
|
.list(parent=parent, pageSize=page_size, pageToken=page_token)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def GetCapacityCommitment(client, reference):
|
|
"""Gets a capacity commitment with the given capacity commitment reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Capacity commitment to get.
|
|
|
|
Returns:
|
|
Capacity commitment object corresponding to the given id.
|
|
"""
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.capacityCommitments()
|
|
.get(name=reference.path())
|
|
.execute()
|
|
)
|
|
|
|
|
|
def DeleteCapacityCommitment(client, reference, force=None):
|
|
"""Deletes a capacity commitment with the given capacity commitment reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Capacity commitment to delete.
|
|
force: Force delete capacity commitment, ignoring commitment end time.
|
|
"""
|
|
client.projects().locations().capacityCommitments().delete(
|
|
name=reference.path(), force=force
|
|
).execute()
|
|
|
|
|
|
def UpdateCapacityCommitment(client, reference, plan, renewal_plan):
|
|
"""Updates a capacity commitment with the given reference.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Capacity commitment to update.
|
|
plan: Commitment plan for this capacity commitment.
|
|
renewal_plan: Renewal plan for this capacity commitment.
|
|
|
|
Returns:
|
|
Capacity commitment object that was updated.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if capacity commitment cannot be updated.
|
|
"""
|
|
if plan is None and renewal_plan is None:
|
|
raise bq_error.BigqueryError('Please specify fields to be updated.')
|
|
capacity_commitment = {}
|
|
update_mask = []
|
|
if plan is not None:
|
|
capacity_commitment['plan'] = plan
|
|
update_mask.append('plan')
|
|
if renewal_plan is not None:
|
|
capacity_commitment['renewal_plan'] = renewal_plan
|
|
update_mask.append('renewal_plan')
|
|
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.capacityCommitments()
|
|
.patch(
|
|
name=reference.path(),
|
|
updateMask=','.join(update_mask),
|
|
body=capacity_commitment,
|
|
)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def SplitCapacityCommitment(
|
|
client,
|
|
reference,
|
|
slots,
|
|
):
|
|
"""Splits a capacity commitment with the given reference into two.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Capacity commitment to split.
|
|
slots: Number of slots in the first capacity commitment after the split.
|
|
|
|
Returns:
|
|
List of capacity commitment objects after the split.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if capacity commitment cannot be updated.
|
|
"""
|
|
if slots is None:
|
|
raise bq_error.BigqueryError('Please specify slots for the split.')
|
|
body = {'slotCount': slots}
|
|
response = (
|
|
client.projects()
|
|
.locations()
|
|
.capacityCommitments()
|
|
.split(name=reference.path(), body=body)
|
|
.execute()
|
|
)
|
|
if 'first' not in response or 'second' not in response:
|
|
raise bq_error.BigqueryError('internal error')
|
|
return [response['first'], response['second']]
|
|
|
|
|
|
def MergeCapacityCommitments(
|
|
client, project_id, location, capacity_commitment_ids
|
|
):
|
|
"""Merges capacity commitments into one.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
project_id: The project ID of the resources to update.
|
|
location: Capacity commitments location.
|
|
capacity_commitment_ids: List of capacity commitment ids.
|
|
|
|
Returns:
|
|
Merged capacity commitment.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if capacity commitment cannot be merged.
|
|
"""
|
|
if not project_id:
|
|
raise bq_error.BigqueryError('project id must be specified.')
|
|
if not location:
|
|
raise bq_error.BigqueryError('location must be specified.')
|
|
if capacity_commitment_ids is None or len(capacity_commitment_ids) < 2:
|
|
raise bq_error.BigqueryError(
|
|
'at least 2 capacity commitments must be specified.'
|
|
)
|
|
parent = 'projects/%s/locations/%s' % (project_id, location)
|
|
body = {'capacityCommitmentIds': capacity_commitment_ids}
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.capacityCommitments()
|
|
.merge(parent=parent, body=body)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def CreateReservationAssignment(
|
|
client,
|
|
reference,
|
|
assignee_type: Optional[str] = None,
|
|
assignee_id: Optional[str] = None,
|
|
job_type: Optional[str] = None,
|
|
priority: Optional[str] = None,
|
|
scheduling_policy_max_slots: Optional[int] = None,
|
|
scheduling_policy_concurrency: Optional[int] = None,
|
|
):
|
|
"""Creates a reservation assignment for a given project/folder/organization.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reference to the project reservation is assigned. Location must
|
|
be the same location as the reservation.
|
|
assignee_type: Type of assignees for the reservation assignment.
|
|
assignee_id: Project/folder/organization ID, to which the reservation is
|
|
assigned.
|
|
job_type: Type of jobs for this assignment.
|
|
priority: Default job priority for this assignment.
|
|
scheduling_policy_max_slots: Soft cap on max slots for a specific project
|
|
within the reservation.
|
|
scheduling_policy_concurrency: Soft cap on job concurrency for a specific
|
|
project within the reservation.
|
|
|
|
Returns:
|
|
ReservationAssignment object that was created.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if assignment cannot be created.
|
|
"""
|
|
reservation_assignment = GetBodyForCreateReservationAssignment(
|
|
reference,
|
|
job_type,
|
|
priority,
|
|
assignee_type,
|
|
assignee_id,
|
|
scheduling_policy_max_slots,
|
|
scheduling_policy_concurrency,
|
|
)
|
|
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.assignments()
|
|
.create(parent=reference.path(), body=reservation_assignment)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def DeleteReservationAssignment(client, reference):
|
|
"""Deletes given reservation assignment.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reference to the reservation assignment.
|
|
"""
|
|
client.projects().locations().reservations().assignments().delete(
|
|
name=reference.path()
|
|
).execute()
|
|
|
|
|
|
def MoveReservationAssignment(
|
|
client,
|
|
id_fallbacks: NamedTuple(
|
|
'IDS',
|
|
[
|
|
('project_id', Optional[str]),
|
|
('api_version', Optional[str]),
|
|
],
|
|
),
|
|
reference,
|
|
destination_reservation_id,
|
|
default_location,
|
|
):
|
|
"""Moves given reservation assignment under another reservation."""
|
|
destination_reservation_reference = bq_client_utils.GetReservationReference(
|
|
id_fallbacks=id_fallbacks,
|
|
identifier=destination_reservation_id,
|
|
default_location=default_location,
|
|
check_reservation_project=False,
|
|
)
|
|
body = {'destinationId': destination_reservation_reference.path()}
|
|
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.assignments()
|
|
.move(name=reference.path(), body=body)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def UpdateReservationAssignment(
|
|
client,
|
|
reference,
|
|
priority,
|
|
scheduling_policy_max_slots,
|
|
scheduling_policy_concurrency,
|
|
):
|
|
"""Updates reservation assignment.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reference to the reservation assignment.
|
|
priority: Default job priority for this assignment.
|
|
scheduling_policy_max_slots: Max slots for the scheduling policy.
|
|
scheduling_policy_concurrency: Concurrency for the scheduling policy.
|
|
|
|
Returns:
|
|
Reservation assignment object that was updated.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if assignment cannot be updated.
|
|
"""
|
|
reservation_assignment = {}
|
|
update_mask = ''
|
|
if priority is not None:
|
|
if not priority:
|
|
priority = 'JOB_PRIORITY_UNSPECIFIED'
|
|
reservation_assignment['priority'] = priority
|
|
update_mask += 'priority,'
|
|
if scheduling_policy_max_slots is not None:
|
|
reservation_assignment['scheduling_policy'] = {}
|
|
reservation_assignment['scheduling_policy'][
|
|
'max_slots'
|
|
] = scheduling_policy_max_slots
|
|
update_mask += 'scheduling_policy.max_slots,'
|
|
if scheduling_policy_concurrency is not None:
|
|
if 'scheduling_policy' not in reservation_assignment:
|
|
reservation_assignment['scheduling_policy'] = {}
|
|
reservation_assignment['scheduling_policy'][
|
|
'concurrency'
|
|
] = scheduling_policy_concurrency
|
|
update_mask += 'scheduling_policy.concurrency,'
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.assignments()
|
|
.patch(
|
|
name=reference.path(),
|
|
updateMask=update_mask,
|
|
body=reservation_assignment,
|
|
)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def ListReservationAssignments(client, reference, page_size, page_token):
|
|
"""Lists reservation assignments for given project and location.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
reference: Reservation reference for the parent.
|
|
page_size: Number of results to show.
|
|
page_token: Token to retrieve the next page of results.
|
|
|
|
Returns:
|
|
ReservationAssignment object that was created.
|
|
"""
|
|
return (
|
|
client.projects()
|
|
.locations()
|
|
.reservations()
|
|
.assignments()
|
|
.list(parent=reference.path(), pageSize=page_size, pageToken=page_token)
|
|
.execute()
|
|
)
|
|
|
|
|
|
|
|
|
|
def SearchAllReservationAssignments(
|
|
client, location: str, job_type: str, assignee_type: str, assignee_id: str
|
|
) -> Dict[str, Any]:
|
|
"""Searches reservations assignments for given assignee.
|
|
|
|
Arguments:
|
|
client: The client used to make the request.
|
|
location: location of interest.
|
|
job_type: type of job to be queried.
|
|
assignee_type: Type of assignees for the reservation assignment.
|
|
assignee_id: Project/folder/organization ID, to which the reservation is
|
|
assigned.
|
|
|
|
Returns:
|
|
ReservationAssignment object if it exists.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: If required parameters are not passed in or
|
|
reservation assignment not found.
|
|
"""
|
|
if not location:
|
|
raise bq_error.BigqueryError('location not specified.')
|
|
if not job_type:
|
|
raise bq_error.BigqueryError('job_type not specified.')
|
|
if not assignee_type:
|
|
raise bq_error.BigqueryError('assignee_type not specified.')
|
|
if not assignee_id:
|
|
raise bq_error.BigqueryError('assignee_id not specified.')
|
|
# assignee_type is singular, that's why we need additional 's' inside
|
|
# format string for assignee below.
|
|
assignee = '%ss/%s' % (assignee_type.lower(), assignee_id)
|
|
query = 'assignee=%s' % assignee
|
|
parent = 'projects/-/locations/%s' % location
|
|
|
|
response = (
|
|
client.projects()
|
|
.locations()
|
|
.searchAllAssignments(parent=parent, query=query)
|
|
.execute()
|
|
)
|
|
if 'assignments' in response:
|
|
for assignment in response['assignments']:
|
|
if assignment['jobType'] == job_type:
|
|
return assignment
|
|
raise bq_error.BigqueryError('Reservation assignment not found')
|
|
|
|
|
|
def CreateReservationGroup(
|
|
reservation_group_client: discovery.Resource,
|
|
reference: bq_id_utils.ApiClientHelper.ReservationGroupReference,
|
|
) -> Dict[str, Any]:
|
|
"""Creates a reservation group with the given reservation group reference.
|
|
|
|
Arguments:
|
|
reservation_group_client: The client used to make the request.
|
|
reference: Reservation group to create.
|
|
|
|
Returns:
|
|
Reservation group object that was created.
|
|
|
|
Raises:
|
|
bq_error.BigqueryError: if reservation group cannot be created.
|
|
"""
|
|
|
|
parent = 'projects/%s/locations/%s' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
reservation_group = {}
|
|
return (
|
|
reservation_group_client.projects()
|
|
.locations()
|
|
.reservationGroups()
|
|
.create(
|
|
parent=parent,
|
|
body=reservation_group,
|
|
reservationGroupId=reference.reservationGroupId,
|
|
)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def ListReservationGroups(
|
|
reservation_group_client: discovery.Resource,
|
|
reference: bq_id_utils.ApiClientHelper.ReservationGroupReference,
|
|
page_size: int,
|
|
page_token: str,
|
|
) -> ...:
|
|
"""Lists reservation groups in the project and location for the given reference.
|
|
|
|
Arguments:
|
|
reservation_group_client: The client used to make the request.
|
|
reference: Reservation group reference containing project and location.
|
|
page_size: Number of results to show.
|
|
page_token: Token to retrieve the next page of results.
|
|
|
|
Returns:
|
|
Reservation group object that was created.
|
|
"""
|
|
parent = 'projects/%s/locations/%s' % (
|
|
reference.projectId,
|
|
reference.location,
|
|
)
|
|
return (
|
|
reservation_group_client.projects()
|
|
.locations()
|
|
.reservationGroups()
|
|
.list(parent=parent, pageSize=page_size, pageToken=page_token)
|
|
.execute()
|
|
)
|
|
|
|
|
|
def GetReservationGroup(
|
|
reservation_group_client: discovery.Resource,
|
|
reference: bq_id_utils.ApiClientHelper.ReservationGroupReference,
|
|
) -> ...:
|
|
"""Gets a reservation group with the given reservation group reference.
|
|
|
|
Arguments:
|
|
reservation_group_client: The client used to make the request.
|
|
reference: Reservation group to get.
|
|
|
|
Returns:
|
|
Reservation group object corresponding to the given id.
|
|
"""
|
|
return (
|
|
reservation_group_client.projects()
|
|
.locations()
|
|
.reservationGroups()
|
|
.get(name=reference.path())
|
|
.execute()
|
|
)
|
|
|
|
|
|
def DeleteReservationGroup(
|
|
reservation_group_client: discovery.Resource,
|
|
reference: bq_id_utils.ApiClientHelper.ReservationGroupReference,
|
|
) -> ...:
|
|
"""Deletes a reservation group with the given reservation group reference.
|
|
|
|
Arguments:
|
|
reservation_group_client: The client used to make the request.
|
|
reference: Reservation group to delete.
|
|
"""
|
|
reservation_group_client.projects().locations().reservationGroups().delete(
|
|
name=reference.path()
|
|
).execute()
|
|
|
|
|
|
def SetReservationIAMPolicy(
|
|
apiclient: discovery.Resource,
|
|
reference: bq_id_utils.ApiClientHelper.ReservationReference,
|
|
policy: str,
|
|
) -> ...:
|
|
"""Sets IAM policy for the given reservation resource.
|
|
|
|
Arguments:
|
|
apiclient: the apiclient used to make the request.
|
|
reference: the ReservationReference for the reservation resource.
|
|
policy: The policy string in JSON format.
|
|
|
|
Returns:
|
|
The updated IAM policy attached to the given reservation resource.
|
|
|
|
Raises:
|
|
BigqueryTypeError: if reference is not a ReservationReference.
|
|
"""
|
|
bq_id_utils.typecheck(
|
|
reference,
|
|
bq_id_utils.ApiClientHelper.ReservationReference,
|
|
method='SetReservationIAMPolicy',
|
|
)
|
|
request = {'policy': policy}
|
|
return (
|
|
apiclient.projects()
|
|
.locations()
|
|
.reservations()
|
|
.setIamPolicy(body=request, resource=reference.path())
|
|
.execute()
|
|
)
|
|
|
|
|
|
def GetReservationIAMPolicy(
|
|
apiclient: discovery.Resource,
|
|
reference: bq_id_utils.ApiClientHelper.ReservationReference,
|
|
) -> ...:
|
|
"""Gets IAM policy for the given reservation resource.
|
|
|
|
Arguments:
|
|
apiclient: the apiclient used to make the request.
|
|
reference: the ReservationReference for the reservation resource.
|
|
|
|
Returns:
|
|
The IAM policy attached to the given reservation resource.
|
|
|
|
Raises:
|
|
BigqueryTypeError: if reference is not a ReservationReference.
|
|
"""
|
|
bq_id_utils.typecheck(
|
|
reference,
|
|
bq_id_utils.ApiClientHelper.ReservationReference,
|
|
method='GetReservationIAMPolicy',
|
|
)
|
|
return (
|
|
apiclient.projects()
|
|
.locations()
|
|
.reservations()
|
|
.getIamPolicy(resource=reference.path())
|
|
.execute()
|
|
)
|
|
|
|
|