SupplyAgent/supply_infra/db/repositories/generated_demand_repo.py

from __future__ import annotations

from sqlalchemy.dialects.mysql import insert

from supply_infra.db.models.generated_demand import GeneratedDemand
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 500


class GeneratedDemandRepository(BaseRepository[GeneratedDemand]):
    """单维度需求产生结果表 repository。"""

    model = GeneratedDemand

    def bulk_insert(self, rows: list[dict]) -> int:
        """批量插入生成结果行。"""
        if not rows:
            return 0

        inserted = 0
        for i in range(0, len(rows), _BATCH_SIZE):
            batch = rows[i : i + _BATCH_SIZE]
            stmt = insert(GeneratedDemand).values(batch)
            result = self.session.execute(stmt)
            inserted += result.rowcount or 0
        return inserted