adding logs api example for elasticsearch (#8)

* adding logs api example for elasticsearch

* updating existing elasticsearch extension with logs api framework

* adding copyright to elasticsearch_producer

* updated readme with updated log output

* removed dependency on the requests module

* removing outdate info from readme

* updating readme to include function deployment configuration requirements

Author: heeki

python-example-elasticsearch-extension/README.md

@@ -1,70 +1,38 @@
-# Example ElasticSearch Extension in Python
-The provided code sample demonstrates how to get a basic extension written in Python 3 up and running.
+# Example Logs API Extension in Python for Elasticsearch
+The provided code sample demonstrates how to get a basic Logs API extension for Elasticsearch written in Python 3 up and running.
 
-> Note: This extension requires the Python 3 runtime to be present in the Lambda execution environment of your function.
+> Note: This extension requires the Python 3 runtime to be present in the Lambda execution environment of your function. This example code is not production ready. Use it with your own discretion after testing thoroughly.
 
-There are two components to this sample:
-* `extensions/`: This sub-directory should be extracted to /opt/extensions where the Lambda platform will scan for executables to launch extensions
-* `python-example-elasticsearch-extension/`: This sub-directory should be extracted to /opt/python-example-extension which is referenced by the `extensions/python-example-elasticsearch-extension` executable and includes a Python executable along with all of its necessary dependencies.
+In this example, we start by developing a simple extension and then add the ability to read logs from the Logs API. For more details on building an extension, please read the Extension API Developer Guide.
 
-## Prep Python Dependencies
-Install the extension dependencies locally, which will be mounted along with the extension code.
-
-```bash
-$ cd python-example-elasticsearch-extension
-$ chmod +x extension.py
-$ pip3 install -r requirements.txt -t .
-$ cd ..
-```
+When the Lambda service sets up the execution environment, it runs the extension (logs_api_elasticsearch_extension.py). This extension first registers as an extension and then subscribes to the Logs API to receive the logs via HTTP protocol. It starts an HTTP listener which receives the logs and processes them.
 
 ## Layer Setup Process
-The extensions .zip file should contain a root directory called `extensions/`, where the extension executables are located and another root directory called `python-example-elasticsearch-extension/`, where the core logic of the extension  and its dependencies are located.
+The extensions .zip file should contain a root directory called `extensions/`, where the extension executables are located. The dependencies for the extension (logs_api_elasticsearch_extension.py) are found in the logs_api_elasticsearch_extension directory.
 
 Creating zip package for the extension:
 ```bash
-$ chmod +x extensions/python-example-elasticsearch-extension
+$ cd python-example-elasticsearch-extension
+$ chmod +x extensions/logs_api_elasticsearch_extension.py
 $ zip -r extension.zip .
 ```
 
-Ensure that you have aws-cli v2 for the commands below.
-Publish a new layer using the `extension.zip`. The output of the following command should provide you a layer arn.
+Publish a new layer using the `extension.zip`. The output of the following command should provides you a layer arn.
 ```bash
 aws lambda publish-layer-version \
- --layer-name "python-example-extension" \
+ --layer-name "python-example-elasticsearch-extension" \
  --region <use your region> \
  --zip-file  "fileb://extension.zip"
 ```
 Note the LayerVersionArn that is produced in the output.
-eg. `"LayerVersionArn": "arn:aws:lambda:<region>:123456789012:layer:<layerName>:1"`
-
-Add the newly created layer version to a Python 3.8 runtime Lambda function. Ensure to include environment variables like `ES_ENDPOINT`="ec2-XXX-XXX-XXX-XXX.compute-1.amazonaws.com" and `ES_INDEX`="extensions"
-
+e.g. `"LayerVersionArn": "arn:aws:lambda:<region>:123456789012:layer:python-example-elasticsearch-extension:1"`
 
-## Function Invocation and Extension Execution
-
-When invoking the function, you should now see log messages from the example extension similar to the following:
-```
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    EXTENSION Name: python-example-elasticsearch-extension State: Ready Events: [INVOKE,SHUTDOWN]
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    START RequestId: 9ca08945-de9b-46ec-adc6-3fe9ef0d2e8d Version: $LATEST
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    python-example-elasticsearch-extension launching extension
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Registering...
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Registered with ID: 6ec8756c-4830-458b-9dda-156e5dda1cc1
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Waiting for event...
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Received event: {"eventType": "INVOKE", "deadlineMs": 1596756305517, "requestId": "4b61d2be-3ba0-4e99-9121-30268b462c77", "invokedFunctionArn": "", "tracing": {"type": "X-Amzn-Trace-Id", "value": ""}}
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Sending payload: {"functionName": "defaultFunctionName", "functionVersion": "$LATEST", "requestId": "4b61d2be-3ba0-4e99-9121-30268b462c77", "waitDuration": 3787.946044921875}
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Attempting POST to: https://ec2-XXX-XXX-XXX-XXX.compute-1.amazonaws.com/extensions/_doc
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Response: {"_index":"extensions","_type":"_doc","_id":"eW8WxnMB4bYnBqSFvK5w","_version":1,"result":"created","_shards":{"total":2,"successful":1,"failed":0},"_seq_no":0,"_primary_term":1}
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    [python-example-elasticsearch-extension] Waiting for event...
-    ...
-    ...
-    Function logs...
-    ...
-    ...
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    END RequestId: 9ca08945-de9b-46ec-adc6-3fe9ef0d2e8d
-    XXXX-XX-XXTXX:XX:XX.XXX-XX:XX    REPORT RequestId: 9ca08945-de9b-46ec-adc6-3fe9ef0d2e8d Duration: 80.36 ms Billed Duration: 100 ms Memory Size: 128 MB Max Memory Used: 67 MB Init Duration: 297.83 ms
+Add the newly created layer version to a Python 3.8 runtime Lambda function.
+```bash
+aws lambda update-function-configuration --region <use your region> --function-name <your function name> --layers <LayerVersionArn from previous step>
 ```
 
-
-
-
-
+## Lambda Function Deployment
+When deploying the Lambda function, be sure to include configure two environment variables that the extension will leverage for communicating with the Elasticsearch cluster.
+ES_ENDPOINT: The endpoint for the Elasticsearch cluster, e.g. elasticsearch.example.com (don't need to include https:// or the trailing /). This endpoint must be reachable from the Lambda function.
+ES_INDEX: The index to which the extension should write log outputs.

python-example-elasticsearch-extension/extensions/logs_api_elasticsearch_extension.py

@@ -0,0 +1,104 @@
+#!/bin/sh
+''''exec python -u -- "$0" ${1+"$@"} # '''
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+import json
+import os
+
+from logs_api_elasticsearch_extension.http_listener import http_server_init, RECEIVER_PORT
+from logs_api_elasticsearch_extension.logs_api_client import LogsAPIClient
+from logs_api_elasticsearch_extension.extensions_api_client import ExtensionsAPIClient
+from logs_api_elasticsearch_extension.elasticsearch_producer import ElasticsearchProducer
+from queue import Queue
+
+"""Here is the sample extension code that stitches all of this together.
+    - The extension will run two threads. The "main" thread, will register to ExtensionAPI and process its invoke and
+    shutdown events (see next call). The second "listener" thread will listen for HTTP Post events that deliver log batches.
+    - The "listener" thread will place every log batch it receives in a synchronized queue; during each execution slice,
+    the "main" thread will make sure to process any event in the queue before returning control by invoking next again.
+    - Note that because of the asynchronous nature of the system, it is possible that logs for one invoke are
+    processed during the next invoke slice. Likewise, it is possible that logs for the last invoke are processed during
+    the SHUTDOWN event.
+
+Note: 
+
+1.  This is a simple example extension to make you help start investigating the Lambda Runtime Logs API.
+    This code is not production ready, and it has never intended to be. Use it with your own discretion after you tested
+    it thoroughly.  
+
+2.  The extension code is starting with a shebang this is to bring Python runtime to the execution environment.
+    This works if the lambda function is a python3.x function therefore it brings python3.x runtime with itself.
+    It may not work for python 2.7 or other runtimes. 
+    The recommended best practice is to compile your extension into an executable binary and not rely on the runtime.
+  
+3.  This file needs to be executable, so make sure you add execute permission to the file 
+    `chmod +x logs_api_http_extension.py`
+
+"""
+
+class LogsAPIHTTPExtension():
+    def __init__(self, agent_name, registration_body, subscription_body):
+        print(f"Initializing LogsAPIExternalExtension {agent_name}")
+        self.agent_name = agent_name
+        self.queue = Queue()
+        self.logs_api_client = LogsAPIClient()
+        self.extensions_api_client = ExtensionsAPIClient()
+        self.es_endpoint = os.environ['ES_ENDPOINT']
+        self.es_index = os.environ['ES_INDEX']
+        self.es_producer = ElasticsearchProducer(self.agent_name, self.es_endpoint, self.es_index)
+
+        # Register early so Runtime could start in parallel
+        self.agent_id = self.extensions_api_client.register(self.agent_name, registration_body)
+
+        # Start listening before Logs API registration
+        http_server_init(self.queue)
+        self.logs_api_client.subscribe(self.agent_id, subscription_body)
+
+    def run_forever(self):
+        print(f"Serving LogsAPIHTTPExternalExtension {self.agent_name}")
+        while True:
+            resp = self.extensions_api_client.next(self.agent_id)
+            # Process the received batches if any.
+            while not self.queue.empty():
+                batch = self.queue.get_nowait()
+                # This line logs the events received to CloudWatch.
+                # Replace it to send logs to elsewhere.
+                # If you've subscribed to extension logs, e.g. "types": ["platform", "function", "extension"],
+                # you'll receive the logs of this extension back from Logs API.
+                # And if you log it again with the line below, it will create a cycle since you receive it back again.
+                # Use `extension` log type if you'll egress it to another endpoint,
+                # or make sure you've implemented a protocol to handle this case.
+                for item in batch:
+                    self.es_producer.send(item)
+
+# Register for the INVOKE events from the EXTENSIONS API
+_REGISTRATION_BODY = {
+    "events": ["INVOKE", "SHUTDOWN"],
+}
+
+# Subscribe to platform logs and receive them on ${local_ip}:4243 via HTTP protocol.
+TIMEOUT_MS = 1000 # Maximum time (in milliseconds) that a batch would be buffered.
+MAX_BYTES = 262144 # Maximum size in bytes that the logs would be buffered in memory.
+MAX_ITEMS = 10000 # Maximum number of events that would be buffered in memory.
+
+_SUBSCRIPTION_BODY = {
+    "destination":{
+        "protocol": "HTTP",
+        "URI": f"http://sandbox:{RECEIVER_PORT}",
+    },
+    "types": ["platform", "function"],
+    "buffering": {
+        "timeoutMs": TIMEOUT_MS,
+        "maxBytes": MAX_BYTES,
+        "maxItems": MAX_ITEMS
+    }
+}
+
+def main():
+    print(f"Starting Extensions {_REGISTRATION_BODY} {_SUBSCRIPTION_BODY}")
+    # Note: Agent name has to be file name to register as an external extension
+    ext = LogsAPIHTTPExtension(os.path.basename(__file__), _REGISTRATION_BODY, _SUBSCRIPTION_BODY)
+    ext.run_forever()
+
+if __name__ == "__main__":
+    main()

python-example-elasticsearch-extension/extensions/logs_api_elasticsearch_extension/__init__.py

@@ -0,0 +1,29 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+import json
+import sys
+import urllib.request
+
+class ElasticsearchProducer():
+    def __init__(self, agent_name, endpoint, index):
+        self.agent_name = agent_name
+        self.endpoint = endpoint
+        self.index = index
+
+    def send(self, payload):
+        url = f"https://{self.endpoint}/{self.index}/_doc"
+        try:
+            if isinstance(payload["record"], str):
+                converted = payload
+                converted["record"] = json.loads(payload["record"].replace("'",'"').rstrip())
+            else:
+                converted = payload
+            req = urllib.request.Request(url)
+            req.method = "POST"
+            req.add_header("Content-Type", "application/json")
+            req.data = json.dumps(converted).encode("utf-8")
+            resp = urllib.request.urlopen(req)
+        except urllib.request.HTTPError as e:
+            print(f"[{self.agent_name}] HTTPError: {e}", flush=True)
+            sys.exit(1)

python-example-elasticsearch-extension/extensions/logs_api_elasticsearch_extension/elasticsearch_producer.py

@@ -0,0 +1,62 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+import json
+import os
+import sys
+import urllib.request
+
+""" Demonstrates code to register your extension as an extension.
+"""
+
+LAMBDA_AGENT_NAME_HEADER_KEY = "Lambda-Extension-Name"
+LAMBDA_AGENT_IDENTIFIER_HEADER_KEY = "Lambda-Extension-Identifier"
+
+class ExtensionsAPIClient():
+    def __init__(self):
+        try:
+            runtime_api_address = os.environ['AWS_LAMBDA_RUNTIME_API']
+            self.runtime_api_base_url = f"http://{runtime_api_address}/2020-01-01/extension"
+        except Exception as e:
+            raise Exception(f"AWS_LAMBDA_RUNTIME_API is not set {e}") from e
+
+    # Register as early as possible - the runtime initialization starts after all extensions have registered.
+    def register(self, agent_unique_name, registration_body):
+        try:
+            print(f"Registering to ExtensionsAPIClient on {self.runtime_api_base_url}")
+            req = urllib.request.Request(f"{self.runtime_api_base_url}/register")
+            req.method = 'POST'
+            req.add_header(LAMBDA_AGENT_NAME_HEADER_KEY, agent_unique_name)
+            req.add_header("Content-Type", "application/json")
+            data = json.dumps(registration_body).encode("utf-8")
+            req.data = data
+            resp = urllib.request.urlopen(req)
+            if resp.status != 200:
+                print(f"/register request to ExtensionsAPIClient failed. Status:  {resp.status}, Response: {resp.read()}")
+                # Fail the extension
+                sys.exit(1)
+            agent_identifier = resp.headers.get(LAMBDA_AGENT_IDENTIFIER_HEADER_KEY)
+            return agent_identifier
+        except Exception as e:
+            raise Exception(f"Failed to register to ExtensionsAPIClient: on {self.runtime_api_base_url}/register \
+                with agent_unique_name:{agent_unique_name}  \
+                and registration_body:{registration_body}\nError: {e}") from e
+
+    # Call the following method when the extension is ready to receive the next invocation
+    # and there is no job it needs to execute beforehand.
+    def next(self, agent_id):
+        try:
+            req = urllib.request.Request(f"{self.runtime_api_base_url}/event/next")
+            req.method = 'GET'
+            req.add_header(LAMBDA_AGENT_IDENTIFIER_HEADER_KEY, agent_id)
+            req.add_header("Content-Type", "application/json")
+            resp = urllib.request.urlopen(req)
+            if resp.status != 200:
+                print(f"/event/next request to ExtensionsAPIClient failed. Status: {resp.status}, Response: {resp.read()} ")
+                # Fail the extension
+                sys.exit(1)
+            data = resp.read()
+            print(f"Received response from ExtensionsAPIClient: {data}")
+            return data
+        except Exception as e:
+            raise Exception(f"Failed to get /event/next from ExtensionsAPIClient: {e}") from e

python-example-elasticsearch-extension/extensions/logs_api_elasticsearch_extension/extensions_api_client.py

@@ -0,0 +1,64 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+import json
+import sys
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from threading import Event, Thread
+
+""" Demonstrates code to set up an HTTP listener and receive log events
+"""
+
+RECEIVER_IP = "0.0.0.0"
+RECEIVER_PORT = 4243
+
+def http_server_init(queue):
+    def handler(*args):
+        LogsHandler(queue, *args)
+    print(f"Initializing HTTP Server on {RECEIVER_IP}:{RECEIVER_PORT}")
+    server = HTTPServer((RECEIVER_IP, RECEIVER_PORT), handler)
+
+    # Ensure that the server thread is scheduled so that the server binds to the port
+    # and starts to listening before subscribe for the logs and ask for the next event.
+    started_event = Event()
+    server_thread = Thread(target=serve, daemon=True, args=(started_event, server,))
+    server_thread.start()
+    rc = started_event.wait(timeout = 9)
+    if rc is not True:
+        raise Exception("server_thread has timedout before starting")
+
+
+# Server implementation
+class LogsHandler(BaseHTTPRequestHandler):
+    def __init__(self, queue, *args):
+        self.queue = queue
+        BaseHTTPRequestHandler.__init__(self, *args)
+
+    def do_POST(self):
+        try:
+            cl = self.headers.get("Content-Length")
+            if cl:
+                data_len = int(cl)
+            else:
+                data_len = 0
+            content = self.rfile.read(data_len)
+            self.send_response(200)
+            self.end_headers()
+            batch = json.loads(content.decode("utf-8"))
+            self.queue.put(batch)
+
+        except Exception as e:
+            print(f"Error processing message: {e}")
+
+# Server thread
+def serve(started_event, server):
+    # Notify that this thread is up and running
+    started_event.set()
+    try:
+        print(f"Serving HTTP Server on {RECEIVER_IP}:{RECEIVER_PORT}")
+        server.serve_forever()
+    except:
+        print(f"Error in HTTP server {sys.exc_info()[0]}")
+    finally:
+        if server:
+            server.shutdown()

python-example-elasticsearch-extension/extensions/logs_api_elasticsearch_extension/http_listener.py

@@ -0,0 +1,40 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+import json
+import os
+import sys
+import urllib.request
+
+""" Demonstrates code to call the Logs API to subscribe to log events
+"""
+
+LAMBDA_AGENT_IDENTIFIER_HEADER_KEY = "Lambda-Extension-Identifier"
+
+class LogsAPIClient:
+    def __init__(self):
+        try:
+            runtime_api_address = os.environ['AWS_LAMBDA_RUNTIME_API']
+            self.logs_api_base_url = f"http://{runtime_api_address}/2020-08-15"
+        except Exception as e:
+            raise Exception(f"AWS_LAMBDA_RUNTIME_API is not set {e}") from e
+
+    # Method to call the Logs API to subscribe to log events.
+    def subscribe(self, agent_id, subscription_body):
+        try:
+            print(f"Subscribing to Logs API on {self.logs_api_base_url}")
+            req = urllib.request.Request(f"{self.logs_api_base_url}/logs")
+            req.method = 'PUT'
+            req.add_header(LAMBDA_AGENT_IDENTIFIER_HEADER_KEY, agent_id)
+            req.add_header("Content-Type", "application/json")
+            data = json.dumps(subscription_body).encode("utf-8")
+            req.data = data
+            resp = urllib.request.urlopen(req)
+            if resp.status != 200:
+                print(f"Could not subscribe to Logs API: {resp.status} {resp.read()}")
+                # Fail the extension
+                sys.exit(1)
+            print(f"Successfully subscribed to Logs API: {resp.read()}")
+        except Exception as e:
+            raise Exception(f"Failed to subscribe to Logs API on {self.logs_api_base_url} with id: {agent_id} \
+                and subscription_body: {json.dumps(subscription_body).encode('utf-8')} \nError:{e}") from e