Добавлена модель Company и обновлены связанные модели для учета компании. Обновлены функции регистрации TgAgent с проверкой ключа компании. Изменены комментарии для улучшения понимания кода.

main.py

@@ -15,6 +15,14 @@ import jwt
 AUTH_DATABASE_ADDRESS = "sqlite:///partner.db"
 
 #SQLModel
+class Company(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    name: str
+    commission: float  # процент комиссии
+    key: str = Field(index=True, unique=True)
+    create_dttm: datetime = Field(default_factory=datetime.utcnow)
+    update_dttm: datetime = Field(default_factory=datetime.utcnow)
+
 class TgAgent(SQLModel, table=True):
     id: Optional[int] = Field(default=None, primary_key=True)
     tg_id: int = Field(index=True, unique=True)
@@ -23,6 +31,7 @@ class TgAgent(SQLModel, table=True):
     name: Optional[str] = None
     login: Optional[str] = None
     hash: Optional[str] = None
+    company_id: int = Field(foreign_key="company.id")
     create_dttm: datetime = Field(default_factory=datetime.utcnow)
     update_dttm: datetime = Field(default_factory=datetime.utcnow)
 
@@ -40,6 +49,7 @@ class Sale(SQLModel, table=True):
     crediting: float  # сколько начислено за продажу
     ref: int = Field(foreign_key="ref.id")
     sale_id: str
+    company_id: int = Field(foreign_key="company.id")
     create_dttm: datetime = Field(default_factory=datetime.utcnow)
     update_dttm: datetime = Field(default_factory=datetime.utcnow)
 
@@ -49,6 +59,7 @@ class Transaction(SQLModel, table=True):
     sum: float
     tg_agent_id: int = Field(foreign_key="tgagent.id")
     status: str  # 'process' || 'done' || 'error' || 'waiting'
+    company_id: int = Field(foreign_key="company.id")
     create_dttm: datetime = Field(default_factory=datetime.utcnow)
     update_dttm: datetime = Field(default_factory=datetime.utcnow)
 
@@ -58,7 +69,9 @@ class Account(SQLModel, table=True):
     password_hash: str  # теперь хранится hash пароля
     name: Optional[str] = None
     email: Optional[str] = None
-    balance: float = 0.0
+    company_id: int = Field(foreign_key="company.id")
+    create_dttm: datetime = Field(default_factory=datetime.utcnow)
+    update_dttm: datetime = Field(default_factory=datetime.utcnow)
 
 # Создание движка базы данных
 AUTH_DB_ENGINE = create_engine(AUTH_DATABASE_ADDRESS, echo=True)
@@ -110,12 +123,25 @@ def register(req: RegisterRequest, db: Session = Depends(get_db)):
     phone = req.phone
     name = getattr(req, 'name', None)
     login = getattr(req, 'login', None)
-    print(f'tg_id: {tg_id}, chat_id: {chat_id}, phone: {phone}, name: {name}, login: {login}')
+    company_key = req.company_key
+    print(f'tg_id: {tg_id}, chat_id: {chat_id}, phone: {phone}, name: {name}, login: {login}, company_key: {company_key}')
     tg_agent = get_tg_agent_by_tg_id(db, tg_id)
     if tg_agent:
         raise HTTPException(status_code=400, detail="tg_id already registered")
+    # Поиск компании по ключу
+    company = db.exec(select(Company).where(Company.key == company_key)).first()
+    if not company:
+        raise HTTPException(status_code=400, detail="Компания с таким ключом не найдена")
     hash_value = sha256(f"{tg_id}sold".encode()).hexdigest()
-    new_tg_agent = TgAgent(tg_id=tg_id, chat_id=chat_id, phone=phone, name=name, login=login, hash=hash_value)
+    new_tg_agent = TgAgent(
+        tg_id=tg_id,
+        chat_id=chat_id,
+        phone=phone,
+        name=name,
+        login=login,
+        hash=hash_value,
+        company_id=company.id
+    )
     db.add(new_tg_agent)
     db.commit()
     db.refresh(new_tg_agent)

models.py

@@ -29,3 +29,4 @@ class RegisterRequest(BaseModel):
     phone: Optional[str] = None
     name: Optional[str] = None
     login: Optional[str] = None
+    company_key: str