zephyrdark
52d9fdf1f7
- 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>
79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
"""
|
|
Base collector class for data collection jobs.
|
|
"""
|
|
from abc import ABC, abstractmethod
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models.stock import JobLog
|
|
|
|
|
|
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
|
|
|
|
def start_job(self) -> JobLog:
|
|
"""Create a job log entry when starting."""
|
|
self.job_log = JobLog(
|
|
job_name=self.job_name,
|
|
status="running",
|
|
started_at=datetime.utcnow(),
|
|
)
|
|
self.db.add(self.job_log)
|
|
self.db.commit()
|
|
return self.job_log
|
|
|
|
def complete_job(self, records_count: int):
|
|
"""Mark job as completed."""
|
|
if self.job_log:
|
|
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:
|
|
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:
|
|
"""
|
|
Perform the data collection.
|
|
Returns the number of records collected.
|
|
"""
|
|
pass
|
|
|
|
def run(self) -> JobLog:
|
|
"""Execute the collection job with logging."""
|
|
self.start_job()
|
|
try:
|
|
records = self.collect()
|
|
self.complete_job(records)
|
|
except Exception as e:
|
|
try:
|
|
self.fail_job(str(e))
|
|
except Exception:
|
|
pass # Log update failed, but original exception is more important
|
|
raise
|
|
return self.job_log
|