db.py

import sqlalchemy

from sqlalchemy import Table, Column

import click

metadata = sqlalchemy.MetaData()

nodes = Table(
    "nodes",
    metadata,
    Column("id", sqlalchemy.BigInteger, primary_key=True),
    Column("node", sqlalchemy.String(64)),
    Column("available", sqlalchemy.Integer),
    Column("current_load", sqlalchemy.Integer),
    Column("capacity", sqlalchemy.Integer),
    Column("downed", sqlalchemy.Integer),
    Column("backoff", sqlalchemy.Integer),
)

users = Table(
    "users",
    metadata,
    Column("nodeid", sqlalchemy.BigInteger),
    Column("replaced_at", sqlalchemy.BigInteger),
)


def cli(f):
    @click.option("--tokenserver-db-host")
    @click.option("--tokenserver-db-user")
    @click.option("--tokenserver-db-password")
    @click.option("--tokenserver-db-database")
    def wrapped(
        tokenserver_db_host,
        tokenserver_db_user,
        tokenserver_db_password,
        tokenserver_db_database,
        *args,
        **kwargs
    ):
        return f(
            *args,
            engine=create_engine(
                tokenserver_db_host,
                tokenserver_db_user,
                tokenserver_db_password,
                tokenserver_db_database,
            ),
            **kwargs
        )

    return wrapped


def create_engine(host, user, password, database):
    return sqlalchemy.create_engine(
        "mysql+pymysql://{}:{}@{}/{}".format(user, password, host, database)
    )

Dockerfile

FROM kennethreitz/pipenv

COPY . /app

find_next_nodes.py

#!/usr/bin/env python3

import click
import db
import sync
import json

def list_nodes():
    return db.nodes.select()

@click.command()
@click.option("--count", help="The number of new node numbers to return", default=1)
@db.cli
def find_next_nodes(engine, count):
    with engine.begin() as connection:
        nodes = connection.execute(list_nodes())
        numbers = [sync.node_number_from_name(node.node) for node in nodes]
        last_node_number = max(numbers)
        next_numbers = list(range(last_node_number+1, last_node_number+1+count))
        print(json.dumps(next_numbers))


if __name__ == '__main__':
    find_next_nodes()

list_node_numbers.py

#!/usr/bin/env python3

import click
import db
import sync
import json

@click.command()
@db.cli
def find_next_nodes(engine):
    with engine.begin() as connection:
        nodes = connection.execute(db.nodes.select())
        numbers = [sync.node_number_from_name(node.node) for node in nodes]
        print(json.dumps(numbers))


if __name__ == '__main__':
    find_next_nodes()

Pipfile

[[source]]
url = "https://pypi.org/simple"
verify_ssl = true
name = "pypi"

[packages]
future = "*"
click = "*"
sqlalchemy = "*"
boto = "*"
pymysql = "*"

[dev-packages]

[requires]
python_version = "3.6"

Pipfile.lock

