offload-sender/main.py

#!/usr/bin/env python
import base64
import datetime
import json
import logging
import os
import pathlib
import shutil
import time
from typing import Optional
import urllib.parse

import click
import minio
import requests

from progress import Progress

logging.basicConfig(level=logging.INFO)
BACKFEED_DELIM = "\n"


# TODO: Add rsync support
# TODO: Add rsync+ssh support
# TODO: Add webdav support.
# TODO: Fix the "ctrl-c handling" logic so it actually cleans up in the s3 bucket.

def retry_failures(fn, msg, *args, **kwargs):
    tries = 0
    while True:
        try:
            return fn(*args, **kwargs)
        except Exception:
            logging.exception(msg)
            delay = min(2 ** tries, 64)
            tries = tries + 1
            logging.info(f"Sleeping {delay} seconds...")
            time.sleep(delay)


@click.group()
def sender():
    pass


def watch_pass(input_directory: pathlib.Path, work_directory: pathlib.Path, ia_collection: str, ia_item_title: str,
               ia_item_prefix: str, ia_item_date: str, project: str, dispatcher: str, delete: bool, backfeed_key: str):
    logging.info("Checking for new items...")
    for original_directory in input_directory.iterdir():
        if original_directory.is_dir():
            original_name = original_directory.name
            new_directory = work_directory.joinpath(original_name)
            try:
                original_directory.rename(new_directory)
            except FileNotFoundError:
                logging.warning(f"Unable to move item {original_directory}")
                continue
            single_impl(new_directory, ia_collection, ia_item_title, ia_item_prefix, ia_item_date, project,
                        dispatcher, delete, backfeed_key)
            return True
    return False


@sender.command()
@click.option('--input-directory', envvar='UPLOAD_QUEUE_DIR', default="/data/upload-queue",
              type=click.Path(exists=True))
@click.option('--work-directory', envvar='UPLOADER_WORKING_DIR', default="/data/uploader-work",
              type=click.Path(exists=True))
@click.option('--ia-collection', envvar='IA_COLLECTION', required=True)
@click.option('--ia-item-title', envvar='IA_ITEM_TITLE', required=True)
@click.option('--ia-item-prefix', envvar='IA_ITEM_PREFIX', required=True)
@click.option('--ia-item-date', envvar='IA_ITEM_DATE', required=False)
@click.option('--project', envvar='PROJECT', required=True)
@click.option('--dispatcher', envvar='DISPATCHER', required=True)
@click.option('--delete/--no-delete', envvar='DELETE', default=False)
@click.option('--backfeed-key', envvar='BACKFEED_KEY', required=True)
def watch(input_directory: pathlib.Path, work_directory: pathlib.Path, ia_collection: str, ia_item_title: str,
          ia_item_prefix: str, ia_item_date: str, project: str, dispatcher: str, delete: bool, backfeed_key: str):
    if not isinstance(input_directory, pathlib.Path):
        input_directory = pathlib.Path(input_directory)
    if not isinstance(work_directory, pathlib.Path):
        work_directory = pathlib.Path(work_directory)

    while True:
        if not watch_pass(input_directory, work_directory, ia_collection, ia_item_title, ia_item_prefix, ia_item_date,
                          project, dispatcher, delete, backfeed_key):
            logging.info("No item found, sleeping...")
            time.sleep(10)


@sender.command()
@click.option('--item-directory', type=click.Path(exists=True), required=True)
@click.option('--ia-collection', envvar='IA_COLLECTION', required=True)
@click.option('--ia-item-title', envvar='IA_ITEM_TITLE', required=True)
@click.option('--ia-item-prefix', envvar='IA_ITEM_PREFIX', required=True)
@click.option('--ia-item-date', envvar='IA_ITEM_DATE', required=False)
@click.option('--project', envvar='PROJECT', required=True)
@click.option('--dispatcher', envvar='DISPATCHER', required=True)
@click.option('--delete/--no-delete', envvar='DELETE', default=False)
@click.option('--backfeed-key', envvar='BACKFEED_KEY', required=True)
def single(item_directory: pathlib.Path, ia_collection: str, ia_item_title: str, ia_item_prefix: str,
           ia_item_date: Optional[str], project: str, dispatcher: str, delete: bool, backfeed_key: str):
    single_impl(item_directory, ia_collection, ia_item_title, ia_item_prefix, ia_item_date, project, dispatcher, delete,
                backfeed_key)


