Amazon Elasticsearch Service に放り込んだレストランデータを検索するコマンドラインツールサンプル with dd-trace-py

restaurant-search.py

# -*- coding: utf-8 -*-

import elasticsearch
import json
import os
import time
import datetime
import argparse
import sys
import codecs
sys.stdout = codecs.getwriter('utf-8')(sys.stdout)
#
from ddtrace import Pin, patch, patch_all
patch(elasticsearch=True)
from ddtrace import tracer
#
ESS_ENDPOINT = 'https://search-xxxxxxxxxxxxxxxxxxxxxxxxxxxxx.ap-northeast-1.es.amazonaws.com:443/'

@tracer.wrap('create_ess_connetcion', service='sample-app')
def es():
    return elasticsearch.Elasticsearch([ESS_ENDPOINT], verify_certs=False)

@tracer.wrap('result_ouput', service='sample-app')
def output(page):
    for hit in page:
        print hit['_source']['name'] + ' | ' + hit['_source']['address'] + ' | ' + hit['_source']['description']

@tracer.wrap('scroll_index', service='sample-app')
def scroll_index(scroll_size, sid):
    try:
        while (scroll_size > 0):
            page = es().scroll(scroll_id=sid, scroll='2m')
            sid = page['_scroll_id']
            scroll_size = len(page['hits']['hits'])
            output(page['hits']['hits'])
    except Exception, e:
        print e

@tracer.wrap('search_index', service='sample-app')
def search_index(name, keyword=None, location=None):
    if location == None:
       body = { "query": { "bool": { "must": [{ "term": { "name": name }}]}}}
    elif keyword == None:
       body = { "query": { "bool": { "must": [{ "term": { "name": name }}, { "term": { "address": location }}]}}}
    else:
       body = { "query": { "bool": { "must": [{ "term": { "name": name }}, { "term": { "description": keyword}}, { "term": { "address": location }}]}}}

    try:
        page = es().search(index='ldgourmet', body=body, scroll='2m', size=1000)
        sid = page['_scroll_id']
        scroll_size = page['hits']['total']
        output(page['hits']['hits'])

        scroll_index(scroll_size, sid)
    except Exception, e:
        print e

    span = tracer.current_span()
    span.set_tag('name', unicode(name, 'utf-8', 'ignore'))
    span.set_tag('keyword', unicode(keyword, 'utf-8', 'ignore')) if keyword != None else ''
    span.set_tag('location', unicode(location, 'utf-8', 'ignore')) if location != None else ''

if __name__ in '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('-k', '--keyword', type=str, default=None)
    p.add_argument('-l', '--location', type=str, default=None)
    p.add_argument('-n', '--name', type=str, default=None)
    args = p.parse_args()

    if args.name == None:
        print 'Please input --name option.'
        exit(1)
    else:
        search_index(args.name, args.keyword, args.location)