Test ElasticSearch Connectivity

es_test.py

"""
Simple test for checking if an EC2 instance or ECS task can access an
ElasticSearch domain. This was mainly lifted from the AWS documentation:
https://docs.aws.amazon.com/elasticsearch-service/latest/developerguide/es-request-signing.html#es-request-signing-python

To use this, first the requisite libraries must be available:
    $ pip install -U \
        boto3 \
        elasticsearch \
        requests \
        requests-aws4auth

Ensure to replace the `host` and `region` variables with appropriate ones.

If attempting to interactively troubleshoot an ECS task, keep in mind that all
the task's environment variables may not have been inherited by the interactive
session being used. In particularly, the
`AWS_CONTAINER_CREDENTIALS_RELATIVE_URI` is required for boto3 to know to
assume the TaskRoleArn. Otherwise, the EC2 instance's role (instance-profile)
will be assumed. There is a commented variable below if this is your use-case.
"""
from elasticsearch import Elasticsearch, RequestsHttpConnection
from requests_aws4auth import AWS4Auth
import boto3
import os

#os.environ['AWS_CONTAINER_CREDENTIALS_RELATIVE_URI'] = '/v2/credentials/UUID-TO-REPLACE'

# For example, my-test-domain.us-east-1.es.amazonaws.com
host = os.environ.get('ES_HOST', '')

# e.g. us-west-1
region = os.environ.get('AWS_DEFAULT_REGION', '')

assert host, \
    'Must configure ES_HOST env var'
assert region, \
    'Must configure AWS_DEFAULT_REGION env var'

service = 'es'
credentials = boto3.Session().get_credentials()
awsauth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service, session_token=credentials.token)

es = Elasticsearch(
    hosts = [{'host': host, 'port': 443}],
    http_auth = awsauth,
    use_ssl = True,
    verify_certs = True,
    connection_class = RequestsHttpConnection
)

if es.ping():
    print("Ping good")
else:
    print("Ping failed")

document = {
    "title": "Moneyball",
    "director": "Bennett Miller",
    "year": "2011"
}

if __name__ == '__main__':
    import sys

    try:
        cmd = sys.argv[1].lower()
    except IndexError:
        sys.exit(0)

    if 'p' in cmd:
        print('Putting...')
        es.index(index="movies", doc_type="movie", id="5", body=document)

    if 'g' in cmd:
        print('Getting...')
        print(es.get(index="movies", doc_type="movie", id="5"))

# vim: ft=python:expandtab:sw=4:ts=4