{
    "_meta": {
        "hash": {
            "sha256": "e7a915f92b6487d7cbfd35d7e53fd927ffc17b42d8c6ae89b5c3835e1bde1fda"
        },
        "pipfile-spec": 6,
        "requires": {
            "python_version": "3.6"
        },
        "sources": [
            {
                "name": "pypi",
                "url": "https://pypi.org/simple",
                "verify_ssl": true
            }
        ]
    },
    "default": {
        "asn1crypto": {
            "hashes": [
                "sha256:2f1adbb7546ed199e3c90ef23ec95c5cf3585bac7d11fb7eb562a3fe89c64e87",
                "sha256:9d5c20441baf0cb60a4ac34cc447c6c189024b6b4c6cd7877034f4965c464e49"
            ],
            "version": "==0.24.0"
        },
        "boto": {
            "hashes": [
                "sha256:13be844158d1bd80a94c972c806ec8381b9ea72035aa06123c5db6bc6a6f3ead",
                "sha256:deb8925b734b109679e3de65856018996338758f4b916ff4fe7bb62b6d7000d1"
            ],
            "index": "pypi",
            "version": "==2.48.0"
        },
        "cffi": {
            "hashes": [
                "sha256:151b7eefd035c56b2b2e1eb9963c90c6302dc15fbd8c1c0a83a163ff2c7d7743",
                "sha256:1553d1e99f035ace1c0544050622b7bc963374a00c467edafac50ad7bd276aef",
                "sha256:1b0493c091a1898f1136e3f4f991a784437fac3673780ff9de3bcf46c80b6b50",
                "sha256:2ba8a45822b7aee805ab49abfe7eec16b90587f7f26df20c71dd89e45a97076f",
                "sha256:3bb6bd7266598f318063e584378b8e27c67de998a43362e8fce664c54ee52d30",
                "sha256:3c85641778460581c42924384f5e68076d724ceac0f267d66c757f7535069c93",
                "sha256:3eb6434197633b7748cea30bf0ba9f66727cdce45117a712b29a443943733257",
                "sha256:495c5c2d43bf6cebe0178eb3e88f9c4aa48d8934aa6e3cddb865c058da76756b",
                "sha256:4c91af6e967c2015729d3e69c2e51d92f9898c330d6a851bf8f121236f3defd3",
                "sha256:57b2533356cb2d8fac1555815929f7f5f14d68ac77b085d2326b571310f34f6e",
                "sha256:770f3782b31f50b68627e22f91cb182c48c47c02eb405fd689472aa7b7aa16dc",
                "sha256:79f9b6f7c46ae1f8ded75f68cf8ad50e5729ed4d590c74840471fc2823457d04",
                "sha256:7a33145e04d44ce95bcd71e522b478d282ad0eafaf34fe1ec5bbd73e662f22b6",
                "sha256:857959354ae3a6fa3da6651b966d13b0a8bed6bbc87a0de7b38a549db1d2a359",
                "sha256:87f37fe5130574ff76c17cab61e7d2538a16f843bb7bca8ebbc4b12de3078596",
                "sha256:95d5251e4b5ca00061f9d9f3d6fe537247e145a8524ae9fd30a2f8fbce993b5b",
                "sha256:9d1d3e63a4afdc29bd76ce6aa9d58c771cd1599fbba8cf5057e7860b203710dd",
                "sha256:a36c5c154f9d42ec176e6e620cb0dd275744aa1d804786a71ac37dc3661a5e95",
                "sha256:a6a5cb8809091ec9ac03edde9304b3ad82ad4466333432b16d78ef40e0cce0d5",
                "sha256:ae5e35a2c189d397b91034642cb0eab0e346f776ec2eb44a49a459e6615d6e2e",
                "sha256:b0f7d4a3df8f06cf49f9f121bead236e328074de6449866515cea4907bbc63d6",
                "sha256:b75110fb114fa366b29a027d0c9be3709579602ae111ff61674d28c93606acca",
                "sha256:ba5e697569f84b13640c9e193170e89c13c6244c24400fc57e88724ef610cd31",
                "sha256:be2a9b390f77fd7676d80bc3cdc4f8edb940d8c198ed2d8c0be1319018c778e1",
                "sha256:ca1bd81f40adc59011f58159e4aa6445fc585a32bb8ac9badf7a2c1aa23822f2",
                "sha256:d5d8555d9bfc3f02385c1c37e9f998e2011f0db4f90e250e5bc0c0a85a813085",
                "sha256:e55e22ac0a30023426564b1059b035973ec82186ddddbac867078435801c7801",
                "sha256:e90f17980e6ab0f3c2f3730e56d1fe9bcba1891eeea58966e89d352492cc74f4",
                "sha256:ecbb7b01409e9b782df5ded849c178a0aa7c906cf8c5a67368047daab282b184",
                "sha256:ed01918d545a38998bfa5902c7c00e0fee90e957ce036a4000a88e3fe2264917",
                "sha256:edabd457cd23a02965166026fd9bfd196f4324fe6032e866d0f3bd0301cd486f",
                "sha256:fdf1c1dc5bafc32bc5d08b054f94d659422b05aba244d6be4ddc1c72d9aa70fb"
            ],
            "markers": "platform_python_implementation != 'pypy'",
            "version": "==1.11.5"
        },
        "click": {
            "hashes": [
                "sha256:29f99fc6125fbc931b758dc053b3114e55c77a6e4c6c3a2674a2dc986016381d",
                "sha256:f15516df478d5a56180fbf80e68f206010e6d160fc39fa508b65e035fd75130b"
            ],
            "index": "pypi",
            "version": "==6.7"
        },
        "cryptography": {
            "hashes": [
                "sha256:3f3b65d5a16e6b52fba63dc860b62ca9832f51f1a2ae5083c78b6840275f12dd",
                "sha256:5251e7de0de66810833606439ca65c9b9e45da62196b0c88bfadf27740aac09f",
                "sha256:551a3abfe0c8c6833df4192a63371aa2ff43afd8f570ed345d31f251d78e7e04",
                "sha256:5cb990056b7cadcca26813311187ad751ea644712022a3976443691168781b6f",
                "sha256:60bda7f12ecb828358be53095fc9c6edda7de8f1ef571f96c00b2363643fa3cd",
                "sha256:64b5c67acc9a7c83fbb4b69166f3105a0ab722d27934fac2cb26456718eec2ba",
                "sha256:6fef51ec447fe9f8351894024e94736862900d3a9aa2961528e602eb65c92bdb",
                "sha256:77d0ad229d47a6e0272d00f6bf8ac06ce14715a9fd02c9a97f5a2869aab3ccb2",
                "sha256:808fe471b1a6b777f026f7dc7bd9a4959da4bfab64972f2bbe91e22527c1c037",
                "sha256:9b62fb4d18529c84b961efd9187fecbb48e89aa1a0f9f4161c61b7fc42a101bd",
                "sha256:9e5bed45ec6b4f828866ac6a6bedf08388ffcfa68abe9e94b34bb40977aba531",
                "sha256:9fc295bf69130a342e7a19a39d7bbeb15c0bcaabc7382ec33ef3b2b7d18d2f63",
                "sha256:abd070b5849ed64e6d349199bef955ee0ad99aefbad792f0c587f8effa681a5e",
                "sha256:ba6a774749b6e510cffc2fb98535f717e0e5fd91c7c99a61d223293df79ab351",
                "sha256:c332118647f084c983c6a3e1dba0f3bcb051f69d12baccac68db8d62d177eb8a",
                "sha256:d6f46e862ee36df81e6342c2177ba84e70f722d9dc9c6c394f9f1f434c4a5563",
                "sha256:db6013746f73bf8edd9c3d1d3f94db635b9422f503db3fc5ef105233d4c011ab",
                "sha256:f57008eaff597c69cf692c3518f6d4800f0309253bb138b526a37fe9ef0c7471",
                "sha256:f6c821ac253c19f2ad4c8691633ae1d1a17f120d5b01ea1d256d7b602bc59887"
            ],
            "version": "==2.2.2"
        },
        "future": {
            "hashes": [
                "sha256:e39ced1ab767b5936646cedba8bcce582398233d6a627067d4c6a454c90cfedb"
            ],
            "index": "pypi",
            "version": "==0.16.0"
        },
        "idna": {
            "hashes": [
                "sha256:156a6814fb5ac1fc6850fb002e0852d56c0c8d2531923a51032d1b70760e186e",
                "sha256:684a38a6f903c1d71d6d5fac066b58d7768af4de2b832e426ec79c30daa94a16"
            ],
            "version": "==2.7"
        },
        "pycparser": {
            "hashes": [
                "sha256:99a8ca03e29851d96616ad0404b4aad7d9ee16f25c9f9708a11faf2810f7b226"
            ],
            "version": "==2.18"
        },
        "pymysql": {
            "hashes": [
                "sha256:08bc389a8d63708ceb883eac928d004d9ebf97fc973abe7ec1c1b4d0e2edfb3c",
                "sha256:8ee053bce47c97786ccf7d2e73e9063c214251cecef846b1219a9169f4f833eb"
            ],
            "index": "pypi",
            "version": "==0.9.0"
        },
        "six": {
            "hashes": [
                "sha256:70e8a77beed4562e7f14fe23a786b54f6296e34344c23bc42f07b15018ff98e9",
                "sha256:832dc0e10feb1aa2c68dcc57dbb658f1c7e65b9b61af69048abc87a2db00a0eb"
            ],
            "version": "==1.11.0"
        },
        "sqlalchemy": {
            "hashes": [
                "sha256:2d5f08f714a886a1382c18be501e614bce50d362384dc089474019ce0768151c"
            ],
            "index": "pypi",
            "version": "==1.2.8"
        }
    },
    "develop": {}
}

