pxify-gke.py

#!/usr/local/bin/python3.6

"""
pxify-gke:

This program takes an existing GKE cluster and makes it "Portworx-ready"
by adding one additional non-root, non-mounted disk.
[Specifying a GKE local SSD disk results in having it automatically mounted]

Prereqs:
Enable API Access (if not done already)
https://support.google.com/cloud/answer/6158841?hl=en&ref_topic=6262490

Download Google API Credentials file:
https://developers.google.com/identity/protocols/application-default-credentials?authuser=1

export GOOGLE_APPLICATION_CREDENTIALS=/Users/jeff/Downloads/project-XXX.json

Usage:
    pxify-gke --project_id PROJECT PROJECT_ID --cluster_id CLUSTER_ID --disk_size DISK_SIZE

Output:

Successful output will look like this:

PXIFYING Project :  jeff-k8s-box   Cluster:  gke-terraform-jeff-k8s
gke-terraform-jeff-k8s
        zone:  us-central1-b
        igroup:  gke-gke-terraform-jeff-k-default-pool-c8f38ca7-grp
            Adding  gke-gke-terraform-jeff-k-default-pool-c8f38ca7-g1h2-px-disk1  to instance:   gke-gke-terraform-jeff-k-default-pool-c8f38ca7-g1h2
            Waiting for Disk :  gke-gke-terraform-jeff-k-default-pool-c8f38ca7-g1h2-px-disk1  to be READY
            Disk :  gke-gke-terraform-jeff-k-default-pool-c8f38ca7-g1h2-px-disk1  is now READY
        zone:  us-central1-c
        igroup:  gke-gke-terraform-jeff-k-default-pool-4cb81159-grp
            Adding  gke-gke-terraform-jeff-k-default-pool-4cb81159-npc9-px-disk1  to instance:   gke-gke-terraform-jeff-k-default-pool-4cb81159-npc9
            Waiting for Disk :  gke-gke-terraform-jeff-k-default-pool-4cb81159-npc9-px-disk1  to be READY
            Disk :  gke-gke-terraform-jeff-k-default-pool-4cb81159-npc9-px-disk1  is now READY
        zone:  us-central1-a
        igroup:  gke-gke-terraform-jeff-k-default-pool-3201c10f-grp
            Adding  gke-gke-terraform-jeff-k-default-pool-3201c10f-v72g-px-disk1  to instance:   gke-gke-terraform-jeff-k-default-pool-3201c10f-v72g
            Waiting for Disk :  gke-gke-terraform-jeff-k-default-pool-3201c10f-v72g-px-disk1  to be READY
            Disk :  gke-gke-terraform-jeff-k-default-pool-3201c10f-v72g-px-disk1  is now READY


This has been tested to work on GKE clusters created through the GUI and through Terraform


TODO:
   - Probably need to update the InstanceTemplate rather than the instance itself.
   - Harden with better try/except handling
   - Allow for a variable number of disks

"""

import argparse
import os
import sys
import time
import re
from pprint import pprint
import googleapiclient.discovery
from googleapiclient.errors import Error, HttpError


def wait_for_ready(compute, project, zone, disk):
      while True:
           req = compute.disks().get(project=project, zone=zone, disk=disk).execute()
           if req['status'] != "READY":
               print ("            Waiting for Disk : ", disk, " to be READY")
               time.sleep(2)
           else:
                print ("            Disk : ", disk, " is now READY")
                return


def pxify (project, cluster, disk_size):
    print ("\n\nPXIFYING Project : ", project, "  Cluster: ", cluster)
    c = googleapiclient.discovery.build('container', 'v1')
    try:
        clusters = c.projects().zones().clusters().list(projectId=project,zone="-").execute()['clusters']
    except:
        print ("Can't obtain clusters from project")
        sys.exit(-1)

    for i in clusters:
       if i['name'] == cluster:
          print (i['name'])
          for ig in i['instanceGroupUrls']:
              igm = re.match ( r'.*/zones/(.*)/.*/(.*)$', ig)
              zone = igm[1]
              igroup = igm[2]
              print("        zone: ", zone)
              print("        igroup: ", igroup)
              #
              # Now list instances
              #
              compute = googleapiclient.discovery.build('compute', 'v1')
              instance = compute.instanceGroups().listInstances(project=project, zone=zone, instanceGroup=igroup, body="{}" ).execute()
              for i in instance['items']:
                   m = re.match(r'.*/(.*)$', i['instance'])
                   iname = m[1]
                   disk_name = iname + "-px-disk1"
                   print ("            Adding ", disk_name, " to instance:  ", iname)
                   disk_body = {
                           "kind": "compute#disk",
                           "name": disk_name,
                           "sizeGb": int(disk_size)
                   }
                   disk = compute.disks().insert(project=project, zone=zone, body=disk_body).execute()
                   # pprint (disk)
                   wait_for_ready (compute, project, zone, disk_name)
                   attach_body = {
                       "source": disk['targetLink'],
                       "autoDelete": "true",
                   }
                   response = compute.instances().attachDisk(project=project, zone=zone, instance=iname, body=attach_body).execute()
                   # pprint(response)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--project_id', help='Your Google Cloud project ID.', required=True)
    parser.add_argument('--cluster_id', help='Your GKE Cluster ID', required=True)
    parser.add_argument('--disk_size',  help="Size of disk to add (in GB)", required=True)

    args = parser.parse_args()

    pxify (args.project_id, args.cluster_id, args.disk_size)