def single_impl(item_directory: pathlib.Path, ia_collection: str, ia_item_title: str, ia_item_prefix: str,
                ia_item_date: Optional[str], project: str, dispatcher: str, delete: bool, backfeed_key: str):
    if not isinstance(item_directory, pathlib.Path):
        item_directory = pathlib.Path(item_directory)

    logging.info(f"Processing item {item_directory}...")

    if ia_item_date is None:
        s = item_directory.name.split("_")
        if len(s) > 0:
            ds = s[0]
            try:
                d = datetime.datetime.strptime(ds, "%Y%m%d%H%M%S")
                ia_item_date = d.strftime("%Y-%m")
            except ValueError:
                pass

    meta_json_loc = item_directory.joinpath('__upload_meta.json')
    if meta_json_loc.exists():
        raise Exception("META JSON EXISTS WTF")
    meta_json = {
        "IA_COLLECTION": ia_collection,
        "IA_ITEM_TITLE": f"{ia_item_title} {item_directory.name}",
        "IA_ITEM_DATE": ia_item_date,
        "IA_ITEM_NAME": f"{ia_item_prefix}{item_directory.name}",
        "PROJECT": project,
    }
    with open(meta_json_loc, 'w') as f:
        f.write(json.dumps(meta_json))
        logging.info("Wrote metadata json.")
    total_size = 0
    files = list(item_directory.glob("**/*"))
    for item in files:
        total_size = total_size + os.path.getsize(item)
    logging.info(f"Item size is {total_size} bytes across {len(files)} files.")
    meta_json["SIZE_HINT"] = str(total_size)

    def assign_target():
        logging.info("Attempting to assign target...")
        r = requests.get(f"{dispatcher}/offload_target", params=meta_json, timeout=60)
        if r.status_code == 200:
            data = r.json()
            return data["url"]
        else:
            raise Exception(f"Invalid status code {r.status_code}: {r.text}")

    url = retry_failures(assign_target, "Failed to fetch target")
    logging.info(f"Assigned target {url}")

    parsed_url = urllib.parse.urlparse(url)
    bf_item = None
    if parsed_url.scheme == "minio+http" or parsed_url.scheme == "minio+https":
        secure = (parsed_url.scheme == "minio+https")
        ep = parsed_url.hostname
        if parsed_url.port is not None:
            ep = f"{ep}:{parsed_url.port}"
        client = None

        def create_client():
            logging.info("Connecting to minio...")
            return minio.Minio(endpoint=ep, access_key=parsed_url.username, secret_key=parsed_url.password,
                               secure=secure)

        client = retry_failures(create_client, "Failed to connect to minio")

        bucket_name = item_directory.name.replace("_", "-")

        def make_bucket():
            logging.info("Attempting to make bucket...")
            if client.bucket_exists(bucket_name=bucket_name):
                raise Exception("Bucket already exists!")
            client.make_bucket(bucket_name=bucket_name)

        retry_failures(make_bucket, "Failed to make bucket")

        logging.info("Starting uploads...")
        for file in files:
            rel_file = file.relative_to(item_directory)

            def upload_file():
                logging.info(f"Uploading file {rel_file}...")
                client.fput_object(bucket_name=bucket_name, object_name=str(rel_file), file_path=file,
                                   progress=Progress())

            retry_failures(upload_file, f"Failed to upload {rel_file}")

        item_data = {"url": url, "item_name": item_directory.name, "bucket_name": bucket_name}
        bf_item_part = base64.urlsafe_b64encode(str(json.dumps(item_data)).encode("UTF-8")).decode("UTF-8")
        bf_item = f"{project}:{parsed_url.hostname}:{bf_item_part}"
    else:
        raise Exception("Unable to upload, don't understand url: {url}")

    if bf_item is None:
        raise Exception("Unable to create backfeed item")

    if backfeed_key == "SKIPBF":
        logging.warning(f"Skipping backfeed! Would have submitted: {bf_item}")
    else:
        def submit_item():
            u = f"https://legacy-api.arpa.li/backfeed/legacy/{backfeed_key}"
            logging.info(f"Attempting to submit bf item {bf_item} to {u}...")
            resp = requests.post(u, params={"skipbloom": "1", "delimiter": BACKFEED_DELIM},
                                 data=f"{bf_item}{BACKFEED_DELIM}".encode("UTF-8"), timeout=60)
            if resp.status_code != 200:
                raise Exception(f"Failed to submit to backfeed {resp.status_code}: {resp.text}")

        retry_failures(submit_item, "Failed to submit to backfeed")
        logging.info("Backfeed submit complete!")

    if delete:
        logging.info("Removing item...")
        shutil.rmtree(item_directory)
    logging.info("Upload complete!")


if __name__ == '__main__':
    sender()