Initial project version
This commit is contained in:
qananasikq
216
iaai_scraper/storage/db.py
Normal file
216
iaai_scraper/storage/db.py
Normal file
@@ -0,0 +1,216 @@
|
||||
import json
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
from typing import Iterator
|
||||
|
||||
from sqlalchemy import create_engine, or_, select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from ..core.config import Settings
|
||||
from .models import Base, Car, Image, ScrapeTask, SyncRun
|
||||
from .schemas import CarRecord
|
||||
|
||||
logger = logging.getLogger("iaai_scraper.db")
|
||||
|
||||
|
||||
# Поля Car, которые приходят из CarRecord (без id, images, relationship).
|
||||
CAR_DB_FIELDS = {
|
||||
col.key for col in Car.__table__.columns
|
||||
if col.key not in ("id",)
|
||||
}
|
||||
|
||||
|
||||
class PersistenceService:
|
||||
|
||||
def __init__(self, settings: Settings) -> None:
|
||||
# Инициализация engine и фабрики сессий.
|
||||
self.settings = settings
|
||||
engine_kwargs = {
|
||||
"echo": settings.database.echo,
|
||||
"future": True,
|
||||
}
|
||||
# Пул только для PostgreSQL.
|
||||
if "postgresql" in settings.database.url:
|
||||
engine_kwargs["pool_size"] = settings.database.pool_size
|
||||
engine_kwargs["max_overflow"] = settings.database.max_overflow
|
||||
self.engine = create_engine(settings.database.url, **engine_kwargs)
|
||||
self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False, future=True)
|
||||
|
||||
def create_tables(self) -> None:
|
||||
try:
|
||||
Base.metadata.create_all(self.engine)
|
||||
except Exception:
|
||||
# Alembic уже создал таблицы/ENUM — пропускаем.
|
||||
logger.debug("create_tables skipped (schema already exists)")
|
||||
|
||||
@contextmanager
|
||||
def session_scope(self) -> Iterator[Session]:
|
||||
# Единая точка commit/rollback для операций записи.
|
||||
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 в `running`,
|
||||
# помечаем его как failed перед новым запуском.
|
||||
now = datetime.now(timezone.utc)
|
||||
stale_runs = session.execute(select(SyncRun).where(SyncRun.status == "running")).scalars().all()
|
||||
for stale in stale_runs:
|
||||
stale.status = "failed"
|
||||
stale.finished_at = now
|
||||
if not stale.error_summary:
|
||||
stale.error_summary = "Recovered stale running sync run before starting a new run"
|
||||
|
||||
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:
|
||||
# Завершаем sync_run и фиксируем итоговую статистику.
|
||||
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
|
||||
|
||||
def get_existing_origin_urls(self, origin_urls: list[str]) -> set[str]:
|
||||
# Возвращает уже существующие в БД origin_url для фильтрации only-new запусков.
|
||||
if not origin_urls:
|
||||
return set()
|
||||
with self.session_scope() as session:
|
||||
rows = session.execute(select(Car.origin_url).where(Car.origin_url.in_(origin_urls))).all()
|
||||
return {str(row[0]) for row in rows if row and row[0]}
|
||||
|
||||
def get_existing_origin_ids(self, origin_ids: list[str]) -> set[str]:
|
||||
# Возвращает уже существующие в БД origin_id для фильтрации только новых авто.
|
||||
if not origin_ids:
|
||||
return set()
|
||||
with self.session_scope() as session:
|
||||
rows = session.execute(select(Car.origin_id).where(Car.origin_id.in_(origin_ids))).all()
|
||||
return {str(row[0]) for row in rows if row and row[0]}
|
||||
|
||||
@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))
|
||||
|
||||
@staticmethod
|
||||
def _car_payload(record: CarRecord) -> dict[str, object]:
|
||||
payload = record.model_dump(mode="python")
|
||||
result = {key: value for key, value in payload.items() if key in CAR_DB_FIELDS}
|
||||
|
||||
if "raw_attributes" in result and isinstance(result["raw_attributes"], dict):
|
||||
result["raw_attributes"] = json.dumps(result["raw_attributes"], ensure_ascii=False, default=str)
|
||||
return result
|
||||
|
||||
def upsert_car(self, record: CarRecord):
|
||||
"""Insert/update/skip по content_hash."""
|
||||
# В БД отправляем только поля, реально существующие в финальной схеме cars.
|
||||
payload = self._car_payload(record)
|
||||
images = [image.model_dump(mode="python") for image in record.images]
|
||||
content_hash = str(payload.get("content_hash") or "")
|
||||
with self.session_scope() as session:
|
||||
# Сначала пытаемся найти по origin_id, а если ранее origin_id был неполный,
|
||||
# подхватываем существующую запись по origin_url, чтобы не плодить дубли.
|
||||
car = session.execute(
|
||||
select(Car).where(or_(Car.origin_id == record.origin_id, Car.origin_url == record.origin_url))
|
||||
).scalar_one_or_none()
|
||||
action = "inserted"
|
||||
if car is None:
|
||||
car = Car(**payload)
|
||||
session.add(car)
|
||||
session.flush()
|
||||
else:
|
||||
# Если контент не менялся, просто обновляем last_seen_at.
|
||||
if content_hash and car.content_hash == content_hash:
|
||||
car.last_seen_at = record.last_seen_at
|
||||
session.flush()
|
||||
return {"car_id": int(car.id), "images_upserted": 0, "action": "skipped"}
|
||||
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}
|
||||
|
||||
# --- ScrapeTask.
|
||||
|
||||
def create_scrape_task(self, celery_task_id: str, task_type: str, vehicle_url: str | None = None) -> int:
|
||||
"""Создаёт запись задачи Celery."""
|
||||
with self.session_scope() as session:
|
||||
task = ScrapeTask(
|
||||
celery_task_id=celery_task_id,
|
||||
task_type=task_type,
|
||||
vehicle_url=vehicle_url,
|
||||
status="pending",
|
||||
)
|
||||
session.add(task)
|
||||
session.flush()
|
||||
return int(task.id)
|
||||
|
||||
def update_scrape_task(self, celery_task_id: str, **kwargs) -> None:
|
||||
"""Обновляет поля задачи по celery_task_id."""
|
||||
with self.session_scope() as session:
|
||||
task = session.execute(
|
||||
select(ScrapeTask).where(ScrapeTask.celery_task_id == celery_task_id)
|
||||
).scalars().first()
|
||||
if task is None:
|
||||
return
|
||||
for key, value in kwargs.items():
|
||||
if hasattr(task, key):
|
||||
setattr(task, key, value)
|
||||
session.flush()
|
||||
|
||||
def get_scrape_task(self, celery_task_id: str) -> dict | None:
|
||||
"""Возвращает информацию о задаче."""
|
||||
with self.session_scope() as session:
|
||||
task = session.execute(
|
||||
select(ScrapeTask).where(ScrapeTask.celery_task_id == celery_task_id)
|
||||
).scalars().first()
|
||||
if task is None:
|
||||
return None
|
||||
return {
|
||||
"id": task.id,
|
||||
"celery_task_id": task.celery_task_id,
|
||||
"task_type": task.task_type,
|
||||
"status": task.status,
|
||||
"vehicle_url": task.vehicle_url,
|
||||
"created_at": task.created_at.isoformat() if task.created_at else None,
|
||||
"started_at": task.started_at.isoformat() if task.started_at else None,
|
||||
"finished_at": task.finished_at.isoformat() if task.finished_at else None,
|
||||
"result_summary": task.result_summary,
|
||||
"error_message": task.error_message,
|
||||
}
|
||||
Reference in New Issue
Block a user