def _GetUpdatedStatefulPolicy(self, client, current_stateful_policy, args):
"""Create an updated stateful policy based on specified args."""
update_disks = args.stateful_disk
remove_device_names = args.remove_stateful_disks
stateful_disks = self._GetUpdatedStatefulPolicyForDisks(
client, current_stateful_policy, update_disks, remove_device_names)
return policy_utils.MakeStatefulPolicy(client.messages, stateful_disks)
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)