wdZddlmZddlmZddlmZddlmZmZmZm Z m Z m Z ddl m Z ddlmZddlmZddlmZdd lmZ dDd ed ed ede ede ede ede ede ede edeeeffdZ dDd ed ed ede ede ede ede ede ede edeeeffdZ dEdddddede ede edef dZdZdZ dFd Zd!efd"Z dGddddd ed ed ede ede ede ede ed#e eeefd$e e ede ede eeefeffd%Z dGd ede ede ede ede ed#e eeefd$e e ede efd&Z d ed'ed(ede edeeeff d)Z!d*Z"d+Z#dEd,Z$d-Z%d.Z&d/Z'd0Z(d1Z)d2ed3d4e efd e efgfd5Z*d6Z+d7Z,d8ed9ed:ed;edeeeff d<Z-d=e j\dej^j`deeeffd>Z1d=e j\dej^j`dededdf d?Z2d=e j\dej^j`ddfd@Z3d=e j\dej^j`ddfdAZ4dBe j\dej^jjddfdCZ6y )Hz,The BigQuery CLI reservation client library.)absolute_import)division)print_function)AnyDict NamedTupleOptionalSetTuple) discovery)utils)bq_error) bq_id_utilsN api_versionslotsignore_idle_slotstarget_job_concurrencymulti_region_auxiliaryautoscale_max_slots max_slots scaling_modereservation_group_namereturnc  i} || d<|| d<||| d<||| d<| i| d<|| dd<||| d<tj||rtjd||||tjd ||| d<||| d <| | | d <| S) aReturn 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. Returns: Reservation object that was created. Raises: bq_error.BigqueryError: if requirements for parameters are not met. slot_capacityrr concurrency autoscalereditionAmax_slots and autoscale_max_slots cannot be set at the same time.z8max_slots and scaling_mode must be set at the same time.rreservation_group)frontend_utilsValidateAtMostOneSelectedr BigqueryError) rrrrrrrrrr reservations )platform/bq/clients/client_reservation.pyGetBodyForCreateReservationr&sF+!&+o%6+!"',BK()'!7K $!K ,?K [) $K --i9LM  K  4L4  B (K ".K''=K#$ c t|||||||| | | } d|jd|j} |jj j j | | |jjS)aCreate 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. projects/ /locations/)parentbody reservationId) r& projectIdlocationprojects locations reservationscreater-execute)clientr referencerrrrrrrrrr$r+s r%CreateReservationr7[sP,   + &  ooy{|~ vk9P9Pwyr'r5.r6 page_size page_tokenr cvd|jd|j}|P|jjj j |||d|zj S|jjj j |||j S)aList 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. r)r*zreservation_group=%s)r+pageSize pageTokenfilterr+r;r<)r.r/r0r1r2listr4)r5r6r8r9r r+s r%ListReservationsr@s* &"    ),==    ooy{|~ t6ItDwy r'cd|jd|jd}|jjj |j }|S)aList 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. r)r*z/biReservationname)r.r/r0r1getBiReservationr4)r5r6r+responses r%ListBiReservationsrFsU &  oo!!#44&4AIIK  /r'c|jjjj|j j S)zGets 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. rB)r0r1r2getpathr4r5r6s r%GetReservationrKs; ooy{|~ s  s!wy r'c|jjjj|j j y)zDeletes a reservation with the given reservation reference. Arguments: client: The client used to make the request. reference: Reservation to delete. rBN)r0r1r2deleterIr4rJs r%DeleteReservationrNsA //,,.55 >> 6 GIr'reservation_sizec|jjdr|ddjr|dd}nV|jjdr|ddjr|dd}n|jr|}n tdt |dzdzdz}i}d}||d <|d z }|j j j|j|| jS) aUpdates 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. GBNGzInvalid 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.isizezsize,rC updateMaskr,) upperendswithisdigit ValueErrorintr0r1updateBiReservationrIr4)r5r6rOreservation_digitsbi_reservation update_masks r%UpdateBiReservationrbs ''- 3B  ' ' ))#2.'',1A#21F1N1N1P)#2.!) , --+,t3d:TA.++. + ooy{~~Knwy r' labels_to_setlabel_keys_to_removec <i} d} | || d<| dz } | || d<| dz } | || d<| dz } ||dk7ri| d <|| d d <| d z } n| d z } | | #t||}d |vr |d | d <ni| d <| dz } | | D]}|| d vs | d |=| | jD] \}}|| d |<tj||rt j d|!||dk7r|dk7rt j d| || d <| dz } | || d<| dz } | | | d<| dz } | | fS)aReturn 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. Returns: Reservation object that was updated. Raises: bq_error.BigqueryError: if parameters are incompatible. rUrzslot_capacity,rzignore_idle_slots,rz concurrency,rrrzautoscale.max_slots,z autoscale,labelszlabels,rSCALING_MODE_UNSPECIFIEDzDscaling_mode and autoscale_max_slots cannot be set at the same time.z max_slots,rz scaling_mode,r zreservation_group,)rKitemsr!&ValidateAtMostOneSelectedAllowsDefaultrr#)r5r6rrrrrrrrcrdrr$ralookup_reservationkeyvalues r%GetParamsForUpdateReservationrm3sD++ #(K ##K"'8K#$''K'!7K >!K$a!#k+.Ak+{+++k\!k%)B' :%%0:k( k(9K%# H% %  !# &$#))+ U#(k(C ,::$  K #6#B005HA5M  N (K <K".K?"K''=K#$''K k !!r'c t|||||||||| | | \} } |jjjj |j | | j S)ahUpdates 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. Returns: Reservation object that was updated. Raises: bq_error.BigqueryError: if autoscale_max_slots is used with other version. rW)rmr0r1r2patchrIr4)r5rr6rrrrrrrcrdrr$ras r%UpdateReservationrpsuJ;   +{ ooy{|~ u).."{uMwy r'plan renewal_plancfi}||d<||d<||d<|||d<|||d<d|jd|j}d} |jr|jd k7r |j} |jj j j ||| } | jS) aCreate 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. slot_countrqrrNrrr)r* )r+r,capacityCommitmentId)r.r/rvr0r1capacityCommitmentsr3r4) r5r6rrrqrrrcapacity_commitmentr+capacity_commitment_idrequests r%CreateCapacityCommitmentr{s.&+l# $f(4n%'4J01 %, " & ## (F(F#(M&;; ooy{ v"5    r'cd|jd|j}|jjj j |||j S)aMLists 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. r)r*r>)r.r/r0r1rwr?r4)r5r6r8r9r+s r%ListCapacityCommitmentsr}sU &  ooy{ t6ItDwy r'c|jjjj|j j S)zGets 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. rB)r0r1rwrHrIr4rJs r%GetCapacityCommitmentr)s= ooy{ s  s!wy r'c|jjjj|j |j y)aDeletes 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. )rCforceN)r0r1rwrMrIr4)r5r6rs r%DeleteCapacityCommitmentr<sC //335<< >> 5= GIr'cn||tjdi}g}|||d<|jd|||d<|jd|jj j j |jdj||jS)aUpdates 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. z$Please specify fields to be updated.rqrr,rW) rr#appendr0r1rwrorIjoinr4)r5r6rqrrrxras r%UpdateCapacityCommitmentrIs \l*  !G HH+ "&v*6'~& ooy{ u~~XXk*"  wy r'c8|tjdd|i}|jjj j |j |j}d|vsd|vrtjd|d|dgS)aSplits 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. z#Please specify slots for the split. slotCountrCr,firstsecondzinternal error)rr#r0r1rwsplitrIr4)r5r6rr,rEs r%SplitCapacityCommitmentrps$ ]  !F GG u $ ooy{ u).."u.wy   H 8  !1 22 7 Xh/ 00r'cV|stjd|stjd|t|dkrtjdd|d|}d|i}|jj j j ||jS) aMerges 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. zproject id must be specified.zlocation must be specified.z2at least 2 capacity commitments must be specified.r)r*capacityCommitmentIdsr+r,)rr#lenr0r1rwmerger4)r5 project_idr/capacity_commitment_idsr+r,s r%MergeCapacityCommitmentsrs"   !@ AA   !> ??$,C(Dq(H  < )I*4X >& !#: ;$ ooy{ uFu&wy r'ci}|stjd||d<|r||d<|stjd|stjd|jd||d<|jj j j j|j|jS) aCreates 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. 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. Returns: ReservationAssignment object that was created. Raises: bq_error.BigqueryError: if assignment cannot be created. job_type not specified.job_typepriorityassignee_type not specified.assignee_id not specified.s/assigneer) rr#lowerr0r1r2 assignmentsr3rIr4)r5r6rr assignee_type assignee_idreservation_assignments r%CreateReservationAssignmentrs*   !: ;;'/$ )1:&   !? @@   != >>($  ooy{|~{} vY^^%,BvCwy r'c|jjjjj |j j y)zDeletes given reservation assignment. Arguments: client: The client used to make the request. reference: Reference to the reservation assignment. rBN)r0r1r2rrMrIr4rJs r%DeleteReservationAssignmentrsK //,,.::<CC >> D GIr' id_fallbacksIDSrc&tj|||d}d|ji}|jj j j j|j|jS)z=Moves given reservation assignment under another reservation.F)r identifierdefault_locationcheck_reservation_project destinationIdr) bq_client_utilsGetReservationReferencerIr0r1r2rmover4)r5rr6destination_reservation_idr!destination_reservation_referencer,s r%MoveReservationAssignmentrsy'6&M&M+' % '# <AAC D$ ooy{|~{} t!t-wy r'ci}d}||sd}||d<|dz }|jjjjj |j ||j S)aYUpdates reservation assignment. Arguments: client: The client used to make the request. reference: Reference to the reservation assignment. priority: Default job priority for this assignment. Returns: Reservation assignment object that was updated. Raises: bq_error.BigqueryError: if assignment cannot be updated. rUJOB_PRIORITY_UNSPECIFIEDrz priority,rW)r0r1r2rrorIr4)r5r6rrras r%UpdateReservationAssignmentrs+  +h)1:&;K ooy{|~{} u~~ %  wy r'c|jjjjj |j ||j S)aYLists 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. r>)r0r1r2rr?rIr4)r5r6r8r9s r%ListReservationAssignmentsr4sF ooy{|~{} t9>>#i:tNwy r'r/rrrc|stjd|stjd|stjd|stjd|jd|}d|z}d|z}|jj j ||j }d |vr|d D]} | d |k(s | cStjd ) aSearches 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. zlocation not specified.rrrrz assignee=%szprojects/-/locations/%s)r+queryrjobTypez Reservation assignment not found)rr#rr0r1searchAllAssignmentsr4) r5r/rrrrrr+rE assignments r%SearchAllReservationAssignmentsrLs(   !: ;;   !: ;;   !? @@   != >>',,. <( ( "% $x /& ooy{67wy  h}- I ( *. ABBr'reservation_group_clientcd|jd|j}i}|jjj j |||j jS)aQCreates 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. r)r*)r+r,reservationGroupId)r.r/r0r1reservationGroupsr3rr4)rr6r+r s r%CreateReservationGroupr{sm$ &'')y{ v &99 wy r'cd|jd|j}|jjj j |||j S)aLists 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. r)r*r>)r.r/r0r1rr?r4)rr6r8r9r+s r%ListReservationGroupsrsW$ & '')y{ t6ItDwy r'c|jjjj|j j S)a 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. rB)r0r1rrHrIr4rr6s r%GetReservationGrouprs?'')y{ s  s!wy r'c|jjjj|j j y)zDeletes 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. rBN)r0r1rrMrIr4rs r%DeleteReservationGrouprsD##%//1CCELL >> M GIr' apiclientctj|tjjd|j j j j|jjS)aZGets 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. GetReservationIAMPolicy)method)resource) r typecheckApiClientHelperReservationReferencer0r1r2 getIamPolicyrIr4)rr6s r%rrsb !!66 & y{|~|Y^^-|.wy r')NNNN)N)r6.)NNNNNN)7__doc__ __future__rrrtypingrrrr r r googleapiclientr clientsr rfrontendr!rrstrr]boolr&r7r@rFrKrNrbrmrpr{r}rrrrrrrrrrrResourcerReservationGroupReferencerrrrrrr'r%rsE2&%>>%,,*.#"&,0FF FF %SM F %TN F"#F}F3-F%SMF #s(^Fd*.#"&,0@@  @  @%SM@%TN@"#@}@3-@%SM@ #s(^@P(, + +++ +  } +  +\(&  ,S,p*.#"&.2/3,0k" k"k"k"  k"  k" %SM k""#k"}k"3-k"DcN+k"#3s8,k"%SMk" 4S>3 k"j*.#"&.2/3,099 %SM 9"#9}93-9DcN+9#3s8,9%SM9x0  0  0  0%TN0 #s(^0f2& $N1B!H,^  8C= ) HSM * @!H0,C,C%(,C9<,CKN,C #s(^,C^ '00 **DD  #s(^ F'00**DD   <'00**DD , '00 **DD   "!!**?? r'