wag-managment-api-service-v.../Services/Postgres/Models/base_model.py

120 lines
4.2 KiB
Python

from contextlib import contextmanager
from typing import Any, Dict, Optional, Generator
from fastapi import HTTPException, status
from sqlalchemy.orm import Session
class BaseModel:
"""Base model class with common utility functions."""
@contextmanager
def db_session(self) -> Generator[Session, None, None]:
"""Context manager for database session."""
session = self.get_session()
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
finally:
session.close()
def update(self, **kwargs: Dict[str, Any]) -> "BaseModel":
"""Update model instance with given attributes."""
with self.db_session() as session:
# Remove unrelated fields
check_kwargs = self.remove_non_related_inputs(kwargs)
# Handle confirmation logic
is_confirmed_argument = kwargs.get("is_confirmed", None)
if is_confirmed_argument and not len(kwargs) == 1:
self.raise_http_exception(
status_code="HTTP_406_NOT_ACCEPTABLE",
error_case="ConfirmError",
data=kwargs,
message="Confirm field cannot be updated with other fields",
)
# Process system fields
check_kwargs = self.extract_system_fields(check_kwargs, create=False)
# Update attributes
for key, value in check_kwargs.items():
setattr(self, key, value)
# Handle user tracking
if hasattr(self, "creds"):
person_id = getattr(self.creds, "person_id", None)
person_name = getattr(self.creds, "person_name", None)
if person_id and person_name:
if is_confirmed_argument:
self.confirmed_by_id = self.creds.get("person_id", "Unknown")
self.confirmed_by = self.creds.get("person_name", "Unknown")
else:
self.updated_by_id = self.creds.get("person_id", "Unknown")
self.updated_by = self.creds.get("person_name", "Unknown")
session.add(self)
session.flush()
return self
@classmethod
@contextmanager
def create_with_session(
cls, **kwargs: Dict[str, Any]
) -> Generator["BaseModel", None, None]:
"""Create new instance with session management."""
instance = cls()
session = instance.get_session()
try:
check_kwargs = cls.remove_non_related_inputs(instance, kwargs)
check_kwargs = cls.extract_system_fields(
instance, check_kwargs, create=True
)
for key, value in check_kwargs.items():
setattr(instance, key, value)
if hasattr(instance, "creds"):
person_id = getattr(instance.creds, "person_id", None)
person_name = getattr(instance.creds, "person_name", None)
if person_id and person_name:
instance.created_by_id = instance.creds.get("person_id", "Unknown")
instance.created_by = instance.creds.get("person_name", "Unknown")
session.add(instance)
session.flush()
yield instance
session.commit()
except Exception:
session.rollback()
raise
finally:
session.close()
# @router.put("/users/{user_id}")
# async def update_user(
# user_id: str,
# update_data: Dict[str, Any],
# db: Session = Depends(get_db_session)
# ):
# with db_session() as session:
# user = session.query(User).filter(User.id == user_id).first()
# if not user:
# raise HTTPException(status_code=404, detail="User not found")
#
# updated_user = user.update(**update_data)
# return updated_user
#
#
# @router.post("/users")
# async def create_user(
# user_data: Dict[str, Any],
# db: Session = Depends(get_db_session)
# ):
# with User.create_with_session(**user_data) as new_user:
# return new_user