registrar.py

#!/usr/bin/env python
# -----------------------------------------------------------------------------
#
# Project: registrar.py
# Authors: Stephan Meissl <stephan.meissl@eox.at>
#
# -----------------------------------------------------------------------------
# Copyright (c) 2019 EOX IT Services GmbH
#
# Python script to register products.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies of this Software or works derived from this Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
# -----------------------------------------------------------------------------


import sys
import os
import argparse
import textwrap
import logging
import traceback
import redis

import lxml.etree
from swiftclient.service import SwiftService

import django
from django.db import transaction
from django.contrib.gis.geos import GEOSGeometry

path = os.path.join(os.getenv('INSTALL_DIR', "/var/www/pvs"), "pvs_instance")
if path not in sys.path:
    sys.path.append(path)

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "pvs_instance.settings")
django.setup()

from eoxserver.backends import access
from eoxserver.resources.coverages import models
from eoxserver.resources.coverages.registration.product import (
    ProductRegistrator
)
from eoxserver.resources.coverages.registration.registrators.gdal import (
    GDALRegistrator
)

logger = logging.getLogger(__name__)


def setup_logging(verbosity):
    # start logging setup
    # get command line level
    verbosity = verbosity
    if verbosity == 0:
        level = logging.CRITICAL
    elif verbosity == 1:
        level = logging.ERROR
    elif verbosity == 2:
        level = logging.WARNING
    elif verbosity == 3:
        level = logging.INFO
    else:
        level = logging.DEBUG
    logger.setLevel(level)
    sh = logging.StreamHandler()
    sh.setLevel(level)
    formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
    sh.setFormatter(formatter)
    logger.addHandler(sh)
    # finished logging setup


def add_mask(product):
    metadata_item = product.metadata_items.all()[0]
    with access.vsi_open(metadata_item) as f:
        tree = lxml.etree.parse(f)
    root = tree.getroot()
    wkt = tree.xpath(
        '//gsc:opt_metadata/gml:metaDataProperty/gsc:EarthObservationMetaData/eop:vendorSpecific/eop:SpecificInformation[eop:localAttribute/text() = "CF_POLY"]/eop:localValue/text()',
        namespaces=root.nsmap
    )[0]
    geometry = GEOSGeometry(wkt)
    mask_type = models.MaskType.objects.get(product_type=product.product_type)
    models.Mask.objects.create(
        product=product,
        mask_type=mask_type,
        geometry=geometry,
    )


def get_product_level(product):
    try:
        metadata_item = product.metadata_items.all()[0]
        with access.vsi_open(metadata_item) as f:
            tree = lxml.etree.parse(f)
        root = tree.getroot()
        xp = '/gsc:report/gsc:opt_metadata/gml:metaDataProperty/gsc:EarthObservationMetaData/eop:parentIdentifier/text()'
        product_type_name = tree.xpath(xp, namespaces=root.nsmap)[0]
        if product_type_name.endswith('Level_1'):
            return 'Level_1'
        if product_type_name.endswith('Level_3'):
            return 'Level_3'
        else:
            raise Exception('Invalid product type name %s' % product_type_name)
    except Exception as e:
        logger.warning(
            'Failed to determine product level for product %s, error was %s'
            % (product.identifier, e)
        )


class RegistrationError(Exception):
    pass


