Script to handle live migrations on GCP with Cockroach on a Debian 9 machine

handle-migration.py

import time

import requests

from subprocess import call


METADATA_URL = 'http://metadata.google.internal/computeMetadata/v1/'
METADATA_HEADERS = {'Metadata-Flavor': 'Google'}


def wait_for_maintenance(callback):
    url = METADATA_URL + 'instance/maintenance-event'
    last_maintenance_event = None
    last_etag = '0'

    while True:
        r = requests.get(
            url,
            params={'last_etag': last_etag, 'wait_for_change': True},
            headers=METADATA_HEADERS)

        # During maintenance the service can return a 503, so these should
        # be retried.
        if r.status_code == 503:
            time.sleep(1)
            continue
        r.raise_for_status()

        last_etag = r.headers['etag']

        if r.text == 'NONE':
            maintenance_event = None
        else:
            # Possible events:
            #   MIGRATE_ON_HOST_MAINTENANCE: instance will be migrated
            #   SHUTDOWN_ON_HOST_MAINTENANCE: instance will be shut down
            maintenance_event = r.text

        if maintenance_event != last_maintenance_event:
            last_maintenance_event = maintenance_event
            callback(maintenance_event)


def maintenance_callback(event):
    if event:
        print('Undergoing host maintenance: {}'.format(event))
        call("systemctl stop cockroach.service", shell=True)
        print('Cockroach shutdown')
    else:
        print('Finished host maintenance')

        call("sntp -S metadata.google.internal", shell=True)
        call("systemctl start cockroach.service", shell=True)

        print('Cockroach started')


def main():
    wait_for_maintenance(maintenance_callback)


if __name__ == '__main__':
    main()