jameszokah commited on
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
Files changed (2) hide show
  1. app/db.py +45 -2
  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
- raise ConnectionError("MongoDB client not initialized")
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