feat: add support for API key validation for self-hosted

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.0.60
+
+* Enable self-hosted authorization using UNSTRUCTURED_API_KEY env variable
+
 ## 0.0.59
 
 * Bump unstructured to 0.11.0

README.md

@@ -323,6 +323,9 @@ As mentioned above, processing a pdf using `hi_res` is currently a slow operatio
 * `UNSTRUCTURED_PARALLEL_MODE_SPLIT_SIZE` - the number of pages to be processed in one request, default is `1`.
 * `UNSTRUCTURED_PARALLEL_RETRY_ATTEMPTS` - the number of retry attempts on a retryable error, default is `2`. (i.e. 3 attempts are made in total)
 
+#### Security
+You may also set the optional `UNSTRUCTURED_API_KEY` env variable to enable request validation for your self-hosted instance of Unstructured. If set, only requests including an `unstructured-api-key` header with the same value will be fulfilled. Otherwise, the server will return a 401 indicating that the request is unauthorized.
+
 #### Controlling Server Load
 Some documents will use a lot of memory as they're being processed. To mitigate OOM errors, the server will return a 503 if the host's available memory drops below 2GB. This is configurable with `UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB`.
 

prepline_general/api/app.py

@@ -11,7 +11,7 @@
 app = FastAPI(
     title="Unstructured Pipeline API",
     description="""""",
-    version="0.0.59",
+    version="0.0.60",
     docs_url="/general/docs",
     openapi_url="/general/openapi.json",
 )

prepline_general/api/general.py

@@ -675,7 +675,7 @@ def return_content_type(filename):
 
 
 @router.post("/general/v0/general")
-@router.post("/general/v0.0.59/general")
+@router.post("/general/v0.0.60/general")
 def pipeline_1(
     request: Request,
     gz_uncompressed_content_type: Optional[str] = Form(default=None),
@@ -697,6 +697,13 @@ def pipeline_1(
     new_after_n_chars: List[str] = Form(default=[]),
     max_characters: List[str] = Form(default=[]),
 ):
+    if api_key_env := os.environ.get("UNSTRUCTURED_API_KEY"):
+        api_key = request.headers.get("unstructured-api-key")
+        if api_key != api_key_env:
+            raise HTTPException(
+                detail=f"API key {api_key} is invalid", status_code=status.HTTP_401_UNAUTHORIZED
+            )
+
     if files:
         for file_index in range(len(files)):
             if files[file_index].content_type == "application/gzip":

preprocessing-pipeline-family.yaml

@@ -1,2 +1,2 @@
 name: general
-version: 0.0.59
+version: 0.0.60

test_general/api/test_app.py

@@ -439,6 +439,33 @@ def test_general_api_returns_503(monkeypatch):
     assert response.status_code == 503
 
 
+def test_general_api_returns_401(monkeypatch):
+    """
+    When UNSTRUCTURED_API_KEY is set, return a 401 if the unstructured-api-key header does not match
+    """
+    monkeypatch.setenv("UNSTRUCTURED_API_KEY", "foobar")
+
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "fake-xml.xml"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+        headers={"unstructured-api-key": "foobar"},
+    )
+
+    assert response.status_code == 200
+
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "fake-xml.xml"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+        headers={"unstructured-api-key": "helloworld"},
+    )
+
+    assert response.status_code == 401
+
+
 class MockResponse:
     def __init__(self, status_code):
         self.status_code = status_code