migrate_group_annotations.py

#!/usr/bin/env python2

import argparse

import elasticsearch
from elasticsearch import helpers


def migrate_annotations(host, port, index, group_hashid, group_pubid):
    es = elasticsearch.Elasticsearch([{'host': host, 'port': port}])
    query = {
        "query": {
            "filtered": {
                "filter": {
                    "term": {"group": group_hashid}
                }
            }
        }
    }
    annotations = helpers.scan(client=es,
                               index=index,
                               doc_type="annotation",
                               query=query)
    actions = []

    for ann in annotations:
        updates = {}
        updates["group"] = group_pubid
        updates["permissions"] = map_permissions(ann["_source"]["permissions"],
                                                 group_hashid, group_pubid)

        actions.append({"_op_type": "update",
                        "_index": index,
                        "_type": "annotation",
                        "_id": ann["_id"],
                        "doc": updates})

    return helpers.bulk(client=es, actions=actions)


def map_permissions(permissions, hashid, pubid):
    result = {}
    for action, principals in permissions.items():
        result[action] = ['group:' + pubid if p == 'group:' + hashid else p
                          for p in principals]
    return result


def load_map(fp):
    result = {}
    for line in fp:
        k, v = line.split()
        k = int(k)
        result[k] = v
    return result


def main():
    parser = argparse.ArgumentParser(
        description="Update annotation group identifiers")
    parser.add_argument("host")
    parser.add_argument("port")
    parser.add_argument("index")
    parser.add_argument("hashid_map", type=argparse.FileType('r'))
    parser.add_argument("pubid_map", type=argparse.FileType('r'))
    args = parser.parse_args()

    hashids = load_map(args.hashid_map)
    pubids = load_map(args.pubid_map)

    for pk in hashids:
        if pk not in pubids:
            raise RuntimeError("mapping missing for group id %d" % pk)

    for pk, group_hashid in hashids.items():
        group_pubid = pubids[pk]
        migrate_annotations(args.host,
                            args.port,
                            args.index,
                            group_hashid,
                            group_pubid)


if __name__ == '__main__':
    main()