codex / src /database /external_batch_crud.py
webhook's picture
Upload 137 files
89419ea verified
Raw History Blame
3.21 kB
from __future__ import annotations
from datetime import datetime
from typing import Optional, Sequence
from sqlalchemy.orm import Session
from .models import ExternalRegistrationBatch, ExternalRegistrationBatchItem
def create_batch(db: Session, **kwargs) -> ExternalRegistrationBatch:
batch = ExternalRegistrationBatch(**kwargs)
db.add(batch)
db.flush()
db.refresh(batch)
return batch
def create_batch_item(db: Session, **kwargs) -> ExternalRegistrationBatchItem:
item = ExternalRegistrationBatchItem(**kwargs)
db.add(item)
db.flush()
db.refresh(item)
return item
def get_batch_by_uuid(db: Session, batch_uuid: str) -> Optional[ExternalRegistrationBatch]:
return db.query(ExternalRegistrationBatch).filter(ExternalRegistrationBatch.batch_uuid == batch_uuid).first()
def get_batch_by_idempotency_key(db: Session, idempotency_key: str) -> Optional[ExternalRegistrationBatch]:
if not idempotency_key:
return None
return db.query(ExternalRegistrationBatch).filter(ExternalRegistrationBatch.idempotency_key == idempotency_key).first()
def list_batch_items(db: Session, batch_uuid: str) -> list[ExternalRegistrationBatchItem]:
return (
db.query(ExternalRegistrationBatchItem)
.join(ExternalRegistrationBatch, ExternalRegistrationBatch.id == ExternalRegistrationBatchItem.batch_id)
.filter(ExternalRegistrationBatch.batch_uuid == batch_uuid)
.order_by(ExternalRegistrationBatchItem.item_index.asc(), ExternalRegistrationBatchItem.id.asc())
.all()
)
def list_recoverable_batches(db: Session) -> list[ExternalRegistrationBatch]:
return (
db.query(ExternalRegistrationBatch)
.filter(ExternalRegistrationBatch.status.in_(("pending", "running")))
.order_by(ExternalRegistrationBatch.created_at.asc(), ExternalRegistrationBatch.id.asc())
.all()
)
def update_batch(db: Session, batch_uuid: str, **kwargs) -> Optional[ExternalRegistrationBatch]:
batch = get_batch_by_uuid(db, batch_uuid)
if not batch:
return None
for key, value in kwargs.items():
setattr(batch, key, value)
batch.updated_at = datetime.utcnow()
db.flush()
db.refresh(batch)
return batch
def update_batch_item(db: Session, item_id: int, **kwargs) -> Optional[ExternalRegistrationBatchItem]:
item = db.query(ExternalRegistrationBatchItem).filter(ExternalRegistrationBatchItem.id == item_id).first()
if not item:
return None
for key, value in kwargs.items():
setattr(item, key, value)
item.updated_at = datetime.utcnow()
db.flush()
db.refresh(item)
return item
def list_items_by_status(db: Session, batch_uuid: str, statuses: Sequence[str]) -> list[ExternalRegistrationBatchItem]:
return (
db.query(ExternalRegistrationBatchItem)
.join(ExternalRegistrationBatch, ExternalRegistrationBatch.id == ExternalRegistrationBatchItem.batch_id)
.filter(
ExternalRegistrationBatch.batch_uuid == batch_uuid,
ExternalRegistrationBatchItem.status.in_(tuple(statuses)),
)
.order_by(ExternalRegistrationBatchItem.item_index.asc())
.all()
)