refactor: Replace deprecated lifespan event (#4187)

j1wonpark · web-flow · commit 8e4412589d24 · 2024-05-11T16:09:23.000Z
diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -3,6 +3,7 @@
 import threading
 import traceback
 import warnings
+from contextlib import asynccontextmanager
 from typing import List, Optional
 
 import pandas as pd
@@ -50,15 +51,16 @@ def get_app(
     registry_ttl_sec: int = DEFAULT_FEATURE_SERVER_REGISTRY_TTL,
 ):
     proto_json.patch()
-
-    app = FastAPI()
     # Asynchronously refresh registry, notifying shutdown and canceling the active timer if the app is shutting down
     registry_proto = None
     shutting_down = False
     active_timer: Optional[threading.Timer] = None
 
-    async def get_body(request: Request):
-        return await request.body()
+    def stop_refresh():
+        nonlocal shutting_down
+        shutting_down = True
+        if active_timer:
+            active_timer.cancel()
 
     def async_refresh():
         store.refresh_registry()
@@ -70,14 +72,16 @@ def async_refresh():
         active_timer = threading.Timer(registry_ttl_sec, async_refresh)
         active_timer.start()
 
-    @app.on_event("shutdown")
-    def shutdown_event():
-        nonlocal shutting_down
-        shutting_down = True
-        if active_timer:
-            active_timer.cancel()
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        async_refresh()
+        yield
+        stop_refresh()
 
-    async_refresh()
+    app = FastAPI(lifespan=lifespan)
+
+    async def get_body(request: Request):
+        return await request.body()
 
     @app.post("/get-online-features")
     def get_online_features(body=Depends(get_body)):
