Script to migrate users from one Algorithmia cluster to another

migrateUsers.py

"""
This module facilitates migrating users from one cluster to another. Note that
as with any user management features on the Algorithmia platform, you'll need
to use Admin API keys on both the source and destination clusters.

In testing mode, the users will be written to files for QA instead of being created on the cluster.

If any user migrations fail, those data will be captured in a file for further investigation or retries.

Non-standard library dependencies: `requests` (installable via `pip install requests`)

"""
import json
import os

import requests

##########################
#### BEGIN USER EDITS ####

# Set to `False` to actually create the users
# also will want to remove "testing counter" logic in `while` loop or set to high number
TESTING = True

# TODO: customize these as needed
#       SOURCE_CLUSTER_DOMAIN <-- cluster being migrated *from* ; goes with `$SOURCE_CLUSTER_ADMIN_API_KEY`
#       DEST_CLUSTER_DOMAIN <-- cluster being migrated *to* ; goes with `$DEST_CLUSTER_ADMIN_API_KEY`
SOURCE_ENDPOINT = "https://api.SOURCE_CLUSTER_DOMAIN/v1/users"
DEST_ENDPOINT = "https://api.DEST_CLUSTER_DOMAIN/v1/users"

# If desired, list subset of users to migrate; else, leave as empty list to migrate all users
user_subset = []
# user_subset = ["user1", "user2", "user3", "user4"]

# We recommend setting API keys as environment variables
# NOTE: User management requires admin API keys (not just standard API keys)
SOURCE_CLUSTER_ADMIN_API_KEY = os.getenv("SOURCE_CLUSTER_ADMIN_API_KEY")
DEST_CLUSTER_ADMIN_API_KEY = os.getenv("DEST_CLUSTER_ADMIN_API_KEY")

# If using a CSRF token, set here. Probably won't need this.
CSRF_TOKEN = ""

#### END USER EDITS ####
########################

# Headers for listing users
headers_list_users = {
    "Authorization": "Simple {}".format(SOURCE_CLUSTER_ADMIN_API_KEY),
}

# Headers for creating users
headers_create_user = {
    "Authorization": "Simple {}".format(DEST_CLUSTER_ADMIN_API_KEY),
    "Content-Type": "application/json",
}

if CSRF_TOKEN != "":
    headers_create_user["X-CSRF-Token"] = CSRF_TOKEN

### GET CURRENT USERS ###

# Query the source cluster and list out all users
source_users = []

# Initiate endpoint, which will change as we paginate through users
endpoint = SOURCE_ENDPOINT
testing_counter = 0
while (endpoint is not None) and (testing_counter < 3):
    print("Requesting users from: {}".format(endpoint))
    user_batch = requests.get(endpoint, headers=headers_list_users).json()
    source_users += user_batch["results"]
    endpoint = user_batch["next_link"]
    testing_counter += 1


# Make new user list with required and important params from old user list
dest_users = []

for source_user in source_users:
    if user_subset:
        if source_user["username"] not in user_subset:
            continue
    dest_user = {}
    # These params are required to create a user
    dest_user["username"] = source_user["username"]
    dest_user["email"] = source_user["email"]
    dest_user["fullname"] = source_user["fullname"]
    # Extract `company_name` and/or `company_role` if they were specified in
    # the previous user's info; these are not required params so this will
    # succeed even if they aren't specified
    company_name = source_user.get("company_name")
    company_role = source_user.get("company_role")
    if (company_name is not None) and (company_name != ""):
        dest_user["company_name"] = company_name
    if (company_role is not None) and (company_role != ""):
        dest_user["company_role"] = company_role
    # In older versions of Algorithmia, `passwordHash` and `shouldCreateHello`
    # weren't optional, so they're included here for backwards compatibility.
    # These parameters won't affect functionality; users should still reset
    # their password at first login by clicking the "Forgot Password?" button.
    dest_user["passwordHash"] = "DummyString"
    dest_user["shouldCreateHello"] = True
    dest_users.append(dest_user)

if TESTING:
    source_user_file = "source_users.json"
    dest_user_file = "dest_users.json"
    # Save intermediate data for QA.
    with open(source_user_file, "w") as outfile:
        json.dump(source_users, outfile)
    with open(dest_user_file, "w") as outfile:
        json.dump(dest_users, outfile)
    print("\n\nIn testing mode; exiting early. "
          "See `{}` and `{}` for QA.".format(source_user_file, dest_user_file))
    print("Change to `TESTING=False` to create users.")
    exit(0)

### CREATE NEW USERS ##

# Can also test creating just one user before doing the bulk migration
if TESTING:
    dest_users = [
        {
            "username": "testuser2109",
            "email": "testuser2109@sample.com",
            "fullname": "Mr. Test User",
        },
    ]

failed_user_additions = []

for dest_user in dest_users:
    dest_user_json = json.dumps(dest_user)
    response = requests.post(
        DEST_ENDPOINT, headers=headers_create_user, data=dest_user_json)
    if response.status_code >= 400:
        print("\ncould not create user: {}".format(dest_user["username"]))
        print(json.loads(response.text)["error"]["message"])
        failed_user_additions.append(dest_user)
    else:
        print("created user {}".format(dest_user["username"]))

failure_file = "failed_user_additions.json"
with open(failure_file, "w") as outfile:
    json.dump(failed_user_additions, outfile)
print("\n\nUse `{}` to debug and retry user migrations that failed.".format(
    failure_file))