Update 1.3.0-beta

api/app/const/plans.py

@@ -2,9 +2,11 @@
     "free": {
         "web_search": 0,
         "link_extraction": 0,
+        "storage": 50 * 1024 * 1024,  # 50 MB
     },
     "premium": {
         "web_search": 200,
         "link_extraction": 1000,
+        "storage": 5 * 1024 * 1024 * 1024,  # 5 GB
     },
 }

api/app/database/pg/auth_ops/verification_crud.py

@@ -0,0 +1,66 @@
+import datetime
+import uuid
+from datetime import timezone
+
+from database.pg.models import User, VerificationToken
+from sqlalchemy import delete, select, update
+from sqlalchemy.ext.asyncio import AsyncEngine as SQLAlchemyAsyncEngine
+from sqlmodel import and_
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+
+async def create_verification_token(
+    pg_engine: SQLAlchemyAsyncEngine, user_id: uuid.UUID, email: str, code: str
+) -> VerificationToken:
+    """
+    Creates a new verification token for a user.
+    """
+    expires_at = datetime.datetime.now(timezone.utc) + datetime.timedelta(minutes=15)
+
+    async with AsyncSession(pg_engine) as session:
+        await delete_verification_tokens_for_user(pg_engine, user_id)
+
+        token = VerificationToken(user_id=user_id, email=email, code=code, expires_at=expires_at)
+        session.add(token)
+        await session.commit()
+        await session.refresh(token)
+        return token
+
+
+async def get_verification_token(
+    pg_engine: SQLAlchemyAsyncEngine, email: str, code: str
+) -> VerificationToken | None:
+    """
+    Retrieves a verification token by email and code.
+    """
+    async with AsyncSession(pg_engine) as session:
+        stmt = select(VerificationToken).where(
+            and_(VerificationToken.email == email, VerificationToken.code == code)
+        )
+        result = await session.exec(stmt)  # type: ignore
+        return result.scalars().one_or_none()  # type: ignore
+
+
+async def delete_verification_tokens_for_user(
+    pg_engine: SQLAlchemyAsyncEngine, user_id: uuid.UUID
+) -> None:
+    """
+    Deletes all verification tokens for a specific user.
+    """
+    async with AsyncSession(pg_engine) as session:
+        stmt = delete(VerificationToken).where(and_(VerificationToken.user_id == user_id))
+        await session.exec(stmt)
+        await session.commit()
+
+
+async def mark_user_as_verified(pg_engine: SQLAlchemyAsyncEngine, user_id: uuid.UUID) -> User:
+    """
+    Updates the user's status to verified.
+    """
+    async with AsyncSession(pg_engine) as session:
+        stmt = update(User).where(and_(User.id == user_id)).values(is_verified=True).returning(User)
+        result = await session.exec(stmt)
+        user = result.scalar_one()
+        await session.commit()
+        await session.refresh(user)
+        return user  # type: ignore

api/app/database/pg/models.py

@@ -21,6 +21,37 @@ class QueryTypeEnum(str, Enum):
     LINK_EXTRACTION = "link_extraction"
 
 
+class UserStorageUsage(SQLModel, table=True):
+    __tablename__ = "user_storage_usage"
+
+    id: Optional[uuid.UUID] = Field(
+        default=None,
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            primary_key=True,
+            server_default=func.uuid_generate_v4(),
+        ),
+    )
+    user_id: uuid.UUID = Field(
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+            unique=True,
+        )
+    )
+    total_bytes_used: int = Field(default=0, nullable=False)
+    updated_at: Optional[datetime.datetime] = Field(
+        default=None,
+        sa_column=Column(
+            TIMESTAMP(timezone=True),
+            server_default=func.now(),
+            onupdate=func.now(),
+            nullable=False,
+        ),
+    )
+
+
 class Folder(SQLModel, table=True):
     __tablename__ = "folders"
 
@@ -253,7 +284,7 @@ class Edge(SQLModel, table=True):
             ["nodes.graph_id", "nodes.id"],
             ondelete="CASCADE",
         ),
-        ForeignKeyConstraint(["graph_id"], ["graphs.id"]),
+        ForeignKeyConstraint(["graph_id"], ["graphs.id"], ondelete="CASCADE"),
         Index("idx_edges_graph_id", "graph_id"),
         Index("idx_edges_source_node", "graph_id", "source_node_id"),
         Index("idx_edges_target_node", "graph_id", "target_node_id"),
@@ -263,8 +294,20 @@ class Edge(SQLModel, table=True):
 class TemplateBookmark(SQLModel, table=True):
     __tablename__ = "template_bookmarks"
 
