storrowed.boston/searchwkr.py

#!/usr/bin/env python3

import os
import sys
import json
from datetime import datetime

import requests

import inventory
import searchlib
import sbenv

def main(argv):
    # Setup
    searchlib.prep_nltk()

    # Actually do the query.
    today = datetime.now()
    recents_before = load_recents()
    res = searchlib.query_day(today, preloadurls=recents_before)
    print(json.dumps(res, indent='  '))

    # Save the urls we've seen now.
    recents_before['urls'].extend(map(lambda e: e['u'], res))
    save_recents(recents_before)

    # Now send them up.
    ok = True
    for e in res:
        print('uploading', e['u'], '...')
        try:
            send_item(today, e)
        except Exception as e:
            print('failed!', str(e))
            ok = False

    return 0 if ok else 1

def send_item(day, ent):
    body = {
        'date': day.strftime(inventory.DATE_FORMAT),
        'desc': ent
    }

    h = {
        'Authorization': 'Bearer %s' % sbenv.get_admin_key(),
    }

    surl = sbenv.get_submit_url()
    r = requests.post(surl, headers=h, data=json.dumps(body))
    r.raise_for_status()

def get_recents_path():
    return os.path.join(sbenv.get_worker_datadir(), 'recents.json')

# Loads the recents file, creating a new one if it's from a different day.
def load_recents():
    cdate = datetime.now().strftime(inventory.DATE_FORMAT)
    rp = get_recents_path()

    if not os.path.exists(rp):
        return {
            'cdate': cdate,
            'urls': []
        }

    with open(rp, 'r') as f:
        j = json.load(f)
        if j['cdate'] != cdate:
            return {
                'cdate': cdate,
                'urls': []
            }
        else:
            return j

# Saves the recents file.
def save_recents(recents):
    with open(get_recents_path(), 'w') as f:
        json.dump(recents, f, indent='  ')

if __name__ == '__main__':
    sys.exit(main(sys.argv))