Update database.py

utils/database.py

@@ -1,104 +1,90 @@
-# api/utils/database.py
-import sqlite3
 import json
 from typing import List, Dict
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, Column, Integer, String, Float, Boolean, Text
+from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, Session
 
-DATABASE = "instagram_ai.db"
+DATABASE_URL = "postgresql://postgres:Lovyelias5584.@db.lgbnxplydqdymepehirg.supabase.co:5432/postgres"
+
+Base = declarative_base()
+
+# Define the posts table using SQLAlchemy ORM
+class Post(Base):
+    __tablename__ = "posts"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    username = Column(String, nullable=False)
+    caption = Column(Text, nullable=True)
+    hashtags = Column(Text, nullable=True)  # Store as JSON string
+    likes = Column(Integer, default=0)
+    comments = Column(Integer, default=0)
+    date = Column(String, nullable=True)
+    image_url = Column(String, unique=True, nullable=False)
+    engagement_rate = Column(Float, default=0.0)
+    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 SQLite database.
+    Initialize the PostgreSQL database by creating tables.
     """
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    cursor.execute('''
-        CREATE TABLE IF NOT EXISTS posts (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            username TEXT NOT NULL,
-            caption TEXT,
-            hashtags TEXT,
-            likes INTEGER,
-            comments INTEGER,
-            date TEXT,
-            image_url TEXT UNIQUE, 
-            engagement_rate REAL,
-            viral_score REAL,
-            promote BOOLEAN
-        )
-    ''')
-    conn.commit()
-    conn.close()
+    Base.metadata.create_all(bind=engine)
+    print("Database initialized.")
 
-def post_exists(image_url: str) -> bool:
+def post_exists(session: Session, image_url: str) -> bool:
     """
     Check if a post with the given image_url already exists in the database.
     """
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    cursor.execute('SELECT id FROM posts WHERE image_url = ?', (image_url,))
-    result = cursor.fetchone()
-    conn.close()
-    return result is not None
+    return session.query(Post.id).filter(Post.image_url == image_url).first() is not None
 
 def save_to_db(data: List[Dict]):
     """
-    Save data to the SQLite database, avoiding duplicates.
+    Save data to the PostgreSQL database, avoiding duplicates.
     """
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    for post in data:
-        if not post_exists(post.get('image_url')):
-            cursor.execute('''
-                INSERT INTO posts (username, caption, hashtags, likes, comments, date, image_url, engagement_rate, viral_score, promote)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-            ''', (
-                post.get('username', ''),
-                post.get('caption', ''),
-                json.dumps(post.get('hashtags', [])), 
-                post.get('likes', 0),
-                post.get('comments', 0),
-                post.get('date', ''), 
-                post.get('image_url', ''),
-                post.get('engagement_rate', 0.0),
-                post.get('viral_score', 0.0),
-                post.get('promote', False)
-            ))
-    conn.commit()
-    conn.close()
-    print(f"Data saved to database: {DATABASE}")
+    with SessionLocal() as session:
+        for post in data:
+            if not post_exists(session, post.get("image_url")):
+                new_post = Post(
+                    username=post.get("username", ""),
+                    caption=post.get("caption", ""),
+                    hashtags=json.dumps(post.get("hashtags", [])),  # Convert list to JSON string
+                    likes=post.get("likes", 0),
+                    comments=post.get("comments", 0),
+                    date=post.get("date", ""),
+                    image_url=post.get("image_url", ""),
+                    engagement_rate=post.get("engagement_rate", 0.0),
+                    viral_score=post.get("viral_score", 0.0),
+                    promote=post.get("promote", False),
+                )
+                session.add(new_post)
+        session.commit()
+    print("Data saved to database.")
 
 def fetch_posts_from_db(username: str) -> List[Dict]:
     """
     Fetch posts from the database for a given username.
     """
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    cursor.execute('SELECT * FROM posts WHERE username = ?', (username,))
-    rows = cursor.fetchall()
-    conn.close()
-
-    
-    posts = []
-    for row in rows:
-        posts.append({
-            "username": row[1],
-            "caption": row[2],
-            "hashtags": json.loads(row[3]), 
-            "likes": row[4],
-            "comments": row[5],
-            "date": row[6],  
-            "image_url": row[7],
-            "engagement_rate": row[8],
-            "viral_score": row[9],
-            "promote": bool(row[10])
-        })
-    return posts
-
-SQLALCHEMY_DATABASE_URL = f"sqlite:///{DATABASE}"
-engine = create_engine(SQLALCHEMY_DATABASE_URL)
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+    with SessionLocal() as session:
+        posts = session.query(Post).filter(Post.username == username).all()
+        return [
+            {
+                "username": post.username,
+                "caption": post.caption,
+                "hashtags": json.loads(post.hashtags),  # Convert JSON string back to list
+                "likes": post.likes,
+                "comments": post.comments,
+                "date": post.date,
+                "image_url": post.image_url,
+                "engagement_rate": post.engagement_rate,
+                "viral_score": post.viral_score,
+                "promote": post.promote,
+            }
+            for post in posts
+        ]
 
 def get_db():
     """
@@ -108,4 +94,4 @@ def get_db():
     try:
         yield db
     finally:
-        db.close()
+        db.close()