47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from sqlalchemy import Column, DateTime, event, func, inspect
|
|
from sqlalchemy.orm import Query
|
|
from database import SessionLocal
|
|
from datetime import datetime
|
|
|
|
class CustomBase:
|
|
create_time = Column(DateTime(timezone=True), server_default=func.now())
|
|
update_time = Column(
|
|
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
|
)
|
|
deleted_at = Column(DateTime(timezone=True), nullable=True)
|
|
|
|
def soft_delete(self, session):
|
|
"""Mark record as deleted without removing from database"""
|
|
self.deleted_at = datetime.now()
|
|
session.add(self)
|
|
session.commit()
|
|
|
|
def restore(self, session):
|
|
"""Restore a soft-deleted record"""
|
|
self.deleted_at = None
|
|
session.add(self)
|
|
session.commit()
|
|
|
|
# Global filter for deleted records
|
|
@event.listens_for(SessionLocal, "do_orm_execute")
|
|
def _add_filtering_criteria(execute_state):
|
|
if (
|
|
execute_state.is_select
|
|
and not execute_state.execution_options.get("include_deleted", False)
|
|
):
|
|
# Check if any of the entities inherit from CustomBase
|
|
for entity in execute_state.statement.column_descriptions:
|
|
entity_class = entity.get("entity", None)
|
|
if entity_class and issubclass(entity_class, CustomBase):
|
|
# Add filter condition to exclude soft-deleted records
|
|
execute_state.statement = execute_state.statement.filter(
|
|
entity_class.deleted_at.is_(None)
|
|
)
|
|
break
|
|
|
|
# Option to include deleted records
|
|
class IncludeDeleted(object):
|
|
name = 'include_deleted'
|
|
|
|
def __init__(self):
|
|
pass |