From e614c14f1b44b1f5de19d0d78469eaf31222e7ba Mon Sep 17 00:00:00 2001
From: Erik Dunteman <44653944+erik-dunteman@users.noreply.github.com>
Date: Tue, 28 Nov 2023 15:15:26 -0600
Subject: [PATCH] Add request ID into request, and include headers (#46)

* version bump

* implemented

* using case insensitive headers
---
 potassium/potassium.py  | 12 ++++++++----
 setup.py                |  2 +-
 tests/test_endpoints.py | 21 +++++++++++++++++++++
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/potassium/potassium.py b/potassium/potassium.py
index c9ee28f..9ad355b 100644
--- a/potassium/potassium.py
+++ b/potassium/potassium.py
@@ -3,6 +3,7 @@
 from typing import Generator, Optional, Union
 from flask import Flask, request, make_response, abort, Response as FlaskResponse
 from werkzeug.serving import make_server
+from werkzeug.datastructures.headers import EnvironHeaders
 from threading import Thread, Lock, Condition
 import functools
 import traceback
@@ -15,18 +16,19 @@ def __init__(self, type, func):
         self.type = type
         self.func = func
 
-
 class Request():
-    def __init__(self, json: dict):
+    def __init__(self, id: str, headers: EnvironHeaders, json: dict):
+        self.id = id
+        self.headers = headers
         self.json = json
 
-
 ResponseBody = Union[bytes, Generator[bytes, None, None]]
 
 class Response():
     def __init__(self, status: int = 200, json: Optional[dict] = None, headers: Optional[dict] = None, body: Optional[ResponseBody] = None):
         assert json == None or body == None, "Potassium Response object cannot have both json and body set"
 
+
         self.headers = headers if headers != None else {}
 
         # convert json to body if not None
@@ -181,7 +183,9 @@ def _handle_generic(self, endpoint, flask_request):
 
         try:
             req = Request(
-                json=flask_request.get_json()
+                headers=flask_request.headers,
+                json=flask_request.get_json(),
+                id=flask_request.headers.get("X-Banana-Request-Id", "")
             )
         except:
             res = make_response()
diff --git a/setup.py b/setup.py
index ef686f2..50b80e7 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@
 setup(
     name='potassium',
     packages=['potassium'],
-    version='0.4.0',
+    version='0.4.1',
     license='Apache License 2.0',
     # Give a short description about your library
     description='The potassium package is a flask-like HTTP server for serving large AI models',
diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py
index 99c9e35..7622b02 100644
--- a/tests/test_endpoints.py
+++ b/tests/test_endpoints.py
@@ -53,6 +53,17 @@ def handler2_id(context: dict, request: potassium.Request) -> potassium.Response
             status=200
         )
 
+    @app.handler("/some_headers_request")
+    def handler5(context: dict, request: potassium.Request) -> potassium.Response:
+        assert request.headers["A"] == "a"
+        assert request.headers["B"] == "b"
+        assert request.headers["X-Banana-Request-Id"] == request.id
+        return potassium.Response(
+            headers={"A": "a", "B": "b", "X-Banana-Request-Id": request.id},
+            json={"hello": "some_headers_request", "id": request.id},
+            status=200
+        )
+
     client = app.test_client()
 
     res = client.post("/", json={})
@@ -77,8 +88,18 @@ def handler2_id(context: dict, request: potassium.Request) -> potassium.Response
     assert res.status_code == 200
     assert res.json == {"hello": "some_path/child_path"}
 
+    # note the capitalization of ID, we're testing that it's case insensitive
+    headers = {"A": "a", "B": "b", "X-Banana-Request-ID": "123"} 
+    res = client.post("/some_headers_request", json={}, headers=headers)
+    assert res.status_code == 200
+    assert res.json == {"hello": "some_headers_request", "id": headers["X-Banana-Request-ID"]}
+    assert res.headers["X-Banana-Request-Id"] == "123"
+    assert res.headers["A"] == "a"
+    assert res.headers["B"] == "b"
+
     res = client.post("/", data='{"key": unquoted_value}', content_type='application/json')
     assert res.status_code == 400
+    
     # check status
     res = client.get("/__status__")
     assert res.status_code == 200