DFiP_Budget_planing/api/src/repository/user_repository.py

171 lines
5.5 KiB
Python

from typing import Optional
from sqlalchemy import delete, select, text, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from src.db.models.org_unit import OrgUnit
from src.db.models.role import Role
from src.db.models.user_org import UserOrg
from src.db.models.app_user import AppUser
from src.core.security import get_password_hash, verify_password
# from src.domain.models import Role, OrgUnit, UserOrg, AppUser
class UserRepository:
def __init__(self, db: AsyncSession):
self.db = db
async def get(self, user_id: int, load_orgs: bool = False) -> Optional[AppUser]:
query = select(AppUser).where(AppUser.id == user_id)
if load_orgs:
query = query.options(
selectinload(AppUser.org_units),
)
return (
(
await self.db.execute(
query.limit(1)
)
)
.scalars()
.first()
)
async def get_by_email(self, email: str) -> Optional[AppUser]:
return (
(
await self.db.execute(
select(AppUser).where(AppUser.email == email).limit(1)
)
)
.scalars()
.first()
)
async def get_by_username(self, username: str) -> Optional[AppUser]:
return (
(
await self.db.execute(
select(AppUser).where(AppUser.username == username).limit(1)
)
)
.scalars()
.first()
)
async def set_app_user_id(self, user_id: int) -> None:
if isinstance(user_id, int):
await self.db.execute(
text(
f"SET LOCAL app.user_id = {user_id}"
),
)
async def get_list(self, skip: int = 0, limit: int = 100) -> list[AppUser]:
query = (
select(AppUser)
.where(AppUser.is_active.is_(True))
.offset(skip)
.limit(limit)
.order_by(AppUser.id)
)
return (await self.db.execute(query)).scalars().all()
async def create(self, user_data: dict) -> AppUser:
hashed_password = (
get_password_hash(user_data["password"]) if user_data.get("password") else None
)
db_user = AppUser(
email=user_data["email"],
username=user_data["username"],
hashed_password=hashed_password,
full_name=user_data.get("full_name"),
role_id=user_data.get("role_id"),
)
self.db.add(db_user)
await self.db.commit()
await self.db.refresh(db_user)
return db_user
async def update(self, user_id: int, user_data: dict) -> Optional[AppUser]:
user = await self.get(user_id)
if not user:
return None
for field, value in user_data.items():
if field == "password" and value:
value = get_password_hash(value)
setattr(user, field, value)
await self.db.commit()
await self.db.refresh(user)
return user
async def logical_delete(self, user_id: int) -> bool:
query = update(AppUser).where(AppUser.id == user_id).values(is_active=False)
result = await self.db.execute(query)
await self.db.commit()
return result.rowcount > 0
async def authenticate(self, username: str, password: str) -> Optional[AppUser]:
user = await self.get_by_username(username)
if not user:
return None
if not verify_password(password, user.hashed_password):
return None
return user
async def authenticate_via_email(self, email: str) -> Optional[AppUser]:
return await self.get_by_email(email)
async def get_roles(self) -> list[Role]:
return (await self.db.execute(select(Role).order_by(Role.id))).scalars().all()
async def get_many_ssp(self, user_id: int) -> list[OrgUnit]:
query = select(OrgUnit).join(UserOrg, UserOrg.ssp_id == OrgUnit.id).where(
UserOrg.user_id == user_id
)
return (await self.db.execute(query)).scalars().all()
async def get_many_ssp_ids(self, user_id: int) -> list[int]:
query = select(UserOrg.ssp_id).where(UserOrg.user_id == user_id)
return (await self.db.execute(query)).scalars().all()
async def set_many_ssp(self, user_id: int, ssp_ids: list[int]) -> bool:
try:
for ssp_id in ssp_ids:
link = UserOrg(user_id=user_id, ssp_id=ssp_id)
self.db.add(link)
await self.db.commit()
return True
except Exception:
await self.db.rollback()
return False
async def unset_many_ssp(self, user_id: int, ssp_ids: list[int]) -> bool:
try:
for ssp_id in ssp_ids:
query = delete(UserOrg).where(
(UserOrg.user_id == user_id) & (UserOrg.ssp_id == ssp_id)
)
await self.db.execute(query)
await self.db.commit()
return True
except Exception:
await self.db.rollback()
return False
async def clear_many_ssp(self, user_id: int) -> bool:
try:
query = delete(UserOrg).where(UserOrg.user_id == user_id)
await self.db.execute(query)
await self.db.commit()
return True
except Exception:
await self.db.rollback()
return False