@@ -605,42 +605,170 @@ def wait_for_ttr(
605605 return False
606606
607607
608+ def get_node_pool_disk_size (node_pool : Info ) -> int :
609+ """Gets the disk size of a GKE node pool using gcloud command.
610+
611+ Args:
612+ node_pool: An instance of the Info class that encapsulates the
613+ configuration and metadata of a GKE node pool.
614+
615+ Returns:
616+ The disk size of the node pool in GB.
617+ """
618+ command = (
619+ f"gcloud container node-pools describe { node_pool .node_pool_name }
620+ f"--project={ node_pool .project_id }
621+ f"--cluster={ node_pool .cluster_name }
622+ f"--location={ node_pool .location }
623+ f'--format="value(config.diskSizeGb)"'
624+ )
625+
626+ result = subprocess .run_exec (command ).strip ()
627+
628+ return int (result )
629+
630+
631+ def get_node_pool_labels (node_pool : Info ) -> dict [str , str ]:
632+ """Gets the labels of a GKE node pool using gcloud command.
633+
634+ Args:
635+ node_pool: An instance of the Info class that encapsulates the
636+ configuration and metadata of a GKE node pool.
637+
638+ Returns:
639+ A dictionary contains the node pool labels.
640+ """
641+ command = (
642+ f"gcloud container node-pools describe { node_pool .node_pool_name }
643+ f"--project={ node_pool .project_id }
644+ f"--cluster={ node_pool .cluster_name }
645+ f"--location={ node_pool .location }
646+ f"--format='json(config.resourceLabels)'"
647+ )
648+
649+ result = (
650+ json .loads (subprocess .run_exec (command ).strip ())
651+ .get ("config" , {})
652+ .get ("resourceLabels" , {})
653+ )
654+
655+ return result
656+
657+
658+ class UpdateTarget (enum .Enum ):
659+ """Defines what to update on the node pool."""
660+
661+ DISK_SIZE = "disk-size"
662+ LABEL = "labels"
663+
664+
665+ @dataclasses .dataclass
666+ class NodePoolUpdateSpec :
667+ """Configuration parameters defining a mutation on a GKE node pool.
668+
669+ Attributes:
670+ target: The specific node pool attribute to update.
671+ delta: The change to apply to the target's current state.
672+ """
673+
674+ target : UpdateTarget
675+ delta : int | dict [str , str ]
676+
677+ @staticmethod
678+ def DiskSize (delta : int ) -> "NodePoolUpdateSpec" :
679+ if not isinstance (delta , int ):
680+ raise TypeError (f"Disk size delta must be an integer. Got: { type (delta )} )
681+
682+ if delta <= 0 :
683+ raise ValueError (f"Disk size delta must be positive. Got: { delta } )
684+
685+ return NodePoolUpdateSpec (
686+ target = UpdateTarget .DISK_SIZE ,
687+ delta = delta ,
688+ )
689+
690+ @staticmethod
691+ def Label (delta : dict [str , str ]) -> "NodePoolUpdateSpec" :
692+ if not isinstance (delta , dict ):
693+ raise TypeError (f"Label delta must be a dictionary. Got: { type (delta )} )
694+
695+ key_pattern = re .compile (r"^[a-z][a-z0-9_-]*$" )
696+ for k , v in delta .items ():
697+ if not isinstance (k , str ) or not isinstance (v , str ):
698+ raise TypeError (
699+ f"All label keys and values must be strings. "
700+ f"Found incompatible item: key='{ k } { type (k )}
701+ f"value='{ v } { type (v )}
702+ )
703+
704+ if not key_pattern .match (k ):
705+ raise ValueError (
706+ f"Invalid label key: '{ k }
707+ "Keys must start with a lowercase letter and contain only "
708+ "lowercase letters ([a-z]), numeric characters ([0-9]), "
709+ "underscores (_) and dashes (-)."
710+ )
711+
712+ return NodePoolUpdateSpec (
713+ target = UpdateTarget .LABEL ,
714+ delta = delta ,
715+ )
716+
717+
608718@task
609- def update_labels (node_pool : Info , node_labels : dict ) -> TimeUtil :
610- """Updates the labels of a GKE node pool using gcloud command .
719+ def update (node_pool : Info , spec : NodePoolUpdateSpec ) -> TimeUtil :
720+ """Applies an update to a GKE node pool based on the provided specification .
611721
612- This task updates GKE node pool labels via gcloud.
613- It captures the current time before execution and returns it as
614- a TimeUtil object for downstream tracking .
722+ This task performs a state-aware update. It retrieves the current node pool
723+ state, resolves the final desired configuration based on the provided `spec`,
724+ and executes the update operation .
615725
616726 Args:
617- node_pool: An instance of the Info class.
618- node_labels: A dictionary of labels to update or remove.
727+ node_pool: An instance of the Info class that encapsulates the
728+ configuration and metadata of a GKE node pool.
729+ spec: An instance of the NodePoolUpdateSpec class defining the
730+ update target and parameters.
619731
620732 Returns:
621733 A TimeUtil object representing the UTC timestamp when the operation started.
734+
735+ Raises:
736+ ValueError: If the target is unsupported.
622737 """
623- operation_start_time = TimeUtil .from_datetime (
624- datetime .datetime .now (datetime .timezone .utc )
625- )
738+ flags : list [str ] = []
626739
627- if not node_labels :
628- logging .info ("The specified label is empty, nothing to update." )
629- return operation_start_time
740+ match spec .target :
741+ case UpdateTarget .DISK_SIZE :
742+ current_disk_size = get_node_pool_disk_size (node_pool = node_pool )
743+ updated_disk_size = current_disk_size + spec .delta
744+ flags .append (f"--{ spec .target .value } { updated_disk_size } )
630745
631- labels = []
746+ case UpdateTarget .LABEL :
747+ current_labels = get_node_pool_labels (node_pool = node_pool )
748+ updated_labels = []
749+ for key , val in spec .delta .items ():
750+ if current_labels .get (key ) == val :
751+ val += val
752+ updated_labels .append (f"{ key } { val } )
753+ flags .append (f"--{ spec .target .value } { ',' .join (updated_labels )} )
632754
633- for key , val in node_labels . items () :
634- labels . append (f"{ key } = { val } )
755+ case _ :
756+ raise ValueError (f"Unsupported target: { spec . target } )
635757
636- command = (
758+ flags_str = " " .join (flags )
759+
760+ update_cmd = (
637761 f"gcloud container node-pools update { node_pool .node_pool_name }
638762 f"--project={ node_pool .project_id }
639763 f"--cluster={ node_pool .cluster_name }
640764 f"--location={ node_pool .location }
641- f"--labels= { ',' . join ( labels ) }
642- "--quiet "
765+ f"--quiet "
766+ f" { flags_str }
643767 )
644768
645- subprocess .run_exec (command )
769+ operation_start_time = TimeUtil .from_datetime (
770+ datetime .datetime .now (datetime .timezone .utc )
771+ )
772+
773+ subprocess .run_exec (update_cmd )
646774 return operation_start_time
0 commit comments