-    user_id: uuid.UUID = Field(foreign_key="users.id", primary_key=True)
-    template_id: uuid.UUID = Field(foreign_key="prompt_templates.id", primary_key=True)
+    user_id: uuid.UUID = Field(
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            primary_key=True,
+        )
+    )
+    template_id: uuid.UUID = Field(
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("prompt_templates.id", ondelete="CASCADE"),
+            primary_key=True,
+        )
+    )
     created_at: datetime.datetime = Field(
         default_factory=datetime.datetime.now,
         sa_column=Column(TIMESTAMP(timezone=True), server_default=func.now(), nullable=False),
@@ -293,6 +336,8 @@ class User(SQLModel, table=True):
         default="free", sa_column=Column(TEXT, nullable=False)
     )  # Options: "premium", "free"
     is_admin: bool = Field(default=False, nullable=False)
+    is_verified: bool = Field(default=False, nullable=False)
+    has_seen_welcome: bool = Field(default=False, nullable=False)
 
     created_at: Optional[datetime.datetime] = Field(
         default=None,
@@ -329,6 +374,45 @@ class User(SQLModel, table=True):
     )
 
 
+class VerificationToken(SQLModel, table=True):
+    __tablename__ = "verification_tokens"
+
+    id: uuid.UUID = Field(
+        default_factory=uuid.uuid4,
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            primary_key=True,
+            server_default=func.uuid_generate_v4(),
+            nullable=False,
+        ),
+    )
+    user_id: uuid.UUID = Field(
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+    )
+    email: str = Field(index=True, nullable=False)
+    code: str = Field(max_length=6, nullable=False)
+    expires_at: datetime.datetime = Field(
+        sa_column=Column(TIMESTAMP(timezone=True), nullable=False)
+    )
+    created_at: datetime.datetime = Field(
+        default_factory=datetime.datetime.now,
+        sa_column=Column(
+            TIMESTAMP(timezone=True),
+            server_default=func.now(),
+            nullable=False,
+        ),
+    )
+
+    __table_args__ = (
+        Index("idx_verification_tokens_user_id", "user_id"),
+        Index("idx_verification_tokens_email", "email"),
+    )
+
+
 class PromptTemplate(SQLModel, table=True):
     __tablename__ = "prompt_templates"
 
@@ -342,9 +426,12 @@ class PromptTemplate(SQLModel, table=True):
         ),
     )
     user_id: uuid.UUID = Field(
-        foreign_key="users.id",
-        nullable=False,
-        index=True,
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        )
     )
     name: str = Field(max_length=255, nullable=False)
     description: Optional[str] = Field(default=None, sa_column=Column(TEXT))
@@ -399,10 +486,13 @@ class Settings(SQLModel, table=True):
         ),
     )
     user_id: uuid.UUID = Field(
-        foreign_key="users.id",
-        nullable=False,
-        index=True,
-        unique=True,  # Ensures one settings entry per user
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+            unique=True,
+        )
     )
 
     # Store all settings as a single JSONB object
@@ -502,8 +592,11 @@ class RefreshToken(SQLModel, table=True):
         ),
     )
     user_id: uuid.UUID = Field(
-        foreign_key="users.id",
-        nullable=False,
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+        )
     )
     token: str = Field(max_length=255, nullable=False)
     expires_at: datetime.datetime = Field(
@@ -518,7 +611,13 @@ class UsedRefreshToken(SQLModel, table=True):
 
     id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
     token: str = Field(index=True, unique=True, nullable=False)
-    user_id: uuid.UUID = Field(foreign_key="users.id", nullable=False)
+    user_id: uuid.UUID = Field(
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+        )
+    )
     expires_at: datetime.datetime = Field(
         sa_column=Column(TIMESTAMP(timezone=True), nullable=False)
     )
@@ -541,9 +640,12 @@ class ProviderToken(SQLModel, table=True):
         ),
     )
     user_id: uuid.UUID = Field(
-        foreign_key="users.id",
-        nullable=False,
-        index=True,
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        )
     )
     provider: str = Field(max_length=50, nullable=False, index=True)  # e.g., 'github'
     access_token: str = Field(sa_column=Column(TEXT, nullable=False))  # Should always be encrypted
@@ -591,7 +693,14 @@ class Repository(SQLModel, table=True):
             nullable=False,
         ),
     )
-    user_id: uuid.UUID = Field(foreign_key="users.id", nullable=False, index=True)
+    user_id: uuid.UUID = Field(
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        )
+    )
     provider: str = Field(default="github", max_length=50, nullable=False)
     repo_name: str = Field(max_length=255, nullable=False)  # e.g., "my-org/my-awesome-project"
     clone_url: str = Field(sa_column=Column(TEXT, nullable=False))
@@ -649,9 +758,12 @@ class UserQueryUsage(SQLModel, table=True):
         ),
     )
     user_id: uuid.UUID = Field(
-        foreign_key="users.id",
-        nullable=False,
-        index=True,
+        sa_column=Column(
+            PG_UUID(as_uuid=True),
+            ForeignKey("users.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        )
     )
     query_type: str = Field(max_length=50, nullable=False, index=True)
     used_queries: int = Field(default=0, nullable=False)
@@ -698,6 +810,7 @@ async def create_initial_users(
                         username=user.username,
                         password=user.password,
                         oauth_provider="userpass",
+                        is_verified=False,
                     )
                     session.add(new_user)
                     users.append(new_user)