SQLAlchemy 爬蟲資料儲存實戰：可上線的資料層設計

from __future__ import annotations

import datetime as dt
from typing import Annotated, Any

from sqlalchemy import (
    BigInteger,
    ForeignKey,
    Index,
    MetaData,
    String,
    TIMESTAMP,
    UniqueConstraint,
    func,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship

# Column recipes declared once and annotated everywhere. Changing the policy
# for every timestamp in the schema is then a one-line edit, not a sweep.
bigint_pk = Annotated[int, mapped_column(BigInteger, primary_key=True)]
utc_now = Annotated[dt.datetime, mapped_column(server_default=func.now())]

# Every constraint gets a deterministic name. Alembic cannot generate a
# migration for a constraint the database named anonymously, so this is what
# makes autogenerate usable later.
NAMING_CONVENTION = {
    "ix": "ix_%(column_0_label)s",
    "uq": "uq_%(table_name)s_%(column_0_name)s",
    "ck": "ck_%(table_name)s_%(constraint_name)s",
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s",
}


class Base(DeclarativeBase):
    metadata = MetaData(naming_convention=NAMING_CONVENTION)

    # A naive timestamp is the classic silent data loss in a crawler that runs
    # in more than one region. Fix it once, for every model in the schema.
    type_annotation_map = {
        dt.datetime: TIMESTAMP(timezone=True),
        dict[str, Any]: JSONB,
    }


class ScrapedPage(Base):
    __tablename__ = "scraped_page"

    id: Mapped[bigint_pk]
    source: Mapped[str] = mapped_column(String(64))
    external_id: Mapped[str] = mapped_column(String(256))
    url: Mapped[str] = mapped_column(String(2048))
    # Mapped[str] is NOT NULL; Mapped[str | None] is NULL. The annotation IS
    # the schema, so mypy and the DDL cannot disagree about optionality.
    title: Mapped[str | None]
    payload: Mapped[dict[str, Any]]
    # Hash what you EXTRACTED, never the raw HTML: hash the page and every
    # rotating ad slot and CSRF token reads as a content change.
    content_hash: Mapped[str] = mapped_column(String(64))
    first_seen_at: Mapped[utc_now]
    last_changed_at: Mapped[utc_now]

    observations: Mapped[list[PriceObservation]] = relationship(
        back_populates="page", cascade="all, delete-orphan"
    )

    __table_args__ = (
        # The natural key: what "the same page, seen again" means. Without it
        # the upsert in the next section has nothing to conflict on.
        UniqueConstraint("source", "external_id", name="uq_scraped_page_natural_key"),
        Index("ix_scraped_page_changed", "source", "last_changed_at"),
    )


class PriceObservation(Base):
    __tablename__ = "price_observation"

    id: Mapped[bigint_pk]
    page_id: Mapped[int] = mapped_column(ForeignKey("scraped_page.id", ondelete="CASCADE"))
    observed_at: Mapped[utc_now]
    # Money as integer minor units. A float price silently loses cents, and
    # you will not notice until someone reconciles a year of history.
    amount_minor: Mapped[int]
    currency: Mapped[str] = mapped_column(String(3))

    page: Mapped[ScrapedPage] = relationship(back_populates="observations")

import os

from sqlalchemy import create_engine

engine = create_engine(
    # Never a literal DSN. Credentials belong in the environment or a secrets
    # manager, and the URL carries the driver: postgresql+psycopg for sync.
    os.environ["SCRAPER_DATABASE_URL"],
    pool_size=5,
    max_overflow=5,          # hard ceiling of 10 per process, not 15
    pool_timeout=10,         # fail fast; a queued worker is a stalled worker
    pool_recycle=1800,       # under a proxy or a managed DB, assume idle kills
    pool_pre_ping=True,      # one cheap round trip beats one lost batch
    insertmanyvalues_page_size=500,
)

from multiprocessing import Pool


def init_worker() -> None:
    # close=False drops the inherited connection references WITHOUT closing
    # them — the parent still owns those sockets. The child then opens its own.
    engine.dispose(close=False)


with Pool(processes=8, initializer=init_worker) as pool:
    pool.map(scrape_one, urls)

from sqlalchemy.orm import sessionmaker

# Once, at import time. The factory is cheap and shareable; sessions are not.
SessionFactory = sessionmaker(engine)


def store_batch(rows: list[dict[str, object]]) -> None:
    # One transaction per BATCH. Per row, commit overhead dominates the run;
    # per crawl, one bad page rolls back an hour of work while every row you
    # did write stays locked away from readers.
    with SessionFactory() as session, session.begin():
        session.execute(upsert_pages(rows))

from sqlalchemy import func
from sqlalchemy.dialects.postgresql import insert


def upsert_pages(rows: list[dict[str, object]]):
    stmt = insert(ScrapedPage).values(rows)
    return stmt.on_conflict_do_update(
        # Infers the unique index behind the natural key.
        index_elements=[ScrapedPage.source, ScrapedPage.external_id],
        # `excluded` is the row PostgreSQL tried to insert — the fresh scrape.
        set_={
            "url": stmt.excluded.url,
            "title": stmt.excluded.title,
            "payload": stmt.excluded.payload,
            "content_hash": stmt.excluded.content_hash,
            "last_changed_at": func.now(),
        },
        # The line that pays for itself: a re-crawl that found nothing new
        # writes no row at all. No dead tuple, no WAL, no index churn.
        where=ScrapedPage.content_hash != stmt.excluded.content_hash,
    )

changed = session.scalars(
    upsert_pages(rows).returning(ScrapedPage),
    # Refreshes any of these objects already living in the session's identity
    # map — unlike a plain INSERT, some of these rows already existed.
    execution_options={"populate_existing": True},
).all()

for page in changed:
    enqueue_downstream(page.id)

from sqlalchemy import insert

session.execute(
    # render_nulls keeps every row in one batch. Scraped records are ragged by
    # nature — an optional field that is None would otherwise split the run
    # into several statements at exactly the rows you have most of.
    insert(ScrapedPage).execution_options(render_nulls=True),
    rows,
)

page_ids = session.scalars(
    # sort_by_parameter_order guarantees the returned ids line up with the
    # input rows, which is what lets you attach child records without a
    # second SELECT. Added in SQLAlchemy 2.0.10.
    insert(ScrapedPage).returning(ScrapedPage.id, sort_by_parameter_order=True),
    rows,
).all()

import asyncio
import os

from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

# The URL carries the async driver: postgresql+asyncpg://...
engine = create_async_engine(os.environ["SCRAPER_DATABASE_URL"], pool_size=10, max_overflow=0)

# expire_on_commit=False so attributes stay readable after commit. The default
# would expire them, and refreshing an expired attribute is implicit IO —
# exactly what asyncio cannot do behind your back.
AsyncSessionFactory = async_sessionmaker(engine, expire_on_commit=False)


async def store(rows: list[dict[str, object]], gate: asyncio.Semaphore) -> None:
    # One AsyncSession per task: a single instance is explicitly documented as
    # unsafe in concurrent tasks. The semaphore keeps the task count from
    # outrunning the pool, which is the other half of the same budget.
    async with gate, AsyncSessionFactory() as session, session.begin():
        await session.execute(upsert_pages(rows))


async def main(batches: list[list[dict[str, object]]]) -> None:
    gate = asyncio.Semaphore(10)
    try:
        await asyncio.gather(*(store(batch, gate) for batch in batches))
    finally:
        # Closes the pool. Skip it and the loop shuts down under open sockets.
        await engine.dispose()

from sqlalchemy import select
from sqlalchemy.orm import raiseload, selectinload

stmt = (
    select(ScrapedPage)
    .options(
        # selectinload is the documented default for COLLECTIONS: a second
        # SELECT with an IN clause, so no join multiplies the parent rows.
        selectinload(ScrapedPage.observations),
        # Everything not eager-loaded above now RAISES instead of quietly
        # emitting a query. An N+1 becomes a test failure, not a slow export.
        raiseload("*"),
    )
    .where(ScrapedPage.source == "example-retailer")
    .order_by(ScrapedPage.last_changed_at.desc())
    .limit(500)
)
pages = session.scalars(stmt).all()

import logging
import time

from sqlalchemy import event

log = logging.getLogger("scraper.sql")
SLOW_QUERY_SECONDS = 0.5


@event.listens_for(engine, "before_cursor_execute")
def _start_timer(conn, cursor, statement, parameters, context, executemany):
    conn.info.setdefault("query_start", []).append(time.perf_counter())


@event.listens_for(engine, "after_cursor_execute")
def _record_duration(conn, cursor, statement, parameters, context, executemany):
    elapsed = time.perf_counter() - conn.info["query_start"].pop()
    if elapsed >= SLOW_QUERY_SECONDS:
        # The statement, never `parameters`: that tuple is the payload.
        log.warning("slow sql %.3fs executemany=%s %s", elapsed, executemany, statement)

import random
import time
from collections.abc import Callable
from typing import TypeVar

from sqlalchemy.exc import DBAPIError

T = TypeVar("T")
# Serialization failure and deadlock detected: transient by definition.
RETRYABLE_SQLSTATES = frozenset({"40001", "40P01"})


def _is_retryable(error: DBAPIError) -> bool:
    if error.connection_invalidated:
        return True
    orig = error.orig
    return getattr(orig, "sqlstate", None) in RETRYABLE_SQLSTATES


def with_retry(work: Callable[[], T], attempts: int = 5) -> T:
    """Re-run a whole unit of work, from BEGIN, on a transient failure."""
    for attempt in range(1, attempts + 1):
        try:
            return work()
        except DBAPIError as error:
            if attempt == attempts or not _is_retryable(error):
                raise
            # Full jitter. A fleet that backs off on one schedule reconnects
            # as a thundering herd and knocks the database over a second time.
            time.sleep(random.uniform(0, min(2**attempt * 0.1, 5.0)))
    raise AssertionError("unreachable")


with_retry(lambda: store_batch(rows))

alembic revision --autogenerate -m "add price_observation"
alembic upgrade head

import pytest
from sqlalchemy.orm import Session


@pytest.fixture()
def session(engine):
    connection = engine.connect()
    transaction = connection.begin()
    # join_transaction_mode="create_savepoint" lets the CODE UNDER TEST call
    # session.commit() for real — the commit lands on a SAVEPOINT inside the
    # outer transaction, which the fixture then throws away wholesale.
    db = Session(bind=connection, join_transaction_mode="create_savepoint")
    try:
        yield db
    finally:
        db.close()
        transaction.rollback()
        connection.close()