51 lines
1.8 KiB
Python
51 lines
1.8 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
|
|
@staticmethod
|
|
@event.listens_for(Query, "before_compile", retval=True)
|
|
def _filter_deleted(query):
|
|
# Skip filtering if this query explicitly asks to include deleted items
|
|
for option in query._with_options:
|
|
if getattr(option, 'name', None) == 'include_deleted':
|
|
return query
|
|
|
|
# Find entities that inherit from CustomBase
|
|
for entity in query._entities:
|
|
if hasattr(entity, 'entity_zero'):
|
|
entity_zero = entity.entity_zero
|
|
if hasattr(entity_zero, 'mapper') and entity_zero.mapper:
|
|
mapper = entity_zero.mapper
|
|
if issubclass(mapper.class_, CustomBase):
|
|
# Apply filter for this entity
|
|
query = query.filter(mapper.class_.deleted_at.is_(None))
|
|
|
|
return query
|
|
|
|
# Option to include deleted records
|
|
class IncludeDeleted(object):
|
|
name = 'include_deleted'
|
|
|
|
def __init__(self):
|
|
pass |