Move teams and user_teams into shared/models

Author: mina-parham

transformerlab/models/users.py

@@ -5,7 +5,8 @@
 from fastapi_users import BaseUserManager, FastAPIUsers, UUIDIDMixin, schemas
 from fastapi_users.authentication import AuthenticationBackend, BearerTransport, JWTStrategy
 from fastapi_users.db import SQLAlchemyUserDatabase
-from transformerlab.shared.models.user_model import User, get_async_session, create_default_team, UserTeam
+from transformerlab.shared.models.user_model import User, get_async_session, create_default_team
+from transformerlab.shared.models.models import UserTeam
 from sqlalchemy.ext.asyncio import AsyncSession
 from jose import jwt as _jose_jwt
 from datetime import datetime, timedelta

transformerlab/routers/auth2.py

@@ -1,5 +1,6 @@
 from fastapi import APIRouter, Depends, HTTPException
-from transformerlab.shared.models.user_model import User, Team, UserTeam, get_async_session, create_default_team
+from transformerlab.shared.models.user_model import User, get_async_session, create_default_team
+from transformerlab.shared.models.models import Team, UserTeam
 from transformerlab.models.users import (
     fastapi_users,
     auth_backend,

transformerlab/routers/teams.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter, Depends, HTTPException
 from sqlalchemy.ext.asyncio import AsyncSession
-from transformerlab.shared.models.user_model import User, Team, UserTeam, get_async_session
+from transformerlab.shared.models.user_model import User, get_async_session
+from transformerlab.shared.models.models import Team, UserTeam
 from transformerlab.models.users import current_active_user
 from pydantic import BaseModel
 from sqlalchemy import select, delete, update

transformerlab/shared/models/models.py

@@ -1,6 +1,7 @@
 from typing import Optional
-from sqlalchemy import String, JSON, DateTime, func, Integer, Index
+from sqlalchemy import String, JSON, DateTime, func, Integer, Index, ForeignKey
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+import uuid
 
 
 class Base(DeclarativeBase):
@@ -82,3 +83,20 @@ class WorkflowRun(Base):
     updated_at: Mapped[DateTime] = mapped_column(
         DateTime, server_default=func.now(), onupdate=func.now(), nullable=False
     )
+
+class Team(Base):
+    """Team model."""
+
+    __tablename__ = "teams"
+
+    id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    name: Mapped[str] = mapped_column(String, nullable=False)
+
+
+class UserTeam(Base):
+    """User-Team association model."""
+
+    __tablename__ = "users_teams"
+
+    user_id: Mapped[str] = mapped_column(String, ForeignKey("user.id"), primary_key=True)
+    team_id: Mapped[str] = mapped_column(String, ForeignKey("teams.id"), primary_key=True)

transformerlab/shared/models/user_model.py

@@ -1,39 +1,20 @@
 # database.py
 from typing import AsyncGenerator
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.ext.declarative import DeclarativeMeta, declarative_base
 from sqlalchemy.orm import sessionmaker
 from fastapi_users.db import SQLAlchemyBaseUserTableUUID
-from sqlalchemy import Column, String, ForeignKey, select
-import uuid
+from sqlalchemy import select
 
 # Replace with your actual database URL (e.g., PostgreSQL, SQLite)
 from transformerlab.db.constants import DATABASE_URL
-
-Base: DeclarativeMeta = declarative_base()
+from .models import Base, Team
 
 
 # 1. Define your User Model (inherits from a FastAPI Users base class)
 class User(SQLAlchemyBaseUserTableUUID, Base):
     pass  # You can add custom fields here later, like 'first_name: str'
 
 
-# 2. Define Team Model
-class Team(Base):
-    __tablename__ = "teams"
-
-    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
-    name = Column(String, nullable=False)
-
-
-# 3. Define User-Team Association Model
-class UserTeam(Base):
-    __tablename__ = "users_teams"
-
-    user_id = Column(String, ForeignKey("user.id"), primary_key=True)
-    team_id = Column(String, ForeignKey("teams.id"), primary_key=True)
-
-
 # 2. Setup the Async Engine and Session
 engine = create_async_engine(DATABASE_URL)
 AsyncSessionLocal = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)