From 47d3257c0f979bf50aeddaf9328ddae4f9f0242b Mon Sep 17 00:00:00 2001
From: Damian Shaw <damian.peter.shaw@gmail.com>
Date: Sun, 20 Oct 2024 15:26:40 -0400
Subject: [PATCH] Move simple API to /simple and add shutdown endpoint

---
 README.md                   |  2 +-
 src/pip_timemachine/main.py | 26 ++++++++++++++++++++++++--
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 594d961..5f2a26d 100644
--- a/README.md
+++ b/README.md
@@ -38,7 +38,7 @@ pip-timemachine 2023-10-18T12:00:00
 
 And then point pip (or any other installer) to use this server as an index:
 ```bash
-pip install requests --index http://127.0.0.1:8040
+pip install requests --index http://127.0.0.1:8040/simple
 ```
 
 ### CLI Options
diff --git a/src/pip_timemachine/main.py b/src/pip_timemachine/main.py
index ef0bc59..68749c2 100644
--- a/src/pip_timemachine/main.py
+++ b/src/pip_timemachine/main.py
@@ -1,5 +1,6 @@
 import datetime as dt
 import logging
+import threading
 from functools import cache
 
 import niquests
@@ -22,6 +23,7 @@
 app = FastAPI()
 
 # Global Vars
+UNICORN_SERVER: uvicorn.Server | None = None
 MOMENT: Instant = Instant.now()
 INDEX: str = "https://pypi.org/simple"
 REQUIRED_API_VERSION = (1, 1)
@@ -78,7 +80,23 @@ def filter_files_by_moment(files: list[dict]) -> tuple[list, dict]:
     return modified_files, modified_versions
 
 
-@app.get("/{package_name}")
+@app.get("/shutdown-pip-timemachine-server")
+async def shutdown_pip_timemachine_server():
+    """
+    Endpoint to gracefully shut down the FastAPI server.
+    """
+    def shutdown():
+        if UNICORN_SERVER:
+            UNICORN_SERVER.should_exit = True
+
+    # Run the shutdown in a separate thread to allow the response to return first
+    shutdown_thread = threading.Thread(target=shutdown)
+    shutdown_thread.start()
+
+    return JSONResponse({"message": "Server is shutting down."})
+
+
+@app.get("/simple/{package_name}")
 async def get_pypi_package_info(package_name: str, request: Request) -> Response:
     # Determine content type based on "Accept" header
     accept_header = request.headers.get("Accept")
@@ -125,10 +143,14 @@ async def get_pypi_package_info(package_name: str, request: Request) -> Response
 def run_server(moment: dt.datetime, index: str = INDEX, port: int = 8040):
     global MOMENT
     global INDEX
+    global UNICORN_SERVER
     MOMENT = Instant.from_py_datetime(moment.replace(tzinfo=dt.UTC))
     INDEX = index.rstrip("/")
 
-    uvicorn.run("pip_timemachine.main:app", port=port)
+    UNICORN_SERVER = uvicorn.Server(
+        uvicorn.Config("pip_timemachine.main:app", port=port)
+    )
+    UNICORN_SERVER.run()
 
 
 def main():