@transaction.atomic
def registrar(
    collection,
    objects_prefix, upload_container, replace=False, client=None, registered_set_key=None
):
    logger.info("Starting registration of product '%s'." % objects_prefix)

    metadata_package, data_package, has_vrt = None, None, None

    with SwiftService() as swift:
        list_parts_gen = swift.list(
            container=upload_container, options={"prefix": objects_prefix},
        )
        for page in list_parts_gen:
            print(page)
            if page["success"]:
                for item in page["listing"]:
                    if item["name"].endswith(".xml"):
                        metadata_package = item["name"]
                    elif item["name"].endswith(".TIF") or \
                            item["name"].endswith(".tif"):
                        if has_vrt is not True:
                            data_package = item["name"]
                    elif item["name"].endswith(".vrt"):
                        data_package = item["name"]
                        has_vrt = True
                    elif not item["name"].endswith(".tar"):
                        raise RegistrationError(
                            "Product with objects prefix '%s' has "
                            "wrong content '%s'."
                            % (objects_prefix, item["name"])
                        )
            else:
                raise RegistrationError(
                    "No product found with objects prefix '%s'."
                    % objects_prefix
                )

    if metadata_package is None or data_package is None:
        raise RegistrationError(
            "Product with objects prefix '%s' has missing content."
            % objects_prefix
        )

    product_type = data_package.split("/")[1]

    product, replaced = ProductRegistrator().register(
        metadata_locations=[[upload_container,
                             metadata_package, ], ],
        type_name="%s_Product_%s" % (collection, product_type),
        replace=replace,
        extended_metadata=True,
        mask_locations=None,
        package_path=None,
        overrides={},
    )
    collection = models.Collection.objects.get(
        identifier=collection
    )
    models.collection_insert_eo_object(collection, product)

    level = get_product_level(product)
    if level == 'Level_1':
        collection_level_1 = models.Collection.objects.get(
            identifier="%s_Level_1" % collection
        )
        models.collection_insert_eo_object(collection_level_1, product)
    elif level == 'Level_3':
        collection_level_3 = models.Collection.objects.get(
            identifier="%s_Level_3" % collection
        )
        models.collection_insert_eo_object(collection_level_3, product)

    report = GDALRegistrator().register(
        data_locations=[[upload_container, data_package, ], ],
        metadata_locations=[[upload_container,
                             metadata_package, ], ],
        coverage_type_name="RGBNir",
        overrides={
            "identifier": "%s__coverage" % product.identifier,
            "footprint": None,
        },
        replace=replace,
    )
    models.product_add_coverage(product, report.coverage)

    try:
        add_mask(product)
    except Exception as e:
        logger.info("Couldn't add mask.")
        logger.debug(traceback.format_exc())
        logger.warning("%s: %s\n" % (type(e).__name__, str(e)))

    if client is not None:
        logger.debug(
            "Storing times in redis queue '%s" % registered_set_key
        )
        client.sadd(
            registered_set_key, "%s/%s"
            % (
                product.begin_time.strftime("%Y%m%dT%H%M%S"),
                product.end_time.strftime("%Y%m%dT%H%M%S")
            )
        )

    logger.info(
        "Successfully finished registration of product '%s'." % objects_prefix
    )


def registrar_redis_wrapper(
    collection,
    upload_container,
    replace=False, host="localhost", port=6379,
    register_queue_key="register_queue",
    registered_set_key="registered_set",
):
    client = redis.Redis(
        host=host, port=port, charset="utf-8", decode_responses=True
    )
    while True:
        logger.debug("waiting for redis queue '%s'..." % register_queue_key)
        value = client.brpop(register_queue_key)
        try:
            registrar(
                collection,
                value[1],
                upload_container,
                replace=replace,
                client=client,
                registered_set_key=registered_set_key
            )
        except Exception as e:
            logger.debug(traceback.format_exc())
            logger.error("%s: %s\n" % (type(e).__name__, str(e)))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.description = textwrap.dedent("""\
    Register products.
    """)

    parser.add_argument(
        "--mode", default="standard", choices=["standard", "redis"],
        help=(
            "The mode to run the registrar. Either one-off (standard) or "
            "reading from a redis queue."
        )
    )
    parser.add_argument(
        "--objects-prefix", default=None,
        help=(
            "Prefix to objects holding the metadata and data of product."
        )
    )
    parser.add_argument(
        "--replace", action="store_true",
        help=(
            "Replace existing products instead of skipping the registration."
        )
    )
    parser.add_argument(
        "--redis-register-queue-key", default="register_queue"
    )
    parser.add_argument(
        "--redis-registered-set-key", default="registered_set"
    )
    parser.add_argument(
        "--redis-host", default="localhost"
    )
    parser.add_argument(
        "--redis-port", type=int, default=6379
    )

    parser.add_argument(
        "-v", "--verbosity", type=int, default=3, choices=[0, 1, 2, 3, 4],
        help=(
            "Set verbosity of log output "
            "(4=DEBUG, 3=INFO, 2=WARNING, 1=ERROR, 0=CRITICAL). (default: 3)"
        )
    )

    arg_values = parser.parse_args()

    setup_logging(arg_values.verbosity)

    collection = os.environ.get('COLLECTION')
    if collection is None:
        logger.critical("Collection environment variable not set.")
        sys.exit(1)

    upload_container = os.environ.get('UPLOAD_CONTAINER')
    if upload_container is None:
        logger.critical("UPLOAD_CONTAINER environment variable not set.")
        sys.exit(1)

    if arg_values.mode == "standard":
        registrar(
            collection,
            arg_values.objects_prefix,
            upload_container,
            replace=arg_values.replace,
        )
    else:
        registrar_redis_wrapper(
            collection,
            upload_container,
            replace=arg_values.replace,
            host=arg_values.redis_host,
            port=arg_values.redis_port,
            register_queue_key=arg_values.redis_register_queue_key,
            registered_set_key=arg_values.redis_registered_set_key,
        )