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:
parent
3abbdfa5b6
commit
52d9fdf1f7
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user