lifespan event to create databases

kscalelabs · Jun 17, 2024 · ac585a6 · ac585a6
1 parent 4ab6553
commit ac585a6
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 37 deletions.
diff --git a/store/app/db.py b/store/app/db.py
@@ -37,38 +37,40 @@ async def create_tables(crud: Crud | None = None, deletion_protection: bool = Fa
             await create_tables(crud)
 
     else:
-        await crud._create_dynamodb_table(
-            name="Users",
-            keys=[
-                ("user_id", "S", "HASH"),
-            ],
-            gsis=[
-                ("emailIndex", "email", "S", "HASH"),
-                ("usernameIndex", "username", "S", "HASH"),
-            ],
-            deletion_protection=deletion_protection,
-        )
-        await crud._create_dynamodb_table(
-            name="Robots",
-            keys=[
-                ("robot_id", "S", "HASH"),
-            ],
-            gsis=[
-                ("ownerIndex", "owner", "S", "HASH"),
-                ("nameIndex", "name", "S", "HASH"),
-            ],
-            deletion_protection=deletion_protection,
-        )
-        await crud._create_dynamodb_table(
-            name="Parts",
-            keys=[
-                ("part_id", "S", "HASH"),
-            ],
-            gsis=[
-                ("ownerIndex", "owner", "S", "HASH"),
-                ("nameIndex", "name", "S", "HASH"),
-            ],
-            deletion_protection=deletion_protection,
+        await asyncio.gather(
+            crud._create_dynamodb_table(
+                name="Users",
+                keys=[
+                    ("user_id", "S", "HASH"),
+                ],
+                gsis=[
+                    ("emailIndex", "email", "S", "HASH"),
+                    ("usernameIndex", "username", "S", "HASH"),
+                ],
+                deletion_protection=deletion_protection,
+            ),
+            crud._create_dynamodb_table(
+                name="Robots",
+                keys=[
+                    ("robot_id", "S", "HASH"),
+                ],
+                gsis=[
+                    ("ownerIndex", "owner", "S", "HASH"),
+                    ("nameIndex", "name", "S", "HASH"),
+                ],
+                deletion_protection=deletion_protection,
+            ),
+            crud._create_dynamodb_table(
+                name="Parts",
+                keys=[
+                    ("part_id", "S", "HASH"),
+                ],
+                gsis=[
+                    ("ownerIndex", "owner", "S", "HASH"),
+                    ("nameIndex", "name", "S", "HASH"),
+                ],
+                deletion_protection=deletion_protection,
+            ),
         )
 
 
@@ -85,10 +87,11 @@ async def delete_tables(crud: Crud | None = None) -> None:
             await delete_tables(crud)
 
     else:
-        await crud._delete_dynamodb_table("Users")
-        await crud._delete_dynamodb_table("UserEmails")
-        await crud._delete_dynamodb_table("Robots")
-        await crud._delete_dynamodb_table("Parts")
+        await asyncio.gather(
+            crud._delete_dynamodb_table("Users"),
+            crud._delete_dynamodb_table("Robots"),
+            crud._delete_dynamodb_table("Parts"),
+        )
 
 
 async def populate_with_dummy_data(crud: Crud | None = None) -> None:

diff --git a/store/app/main.py b/store/app/main.py
@@ -1,16 +1,31 @@
 """Defines the main entrypoint for the FastAPI app."""
 
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
 from fastapi import FastAPI, Request, status
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
+from store.app.db import create_tables
 from store.app.routers.image import image_router
 from store.app.routers.part import parts_router
 from store.app.routers.robot import robots_router
 from store.app.routers.users import users_router
 from store.settings import settings
 
-app = FastAPI()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+    """Initializes the app and creates the database tables."""
+    await create_tables()
+    try:
+        yield
+    finally:
+        pass
+
+
+app = FastAPI(lifespan=lifespan)
 
 # Adds CORS middleware.
 app.add_middleware(