diff --git a/main.py b/main.py index da20bad..0054859 100644 --- a/main.py +++ b/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) diff --git a/models.py b/models.py index c04bd76..c4fd30c 100644 --- a/models.py +++ b/models.py @@ -28,4 +28,5 @@ class RegisterRequest(BaseModel): chat_id: Optional[int] = None phone: Optional[str] = None name: Optional[str] = None - login: Optional[str] = None \ No newline at end of file + login: Optional[str] = None + company_key: str \ No newline at end of file