fix: add transaction rollback and session validation to BaseCollector

- Add session validation in __init__ to ensure database session is not None
- Implement transaction rollback handling in complete_job() for exception safety
- Implement transaction rollback handling in fail_job() for exception safety
- Improve exception handling in run() to gracefully handle fail_job failures while preserving original exception

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
zephyrdark 2026-02-02 23:36:39 +09:00
parent 3abbdfa5b6
commit 52d9fdf1f7

View File

@ -14,6 +14,8 @@ class BaseCollector(ABC):
"""Base class for all data collectors."""
def __init__(self, db: Session):
if db is None:
raise ValueError("Database session cannot be None")
self.db = db
self.job_name = self.__class__.__name__
self.job_log: Optional[JobLog] = None
@ -32,18 +34,26 @@ class BaseCollector(ABC):
def complete_job(self, records_count: int):
"""Mark job as completed."""
if self.job_log:
self.job_log.status = "success"
self.job_log.finished_at = datetime.utcnow()
self.job_log.records_count = records_count
self.db.commit()
try:
self.job_log.status = "success"
self.job_log.finished_at = datetime.utcnow()
self.job_log.records_count = records_count
self.db.commit()
except Exception:
self.db.rollback()
raise
def fail_job(self, error_msg: str):
"""Mark job as failed."""
if self.job_log:
self.job_log.status = "failed"
self.job_log.finished_at = datetime.utcnow()
self.job_log.error_msg = error_msg
self.db.commit()
try:
self.job_log.status = "failed"
self.job_log.finished_at = datetime.utcnow()
self.job_log.error_msg = error_msg
self.db.commit()
except Exception:
self.db.rollback()
raise
@abstractmethod
def collect(self) -> int:
@ -60,6 +70,9 @@ class BaseCollector(ABC):
records = self.collect()
self.complete_job(records)
except Exception as e:
self.fail_job(str(e))
try:
self.fail_job(str(e))
except Exception:
pass # Log update failed, but original exception is more important
raise
return self.job_log