init

database.py

@@ -0,0 +1,138 @@
+from sqlmodel import SQLModel, create_engine, Session, select
+from sqlalchemy import text
+from rich.progress import Progress
+
+from faker import Faker
+
+fake = Faker()
+import random
+
+from models import User, Post
+
+# Setup SQLite DB
+sqlite_url = "sqlite:///data.db"
+engine = create_engine(sqlite_url, echo=True)
+
+
+def create_db_and_tables():
+    SQLModel.metadata.create_all(engine)
+
+
+def setup_db():
+    with engine.connect() as conn:
+        conn.execute(text("PRAGMA foreign_keys = ON;"))
+        conn.execute(text("PRAGMA journal_mode = WAL;"))
+        conn.execute(text("PRAGMA synchronous = NORMAL;"))
+        conn.execute(text("PRAGMA temp_store = MEMORY;"))
+        conn.execute(text("PRAGMA cache_size = 50000;"))
+
+
+post_templates = [
+    "Just published a new blog about {topic}. Check it out!",
+    "Today's thoughts: {quote}",
+    "Here's a photo from my trip to {location} 🌍",
+    "Shoutout to all the amazing devs working with {topic}!",
+    "Feeling inspired by {quote}",
+    "If you're into {topic}, we should connect!",
+    "Throwback to {location} last summer ☀️",
+]
+topics = ["AI", "SQLModel", "Linux", "Neovim", "Docker", "Kubernetes", "FastAPI"]
+quotes = [
+    "Keep it simple.",
+    "Fail fast, learn faster.",
+    "In code we trust.",
+    "Life is short, automate it.",
+    "Stay curious.",
+]
+locations = ["Tokyo", "Paris", "New York", "Berlin", "São Paulo", "Cairo"]
+
+
+def add_posts_to_user(user_id, num_posts=5):
+    with Session(engine) as session:
+        for _ in range(num_posts):
+            post = Post(
+                content=random.choice(post_templates).format(
+                    topic=random.choice(topics),
+                    quote=random.choice(quotes),
+                    location=random.choice(locations),
+                ),
+                timestamp=fake.date_time_this_year(),
+                user_id=user_id,
+            )
+            session.add(post)
+        session.commit()
+
+
+def populate_data(
+    num_users: int = 100, posts_per_user_range=(5, 20), batch_size=10_000
+):
+    total_posts = 0
+    with Progress(transient=True) as progress:
+        user_task = progress.add_task("[bold blue]Creating users...", total=num_users)
+        users = []
+
+        for _ in range(num_users):
+            user = User(
+                username=fake.user_name(),
+                display_name=fake.name(),
+            )
+            users.append(user)
+            progress.advance(user_task)
+
+        # Bulk insert users
+        with Session(engine) as session:
+            for i in range(0, len(users), batch_size):
+                session.bulk_save_objects(users[i : i + batch_size])
+                session.commit()
+
+        # Fetch user ids after bulk insert
+        with Session(engine) as session:
+            user_ids = session.exec(select(User.id)).all()
+
+        # Generate posts
+        post_task = progress.add_task("[green]Creating posts...", total=num_users)
+        posts = []
+
+        for user_id in user_ids:
+            for _ in range(random.randint(*posts_per_user_range)):
+                content = random.choice(post_templates).format(
+                    topic=random.choice(topics),
+                    quote=random.choice(quotes),
+                    location=random.choice(locations),
+                )
+                post = Post(
+                    content=content,
+                    timestamp=fake.date_time_this_year(),
+                    user_id=user_id,
+                )
+                posts.append(post)
+            progress.advance(post_task)
+
+        # Bulk insert posts
+        with Session(engine) as session:
+            for i in range(0, len(posts), batch_size):
+                session.bulk_save_objects(posts[i : i + batch_size])
+                session.commit()
+
+    print(f"✅ Done! Inserted {len(users)} users and {len(posts)} posts.")
+
+
+def query_example():
+    with Session(engine) as session:
+        statement = select(User)
+        users = session.exec(statement).all()
+
+        for user in users:
+            print(f"User: {user.display_name} ({user.username})")
+            for post in user.posts:
+                print(f"  - {post.content} [{post.timestamp.isoformat()}]")
+
+
+if __name__ == "__main__":
+    import sys
+
+    n_users = sys.argv[1] if len(sys.argv) > 1 else 100
+    n_posts = sys.argv[2] if len(sys.argv) > 2 else 20
+    create_db_and_tables()
+    setup_db()
+    populate_data(int(n_users), (int(5), int(n_posts)))