retire.py

#!/usr/bin/env python3
from future.standard_library import install_aliases
install_aliases()

import click
import db
import sqlalchemy
import boto.route53
import urllib.parse as urlparse
import logging

logging.basicConfig(level=logging.INFO)


def region_for_node_name(node_name):
    parsed = urlparse.urlparse(node_name)
    name = parsed.netloc.split(".")[0]
    # Names are formatted sync-$nodeid-$region
    region = "-".join(name.split("-")[2:])
    return region


def get_node_id(node_name):
    """Returns a query that returns the node ID for a given node name."""
    return (
        sqlalchemy.select([db.nodes.c.id]).where(db.nodes.c.node == node_name).limit(1)
    )


def set_node_down(node):
    """Returns a query that sets the node as downed in the token server database."""
    return (
        db.nodes.update()
        .where(db.nodes.c.id == node)
        .values(available=0, capacity=0, downed=1, backoff=1, current_load=0)
    )


def migrate_users(node):
    """Marks all users in the node as replaced."""
    return (
        db.users.update()
        .where(db.users.c.nodeid == node)
        .values(replaced_at=sqlalchemy.func.unix_timestamp() * 1000)
    )


def redirect_users(r53, node_name, to):
    """Creates a CNAME DNS record for the node, pointing all traffic to 401inator so that
    clients talk to the token server and get a new node assigned."""
    record_name = urlparse.urlparse(node_name).netloc + "."
    zone_name = ".".join(record_name.split(".")[1:])
    zone = r53.get_zone(zone_name)
    for record in zone.get_records():
        if record.name == record_name:
            zone.delete_record(record)
            logging.info("Deleted record {}".format(record.name))
    zone.add_cname(record_name, to)
    logging.info("Added {} CNAME record pointing to {}".format(record_name, to))


@click.command()
@click.option("--node-name", help="The name of the node to be retired")
@click.option(
    "--redirect-to",
    help="The domain name to redirect all requests to",
    default="401inator.sync.services.mozilla.com",
)
@db.cli
def retire(engine, node_name, redirect_to):
    region = region_for_node_name(node_name)
    r53 = boto.route53.connect_to_region(region)
    with engine.begin() as connection:
        with connection.begin() as trans:
            node_id = connection.execute(get_node_id(node_name)).fetchone()[0]
            logging.info("Got id {} for node name {}".format(node_id, node_name))
            connection.execute(set_node_down(node_id))
            logging.info("Marked node as downed")
            connection.execute(migrate_users(node_id))
            logging.info("Migrated users away from node")
    redirect_users(r53, node_name, redirect_to)
    logging.info("Node retired successfully")


if __name__ == "__main__":
    retire()

sync.py

from future.standard_library import install_aliases
install_aliases()
import urllib.parse as urlparse

def node_number_from_name(name):
    netloc = urlparse.urlparse(name).netloc
    return int(netloc.split('-')[1])