Update utils/database.py

utils/database.py

@@ -1,10 +1,22 @@
 import json
 from typing import List, Dict
-from sqlalchemy import create_engine, Column, Integer, String, Float, Boolean, Text
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import sessionmaker, declarative_base
+from sqlalchemy import Column, Integer, String, Float, Boolean, Text, select
+import asyncio
 
-DATABASE_URL = "postgresql://postgres.lgbnxplydqdymepehirg:[email protected]:5432/postgres"
+# Use the async database URL (replace "postgresql" with "postgresql+asyncpg")
+DATABASE_URL = "postgresql+asyncpg://postgres.lgbnxplydqdymepehirg:[email protected]:5432/postgres"
+
+# Initialize the async engine
+engine = create_async_engine(DATABASE_URL, echo=True)
+
+# Create an async session maker
+AsyncSessionLocal = sessionmaker(
+    bind=engine,
+    class_=AsyncSession,
+    expire_on_commit=False,
+)
 
 Base = declarative_base()
 
@@ -24,15 +36,13 @@ class Post(Base):
     viral_score = Column(Float, default=0.0)
     promote = Column(Boolean, default=False)
 
-# Initialize SQLAlchemy engine and session maker
-engine = create_engine(DATABASE_URL)
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
-
-def init_db():
+# Initialize the database (create tables)
+async def init_db():
     """
     Initialize the PostgreSQL database by creating tables.
     """
-    Base.metadata.create_all(bind=engine)
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
     print("Database initialized.")
 
 async def post_exists(session: AsyncSession, image_url: str) -> bool:
@@ -65,12 +75,13 @@ async def save_to_db(data: List[Dict]):
         await session.commit()
     print("Data saved to database.")
 
-def fetch_posts_from_db(username: str) -> List[Dict]:
+async def fetch_posts_from_db(username: str) -> List[Dict]:
     """
     Fetch posts from the database for a given username.
     """
-    with SessionLocal() as session:
-        posts = session.query(Post).filter(Post.username == username).all()
+    async with AsyncSessionLocal() as session:
+        result = await session.execute(select(Post).filter(Post.username == username))
+        posts = result.scalars().all()
         return [
             {
                 "username": post.username,
@@ -87,12 +98,37 @@ def fetch_posts_from_db(username: str) -> List[Dict]:
             for post in posts
         ]
 
-def get_db():
+async def get_db():
     """
     Dependency to get a database session.
     """
-    db = SessionLocal()
-    try:
-        yield db
-    finally:
-        db.close()
+    async with AsyncSessionLocal() as session:
+        yield session
+
+# Example usage
+async def main():
+    # Initialize the database
+    await init_db()
+
+    # Example data to save
+    example_data = [
+        {
+            "username": "test_user",
+            "caption": "This is a test post",
+            "hashtags": ["test", "example"],
+            "likes": 10,
+            "comments": 2,
+            "date": "2025-01-27",
+            "image_url": "https://example.com/image1.jpg",
+            "engagement_rate": 0.5,
+            "viral_score": 0.8,
+            "promote": False,
+        }
+    ]
+
+    # Save data to the database
+    await save_to_db(example_data)
+
+    # Fetch posts from the database
+    posts = await fetch_posts_from_db("test_user")
+    print("Fetched posts:", posts)