Repository: PrefectHQ/prefect

Issue:
Switch from deprecated server on_startup/shutdown hooks to lifespan
### First check

- [X] I am a contributor to the Prefect codebase

### Description

Starlette has deprecated the `on_startup` and `on_shutdown` hooks we are using for managing services. We should use the new interface.

See
- https://www.starlette.io/lifespan/
- https://github.com/tiangolo/fastapi/issues/617#issuecomment-1458402644
- https://github.com/encode/starlette/discussions/2067
- https://github.com/encode/starlette/pull/2065


### Impact

This will make us compatible with future versions of Starlette.

### Additional context

We should remove the ignore added in https://github.com/PrefectHQ/prefect/pull/8785


Discussion:

Resolving patch:
diff --git a/src/prefect/cli/cloud.py b/src/prefect/cli/cloud.py
--- a/src/prefect/cli/cloud.py
+++ b/src/prefect/cli/cloud.py
@@ -5,6 +5,7 @@
 import traceback
 import urllib.parse
 import webbrowser
+from contextlib import asynccontextmanager
 from typing import Hashable, Iterable, List, Optional, Tuple, Union
 
 import anyio
@@ -55,7 +56,16 @@ def set_login_api_ready_event():
     login_api.extra["ready-event"].set()
 
 
-login_api = FastAPI(on_startup=[set_login_api_ready_event])
+@asynccontextmanager
+async def lifespan(app):
+    try:
+        set_login_api_ready_event()
+        yield
+    finally:
+        pass
+
+
+login_api = FastAPI(lifespan=lifespan)
 """
 This small API server is used for data transmission for browser-based log in.
 """
diff --git a/src/prefect/server/api/server.py b/src/prefect/server/api/server.py
--- a/src/prefect/server/api/server.py
+++ b/src/prefect/server/api/server.py
@@ -5,6 +5,7 @@
 import asyncio
 import mimetypes
 import os
+from contextlib import asynccontextmanager
 from functools import partial, wraps
 from hashlib import sha256
 from typing import Awaitable, Callable, Dict, List, Mapping, Optional, Tuple
@@ -426,6 +427,16 @@ async def stop_services():
                 # `on_service_exit` should handle logging exceptions on exit
                 pass
 
+    @asynccontextmanager
+    async def lifespan(app):
+        try:
+            await run_migrations()
+            await add_block_types()
+            await start_services()
+            yield
+        finally:
+            await stop_services()
+
     def on_service_exit(service, task):
         """
         Added as a callback for completion of services to log exit
@@ -441,12 +452,7 @@ def on_service_exit(service, task):
     app = FastAPI(
         title=TITLE,
         version=API_VERSION,
-        on_startup=[
-            run_migrations,
-            add_block_types,
-            start_services,
-        ],
-        on_shutdown=[stop_services],
+        lifespan=lifespan,
     )
     api_app = create_orion_api(
         fast_api_app_kwargs={


Test patch:
diff --git a/tests/client/test_orion_client.py b/tests/client/test_orion_client.py
--- a/tests/client/test_orion_client.py
+++ b/tests/client/test_orion_client.py
@@ -2,6 +2,7 @@
 import os
 import random
 import threading
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta, timezone
 from typing import Generator, List
 from unittest.mock import ANY, MagicMock, Mock
@@ -224,6 +225,17 @@ def not_enough_open_files() -> bool:
     return soft_limit < 512 or hard_limit < 512
 
 
+def make_lifespan(startup, shutdown) -> callable:
+    async def lifespan(app):
+        try:
+            startup()
+            yield
+        finally:
+            shutdown()
+
+    return asynccontextmanager(lifespan)
+
+
 class TestClientContextManager:
     async def test_client_context_cannot_be_reentered(self):
         client = PrefectClient("http://foo.test")
@@ -243,7 +255,7 @@ async def test_client_context_cannot_be_reused(self):
 
     async def test_client_context_manages_app_lifespan(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         client = PrefectClient(app)
         startup.assert_not_called()
@@ -258,7 +270,7 @@ async def test_client_context_manages_app_lifespan(self):
 
     async def test_client_context_calls_app_lifespan_once_despite_nesting(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         startup.assert_not_called()
         shutdown.assert_not_called()
@@ -274,7 +286,7 @@ async def test_client_context_calls_app_lifespan_once_despite_nesting(self):
 
     async def test_client_context_manages_app_lifespan_on_sequential_usage(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         async with PrefectClient(app):
             pass
@@ -293,7 +305,7 @@ async def test_client_context_lifespan_is_robust_to_async_concurrency(self):
         startup = MagicMock(side_effect=lambda: print("Startup called!"))
         shutdown = MagicMock(side_effect=lambda: print("Shutdown called!!"))
 
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         one_started = anyio.Event()
         one_exited = anyio.Event()
@@ -338,7 +350,7 @@ async def two():
     @pytest.mark.skipif(not_enough_open_files(), reason=not_enough_open_files.__doc__)
     async def test_client_context_lifespan_is_robust_to_threaded_concurrency(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         async def enter_client(context):
             # We must re-enter the profile context in the new thread
@@ -366,7 +378,7 @@ async def enter_client(context):
 
     async def test_client_context_lifespan_is_robust_to_high_async_concurrency(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         async def enter_client():
             # Use random sleeps to interleave clients
@@ -386,7 +398,7 @@ async def enter_client():
     @pytest.mark.skipif(not_enough_open_files(), reason=not_enough_open_files.__doc__)
     async def test_client_context_lifespan_is_robust_to_mixed_concurrency(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         async def enter_client():
             # Use random sleeps to interleave clients
@@ -437,7 +449,7 @@ async def test_client_context_lifespan_is_robust_to_dependency_deadlocks(self):
         startup = MagicMock(side_effect=lambda: print("Startup called!"))
         shutdown = MagicMock(side_effect=lambda: print("Shutdown called!!"))
 
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         one_started = anyio.Event()
         one_exited = anyio.Event()
@@ -478,7 +490,7 @@ async def two():
 
     async def test_client_context_manages_app_lifespan_on_exception(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         client = PrefectClient(app)
 
@@ -491,7 +503,7 @@ async def test_client_context_manages_app_lifespan_on_exception(self):
 
     async def test_client_context_manages_app_lifespan_on_anyio_cancellation(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         async def enter_client(task_status):
             async with PrefectClient(app):
@@ -510,7 +522,7 @@ async def enter_client(task_status):
 
     async def test_client_context_manages_app_lifespan_on_exception_when_nested(self):
         startup, shutdown = MagicMock(), MagicMock()
-        app = FastAPI(on_startup=[startup], on_shutdown=[shutdown])
+        app = FastAPI(lifespan=make_lifespan(startup, shutdown))
 
         with pytest.raises(ValueError):
             async with PrefectClient(app):