101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
import logging
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timezone
|
|
from typing import Iterator
|
|
|
|
from sqlalchemy import create_engine, select
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from ..core.config import Settings
|
|
from .models import Base, Car, Image, SyncRun
|
|
from .schemas import CarRecord
|
|
|
|
logger = logging.getLogger("iaai_scraper.db")
|
|
|
|
|
|
class PersistenceService:
|
|
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
self.engine = create_engine(settings.database.url, echo=settings.database.echo, future=True)
|
|
self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False, future=True)
|
|
|
|
def create_tables(self) -> None:
|
|
Base.metadata.create_all(self.engine)
|
|
|
|
@contextmanager
|
|
def session_scope(self) -> Iterator[Session]:
|
|
session = self.session_factory()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
def start_sync_run(self, lane: str) -> int:
|
|
with self.session_scope() as session:
|
|
run = SyncRun(status="running", lane=lane, ids_fetched=0, cars_upserted=0, cars_failed=0, images_upserted=0)
|
|
session.add(run)
|
|
session.flush()
|
|
return int(run.id)
|
|
|
|
def finish_sync_run(self, run_id: int, *, status: str, ids_fetched: int, cars_upserted: int, cars_failed: int, images_upserted: int, error_summary: str | None = None) -> None:
|
|
with self.session_scope() as session:
|
|
run = session.get(SyncRun, run_id)
|
|
if run is None:
|
|
return
|
|
run.finished_at = datetime.now(timezone.utc)
|
|
run.status = status
|
|
run.ids_fetched = ids_fetched
|
|
run.cars_upserted = cars_upserted
|
|
run.cars_failed = cars_failed
|
|
run.images_upserted = images_upserted
|
|
run.error_summary = error_summary
|
|
|
|
@staticmethod
|
|
def _add_images(session: Session, car_id: int, images: list[dict[str, object]]) -> None:
|
|
for image_payload in images:
|
|
session.add(Image(fullres_image=str(image_payload["fullres_image"]), preview_image=str(image_payload["preview_image"]), order_index=int(image_payload.get("order_index", 0)), car_id=car_id))
|
|
|
|
def upsert_car(self, record: CarRecord):
|
|
"""Insert/update/skip по content_hash."""
|
|
payload = record.model_dump(mode="python")
|
|
images = payload.pop("images", [])
|
|
payload.pop("raw_attributes", None)
|
|
payload.pop("mapping_notes", None)
|
|
content_hash = payload.pop("content_hash", "")
|
|
with self.session_scope() as session:
|
|
# поиск по origin_id
|
|
car = session.execute(select(Car).where(Car.origin_id == record.origin_id)).scalar_one_or_none()
|
|
action = "inserted"
|
|
if car is None:
|
|
car = Car(**payload)
|
|
session.add(car)
|
|
session.flush()
|
|
else:
|
|
action = "updated"
|
|
for key, value in payload.items():
|
|
setattr(car, key, value)
|
|
car.last_seen_at = record.last_seen_at
|
|
session.flush()
|
|
# замена картинок в savepoint
|
|
nested = session.begin_nested()
|
|
try:
|
|
for image in list(car.images):
|
|
session.delete(image)
|
|
session.flush()
|
|
self._add_images(session, int(car.id), images)
|
|
session.flush()
|
|
nested.commit()
|
|
except Exception:
|
|
nested.rollback()
|
|
logger.warning("Image replacement failed for car %s, keeping old images", record.origin_id)
|
|
images = []
|
|
return {"car_id": int(car.id), "images_upserted": len(images), "action": action}
|
|
self._add_images(session, int(car.id), images)
|
|
session.flush()
|
|
return {"car_id": int(car.id), "images_upserted": len(images), "action": action}
|