Refactored lifespan and update imports

Signed-off-by: samadpls <[email protected]>

app/main.py

@@ -1,8 +1,8 @@
 """Main app."""
 
-from contextlib import asynccontextmanager
 import os
 import warnings
+from contextlib import asynccontextmanager
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
@@ -15,8 +15,58 @@
 from .api import utility as util
 from .api.routers import attributes, query
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """
+    Load and set up resources before the application starts receiving requests.
+    Clean up resources after the application finishes handling requests.
+    """
+    # Check if username and password environment variables are set
+    if (
+        # TODO: Check if this error is still raised when variables are empty strings
+        os.environ.get(util.GRAPH_USERNAME.name) is None
+        or os.environ.get(util.GRAPH_PASSWORD.name) is None
+    ):
+        raise RuntimeError(
+            f"The application was launched but could not find the {util.GRAPH_USERNAME.name} and / or {util.GRAPH_PASSWORD.name} environment variables."
+        )
+
+    # Raises warning if allowed origins environment variable has not been set or is an empty string.
+    if os.environ.get(util.ALLOWED_ORIGINS.name, "") == "":
+        warnings.warn(
+            f"The API was launched without providing any values for the {util.ALLOWED_ORIGINS.name} environment variable. "
+            "This means that the API will only be accessible from the same origin it is hosted from: https://developer.mozilla.org/en-US/docs/Web/Security/Same-origin_policy. "
+            f"If you want to access the API from tools hosted at other origins such as the Neurobagel query tool, explicitly set the value of {util.ALLOWED_ORIGINS.name} to the origin(s) of these tools (e.g. http://localhost:3000). "
+            "Multiple allowed origins should be separated with spaces in a single string enclosed in quotes. "
+        )
+
+    # We use Starlette's ability (FastAPI is Starlette underneath) to store arbitrary state on the app instance (https://www.starlette.io/applications/#storing-state-on-the-app-instance)
+    # to store a temporary directory object and its corresponding path. These data are local to the instance and will be recreated on every app launch (i.e. not persisted).
+    app.state.vocab_dir = TemporaryDirectory()
+    app.state.vocab_dir_path = Path(app.state.vocab_dir.name)
+
+    # TODO: Maybe store these paths in one dictionary on the app instance instead of separate variables?
+    app.state.cogatlas_term_lookup_path = (
+        app.state.vocab_dir_path / "cogatlas_task_term_labels.json"
+    )
+    app.state.snomed_term_lookup_path = (
+        app.state.vocab_dir_path / "snomedct_disorder_term_labels.json"
+    )
+
+    util.fetch_and_save_cogatlas(app.state.cogatlas_term_lookup_path)
+    util.create_snomed_term_lookup(app.state.snomed_term_lookup_path)
+
+    yield
+    # Clean up the temporary directory created on startup
+    app.state.vocab_dir.cleanup()
+
+
 app = FastAPI(
-    default_response_class=ORJSONResponse, docs_url=None, redoc_url=None
+    default_response_class=ORJSONResponse,
+    docs_url=None,
+    redoc_url=None,
+    lifespan=lifespan,
 )
 favicon_url = "https://raw.githubusercontent.com/neurobagel/documentation/main/docs/imgs/logo/neurobagel_favicon.png"
 
@@ -60,54 +110,6 @@ def overridden_redoc():
         redoc_favicon_url=favicon_url,
     )
 
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    """
-    Load and set up resources before the application starts receiving requests.
-    Clean up resources after the application finishes handling requests.
-    """
-    # to store a temporary directory object and its corresponding path. These data are local to the instance and will be recreated on every app launch (i.e. not persisted).
-    vocab_dir = TemporaryDirectory()  
-    try:
-        # Check if username and password environment variables are set
-        if (
-            # TODO: Check if this error is still raised when variables are empty strings
-            os.environ.get(util.GRAPH_USERNAME.name) is None
-            or os.environ.get(util.GRAPH_PASSWORD.name) is None
-        ):
-            raise RuntimeError(
-                f"The application was launched but could not find the {util.GRAPH_USERNAME.name} and / or {util.GRAPH_PASSWORD.name} environment variables."
-            )
-
-        # Raises warning if allowed origins environment variable has not been set or is an empty string.
-        if os.environ.get(util.ALLOWED_ORIGINS.name, "") == "":
-            warnings.warn(
-                f"The API was launched without providing any values for the {util.ALLOWED_ORIGINS.name} environment variable. "
-                "This means that the API will only be accessible from the same origin it is hosted from: https://developer.mozilla.org/en-US/docs/Web/Security/Same-origin_policy. "
-                f"If you want to access the API from tools hosted at other origins such as the Neurobagel query tool, explicitly set the value of {util.ALLOWED_ORIGINS.name} to the origin(s) of these tools (e.g. http://localhost:3000). "
-                "Multiple allowed origins should be separated with spaces in a single string enclosed in quotes. "
-            )
-
-        # Fetch and store vocabularies to a temporary directory
-        # We use Starlette's ability (FastAPI is Starlette underneath) to store arbitrary state on the app instance (https://www.starlette.io/applications/#storing-state-on-the-app-instance)
-        vocab_dir_path = Path(vocab_dir.name)
-
-        # TODO: Maybe store these paths in one dictionary on the app instance instead of separate variables?
-        app.cogatlas_term_lookup_path = (
-            vocab_dir_path / "cogatlas_task_term_labels.json"
-        )
-        app.snomed_term_lookup_path = (
-            vocab_dir_path / "snomedct_disorder_term_labels.json"
-        )
-
-        util.fetch_and_save_cogatlas(app.cogatlas_term_lookup_path)
-        util.create_snomed_term_lookup(app.snomed_term_lookup_path)
-
-        yield
-    finally:
-        # Clean up the temporary directory created on startup
-        vocab_dir.cleanup()
-
 
 app.include_router(query.router)
 app.include_router(attributes.router)