def _CreateStatefulPolicy(self, args, client): """Create stateful policy from disks of args --stateful-disk.""" stateful_disks = [] for stateful_disk_dict in (args.stateful_disk or []): stateful_disks.append( policy_utils.MakeStatefulPolicyPreservedStateDiskEntry( client.messages, stateful_disk_dict)) stateful_disks.sort(key=lambda x: x.key) return policy_utils.MakeStatefulPolicy(client.messages, stateful_disks)
def _GetUpdatedStatefulPolicy(self, client, current_stateful_policy, update_disks=None, remove_device_names=None): """Create an updated stateful policy with the updated disk data and removed disks as specified.""" # Extract disk protos from current stateful policy proto if current_stateful_policy and current_stateful_policy.preservedState \ and current_stateful_policy.preservedState.disks: current_disks = current_stateful_policy \ .preservedState.disks.additionalProperties else: current_disks = [] # Map of disks to have in the stateful policy, after updating and removing # the disks specified by the update and remove flags. final_disks_map = { disk_entry.key: disk_entry for disk_entry in current_disks } # Update the disks specified in --update-stateful-disk for update_disk in (update_disks or []): device_name = update_disk.get('device-name') updated_preserved_state_disk = ( policy_utils.MakeStatefulPolicyPreservedStateDiskEntry( client.messages, update_disk)) # Patch semantics on the `--update-stateful-disk` flag if device_name in final_disks_map: policy_utils.PatchStatefulPolicyDisk( final_disks_map[device_name], updated_preserved_state_disk) else: final_disks_map[device_name] = updated_preserved_state_disk # Remove the disks specified in --remove-stateful-disks for device_name in remove_device_names or []: del final_disks_map[device_name] stateful_disks = sorted([ stateful_disk for _, stateful_disk in six.iteritems(final_disks_map) ], key=lambda x: x.key) return policy_utils.MakeStatefulPolicy(client.messages, stateful_disks)