irc-books/bookworm/batch_unpacker.py

import os
import json
import socket
import sys
import time
import tempfile
from threading import Thread
from subprocess import check_output
from bookworm import s3, parse
from bookworm.logger import log, setup_logger
from bookworm.constants import UNPACKABLE_EXTENSIONS, REDIS
import redis

def should_unpack(fname):
    fname = fname.lower()
    return fname.endswith('rar') or fname.endswith('zip')

def archive_contents(fd):
    to_extract = {}
    contents = check_output(['lsar', '-j', fd.name]).decode('utf-8')
    contents = json.loads(contents)
    
    log.debug(contents['lsarContents'])
    for f in contents['lsarContents']:
        fname = f['XADFileName']
        log.info('Found file %s in the archive', fname)
        if any([extension in fname.lower() for extension in ['.txt'] + UNPACKABLE_EXTENSIONS]):
            log.info('Extracting %s from the archive', fname)
            to_extract[f['XADIndex']] = fname
    return to_extract

def delete_raw_file(s3client, s3key, meta):
    log.info('Deleting %s from %s', s3key, meta['raw_file_bucket'])
    s3client.delete_object(Bucket=meta['raw_file_bucket'], Key=s3key)

def unpack_and_store(job_key, s3key, s3client, redis, meta):
    redis.hset(job_key, REDIS.STEP_KEY, 'UNPACKING')
    unpacked_files = unpack(s3key, s3client, redis, meta)
    redis.hset(job_key, REDIS.STEP_KEY, 'UNPACK_DONE')
    log.info('Done unpacking job %s', job_key)

    for fname, data in unpacked_files:
        log.info('Batch process file: %s', fname)
        try:
            decoded = data.decode('utf-8')
        except UnicodeDecodeError:
            decoded = data.decode('latin-1')

        books = parse.lines_to_dicts(decoded.replace('\r', '').splitlines())
        parse.insert_books(books)

    #delete_raw_file(s3client, s3key, meta)
    redis.delete(job_key)

def unpack(s3key, s3client, redis, meta):
    log.info('Got a request to unpack %s', s3key)
    data = s3client.get_object(Key=s3key, Bucket=meta['raw_file_bucket'])
    with tempfile.NamedTemporaryFile() as fd:
        raw_file_contents = data['Body'].read()
        fd.write(raw_file_contents)
        fd.flush()

        if not should_unpack(s3key):
            log.info("Not unpacking %s", s3key)
            return [(s3key, raw_file_contents)]

        to_extract = archive_contents(fd)
        if not to_extract:
            log.error("Could not find any valid file")
            return []

        ret = []
        for index, fname in to_extract.items():
            log.info('Processing %s %s', index, fname)
            file_contents = check_output(['unar', '-o', '-', '-i', fd.name, str(index)])
            log.info('Got %d bytes', len(file_contents))
            ret.append((fname, file_contents))
        return ret


def main():
    r = redis.StrictRedis(host='localhost', port=6379)
    setup_logger()
    s3client = s3.client()
    while True:
        log.info('Waiting for message on %s', REDIS.Q_PROCESS_BATCH_FILE)
        topic, message = r.blpop(REDIS.Q_PROCESS_BATCH_FILE)

        log.info('got message: %s', message)
        params = json.loads(message.decode('utf-8'))
        log.info('params for unpacker: %s', params)
        params['s3client'] = s3client
        params['redis'] = r

        t = Thread(target=unpack_and_store, kwargs=params)
        t.daemon = True
        t.start()