Datashare task manager migration script

migrate_redis_task_manager.py

import asyncio
import sys
from json import loads, dumps

import redis.asyncio as redis


async def main(argv):
    pool = redis.ConnectionPool.from_url("redis://redis")
    client = redis.Redis.from_pool(pool)
    tasks = await client.hgetall("ds:task:manager")
    for k, v in tasks.items():
        task = loads(v)
        if b"BatchDownloadRunner" in k:
            name, id = task['name'].split('@')
            task['name'], task['id'] = name, id
            task['properties']['@class'] = "java.util.Collections$UnmodifiableMap"
            if 'result' in task:
                task['result'] = {
                    'file': task['result'],
                    '@class': 'org.icij.datashare.tasks.FileResult',
                    'size': task['properties']['batchDownload']['zipSize']
                }
                del task['properties']['batchDownload']['zipSize']

            task['properties']['batchDownload']['query'] = {
                '@class': "org.icij.datashare.text.indexing.SearchQuery",
                'query': task['properties']['batchDownload']['query']
            }
            if task['state'] != 'ERROR':
                await client.hset('ds:task:manager2', key=id, value=dumps(task))
        else:
            await client.hset('ds:task:manager2', key=k, value=v)


if __name__ == '__main__':
    asyncio.get_event_loop().run_until_complete(main(sys.argv))