Search for UBC handle

handle-search.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Script to connect to an Elasticsearch instance to basically, perform a
search for ANY phrase across the index.

Bzsed on: https://gist.github.com/ross-spencer/895b5a346729075dd98f76cd5314728c
"""

from __future__ import print_function
import collections
import json
import logging
import sys

from elasticsearch import Elasticsearch, NotFoundError

# Initial fields to complete for the rest of the script to work...
SEARCH_TERM = "\"hdl:2429/24\""
ELASTICSEARCH_SERVER = "http://127.0.0.1:62002"

# Logging configuration
FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
logging.basicConfig(format=FORMAT)
logger = logging.getLogger("aip.duplicates")
logger.setLevel("INFO")

# We are only interested in transferred objects.
TRANSFER_DIR = "%transferDirectory%objects"

# AIP Reference.
Aip = collections.namedtuple(
    "Aip", "aip original_name checksum checksum_type uuid"
)

# ES configuration
PAGING_SIZE = 100

# Connect to Elasticsearch
es = Elasticsearch([ELASTICSEARCH_SERVER])

# Edit this... can be pretty much anything in fact...
handle = SEARCH_TERM

filter_ = {
    "filter": {},
    "query": {
        "bool": {
            "must": [],
            "must_not": [],
            "should": [
                {
                    "query_string": {
                        "query": handle,
                        "fields": []
                    }
                }
            ]
        }
    }
}


def data_processed(mets):
    """Maintains a running total of all the data processed."""
    return sys.getsizeof(mets)


def pretty_print(mets_dict):
    """A nice wrapper to pretty print any JSON we need to debug."""
    return json.dumps(mets_dict, indent=4)


def count_results():
    """Count the number of results we're expecting to return form ES."""
    try:
        results = es.search(
            index="aips", doc_type="aip", body=filter_, from_=0, size=0
        )
    except NotFoundError:
        return 0

    total = results["hits"]["total"]
    logger.info("Returning %s results", total)
    return total


def get_mets(page=0):
    """Connect to Elasticsearch per number of pages we require and retrieve
    results.
    """
    logger.info("Connecting to ES for results, page %s", page)
    # We query the AIP doc_type because size is associated with that, but
    # perhaps the AIPFILE is more appropriate?
    try:
        # Connect to the ES client and perform a search.
        return (
            es.search(
                index="aips",
                doc_type="aip",
                body=filter_,
                from_=page,
                size=PAGING_SIZE,
            ),
            (page + PAGING_SIZE),
        )
    except NotFoundError:
        return None, None


def main():
    """Primary entry point for this script."""

    # happy path...
    expected_results = count_results()
    if expected_results is 0:
        return

    # happy path...
    results, page = get_mets()
    if not results:
        logger.info("no results to return")
        return

    print("AIP UUID, AIP NAME, SIZE")
    while results["hits"]["hits"]:
        for hit in results["hits"]["hits"]:
            try:
                name = hit["_source"]["name"]
                uuid = hit["_source"]["uuid"]
                size = hit["_source"]["size"]
                print("{}, {}, {}".format(uuid, name, size))
            except KeyError:
                pass
        if page >= expected_results:
            break
        results, page = get_mets(page)


if __name__ == "__main__":
    main()