Spaces:
Running
Running
Commit
·
9b31d36
1
Parent(s):
6fae15f
Implement asynchronous database initialization: add init_db function to create collections and indexes in MongoDB, update get_db to support async, and modify main.py to await database initialization during startup.
Browse files- app/db.py +45 -2
- app/main.py +1 -1
app/db.py
CHANGED
@@ -6,6 +6,7 @@ from motor.motor_asyncio import AsyncIOMotorClient
|
|
6 |
from pymongo.errors import ConnectionFailure
|
7 |
from dotenv import load_dotenv
|
8 |
import certifi
|
|
|
9 |
|
10 |
# Load environment variables
|
11 |
load_dotenv()
|
@@ -20,6 +21,47 @@ DB_NAME = os.getenv("DB_NAME", "tts_api")
|
|
20 |
# MongoDB client instance
|
21 |
client: Optional[AsyncIOMotorClient] = None
|
22 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
23 |
async def connect_to_mongo():
|
24 |
"""Connect to MongoDB."""
|
25 |
global client
|
@@ -46,10 +88,11 @@ async def close_mongo_connection():
|
|
46 |
client.close()
|
47 |
logger.info("MongoDB connection closed")
|
48 |
|
49 |
-
def get_db():
|
50 |
"""Get database instance."""
|
|
|
51 |
if not client:
|
52 |
-
|
53 |
return client[DB_NAME]
|
54 |
|
55 |
# Collection names
|
|
|
6 |
from pymongo.errors import ConnectionFailure
|
7 |
from dotenv import load_dotenv
|
8 |
import certifi
|
9 |
+
from datetime import datetime
|
10 |
|
11 |
# Load environment variables
|
12 |
load_dotenv()
|
|
|
21 |
# MongoDB client instance
|
22 |
client: Optional[AsyncIOMotorClient] = None
|
23 |
|
24 |
+
async def init_db():
|
25 |
+
"""Initialize the database with required collections and indexes."""
|
26 |
+
try:
|
27 |
+
# Get database instance
|
28 |
+
db = await get_db()
|
29 |
+
|
30 |
+
# Create collections if they don't exist
|
31 |
+
collections = await db.list_collection_names()
|
32 |
+
|
33 |
+
# Audiobooks collection
|
34 |
+
if AUDIOBOOKS_COLLECTION not in collections:
|
35 |
+
logger.info(f"Creating collection: {AUDIOBOOKS_COLLECTION}")
|
36 |
+
await db.create_collection(AUDIOBOOKS_COLLECTION)
|
37 |
+
# Create indexes
|
38 |
+
await db[AUDIOBOOKS_COLLECTION].create_index("id", unique=True)
|
39 |
+
await db[AUDIOBOOKS_COLLECTION].create_index("created_at")
|
40 |
+
await db[AUDIOBOOKS_COLLECTION].create_index("status")
|
41 |
+
|
42 |
+
# Voices collection
|
43 |
+
if VOICES_COLLECTION not in collections:
|
44 |
+
logger.info(f"Creating collection: {VOICES_COLLECTION}")
|
45 |
+
await db.create_collection(VOICES_COLLECTION)
|
46 |
+
# Create indexes
|
47 |
+
await db[VOICES_COLLECTION].create_index("id", unique=True)
|
48 |
+
await db[VOICES_COLLECTION].create_index("name")
|
49 |
+
await db[VOICES_COLLECTION].create_index("type")
|
50 |
+
|
51 |
+
# Audio cache collection
|
52 |
+
if AUDIO_CACHE_COLLECTION not in collections:
|
53 |
+
logger.info(f"Creating collection: {AUDIO_CACHE_COLLECTION}")
|
54 |
+
await db.create_collection(AUDIO_CACHE_COLLECTION)
|
55 |
+
# Create indexes
|
56 |
+
await db[AUDIO_CACHE_COLLECTION].create_index("hash", unique=True)
|
57 |
+
await db[AUDIO_CACHE_COLLECTION].create_index("created_at")
|
58 |
+
|
59 |
+
logger.info("Database initialization completed successfully")
|
60 |
+
|
61 |
+
except Exception as e:
|
62 |
+
logger.error(f"Error initializing database: {str(e)}")
|
63 |
+
raise
|
64 |
+
|
65 |
async def connect_to_mongo():
|
66 |
"""Connect to MongoDB."""
|
67 |
global client
|
|
|
88 |
client.close()
|
89 |
logger.info("MongoDB connection closed")
|
90 |
|
91 |
+
async def get_db():
|
92 |
"""Get database instance."""
|
93 |
+
global client
|
94 |
if not client:
|
95 |
+
await connect_to_mongo()
|
96 |
return client[DB_NAME]
|
97 |
|
98 |
# Collection names
|
app/main.py
CHANGED
@@ -57,7 +57,7 @@ async def lifespan(app: FastAPI):
|
|
57 |
# Initialize database
|
58 |
from app.db import init_db
|
59 |
logger.info("Initializing database...")
|
60 |
-
init_db()
|
61 |
logger.info("Database initialized")
|
62 |
|
63 |
# Create necessary directories - use persistent locations
|
|
|
57 |
# Initialize database
|
58 |
from app.db import init_db
|
59 |
logger.info("Initializing database...")
|
60 |
+
await init_db()
|
61 |
logger.info("Database initialized")
|
62 |
|
63 |
# Create necessary directories - use persistent locations
|