Cleanup storage integration and write pending batches on shutdown (#76)

* Promote close() method to the _Storage class definition

* Switch to asynccontextmanager method for handling startup/shutdown events in FastAPI
Add storage.close() to the shutdown phase of the lifespan method
Use fully-qualified names for create_storage() and create_translator() methods so they can be mocked for unit tests

* Add unit test for startup/shutdown handling method

src/anyvar/restapi/main.py

@@ -1,6 +1,7 @@
 """Provide core route definitions for REST service."""
 import logging
 import tempfile
+from contextlib import asynccontextmanager
 from http import HTTPStatus
 
 import ga4gh.vrs
@@ -9,7 +10,7 @@
 from pydantic import StrictStr
 
 import anyvar
-from anyvar.anyvar import AnyVar, create_storage, create_translator
+from anyvar.anyvar import AnyVar
 from anyvar.extras.vcf import VcfRegistrar
 from anyvar.restapi.schema import (
     AnyVarStatsResponse,
@@ -29,25 +30,36 @@
 _logger = logging.getLogger(__name__)
 
 
+@asynccontextmanager
+async def app_lifespan(param_app: FastAPI):
+    """Initialize AnyVar instance and associate with FastAPI app on startup
+    and teardown the AnyVar instance on shutdown"""
+
+    # create anyvar instance
+    storage = anyvar.anyvar.create_storage()
+    translator = anyvar.anyvar.create_translator()
+    anyvar_instance = AnyVar(object_store=storage, translator=translator)
+
+    # associate anyvar with the app state
+    param_app.state.anyvar = anyvar_instance
+
+    yield
+
+    # close storage connector on shutdown
+    storage.close()
+
+
 app = FastAPI(
     title="AnyVar",
     version=anyvar.__version__,
     docs_url="/",
     openapi_url="/openapi.json",
     swagger_ui_parameters={"tryItOutEnabled": True},
     description="Register and retrieve VRS value objects.",
+    lifespan=app_lifespan,
 )
 
 
-@app.on_event("startup")
-async def startup():
-    """Initialize AnyVar instance and associate with FastAPI app"""
-    storage = create_storage()
-    translator = create_translator()
-    anyvar_instance = AnyVar(object_store=storage, translator=translator)
-    app.state.anyvar = anyvar_instance
-
-
 @app.get(
     "/info",
     response_model=InfoResponse,

src/anyvar/storage/__init__.py

@@ -35,6 +35,10 @@ def get_variation_count(self, variation_type: VariationStatisticType) -> int:
     def wipe_db(self):
         """Empty database of all stored records."""
 
+    @abstractmethod
+    def close(self):
+        """Closes the storage integration and cleans up any resources"""
+
 
 class _BatchManager(AbstractContextManager):
     """Base context management class for batch writing.

tests/conftest.py

@@ -12,7 +12,7 @@
 
 def pytest_collection_modifyitems(items):
     """Modify test items in place to ensure test modules run in a given order."""
-    MODULE_ORDER = ["test_variation", "test_general", "test_location", "test_search", "test_vcf", "test_storage_mapping", "test_snowflake"]
+    MODULE_ORDER = ["test_lifespan", "test_variation", "test_general", "test_location", "test_search", "test_vcf", "test_storage_mapping", "test_snowflake"]
     # remember to add new test modules to the order constant:
     assert len(MODULE_ORDER) == len(list(Path(__file__).parent.rglob("test_*.py")))
     items.sort(key=lambda i: MODULE_ORDER.index(i.module.__name__))

tests/test_lifespan.py

@@ -0,0 +1,26 @@
+from fastapi.testclient import TestClient
+from fastapi import FastAPI
+from anyvar.restapi.main import app_lifespan
+from anyvar.storage import _Storage
+
+
+def test_lifespan(mocker):
+    """Test the app_lifespan method in anyvar.restapi.main"""
+    create_storage_mock = mocker.patch("anyvar.anyvar.create_storage")
+    storage_mock = mocker.Mock(spec=_Storage)
+    create_storage_mock.return_value = storage_mock
+    create_translator_mock = mocker.patch("anyvar.anyvar.create_translator")
+    create_translator_mock.return_value = {}
+    app = FastAPI(
+        title="AnyVarTest",
+        docs_url="/",
+        openapi_url="/openapi.json",
+        description="Test app",
+        lifespan=app_lifespan,
+    )
+    with TestClient(app) as client:
+        create_storage_mock.assert_called_once()
+        create_translator_mock.assert_called_once()
+        assert app.state.anyvar is not None
+
+    storage_mock.close.assert_called_once()