move to postgresql
This commit is contained in:
@@ -1,175 +1,213 @@
|
||||
"""Database operations for DetectJob and StageCheckpoint."""
|
||||
"""Database operations for detection pipeline — SQLModel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
from uuid import UUID
|
||||
|
||||
from sqlmodel import select
|
||||
|
||||
from .connection import get_session
|
||||
from .models import (
|
||||
DetectJob, StageCheckpoint, KnownBrand, SourceBrandSighting,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DetectJob
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def create_detect_job(**fields):
|
||||
from admin.mpr.media_assets.models import DetectJob
|
||||
return DetectJob.objects.create(**fields)
|
||||
def create_detect_job(**fields) -> DetectJob:
|
||||
job = DetectJob(**fields)
|
||||
with get_session() as session:
|
||||
session.add(job)
|
||||
session.commit()
|
||||
session.refresh(job)
|
||||
return job
|
||||
|
||||
|
||||
def get_detect_job(id: UUID):
|
||||
from admin.mpr.media_assets.models import DetectJob
|
||||
return DetectJob.objects.get(id=id)
|
||||
def get_detect_job(id: UUID) -> DetectJob | None:
|
||||
with get_session() as session:
|
||||
return session.get(DetectJob, id)
|
||||
|
||||
|
||||
def update_detect_job(job_id: UUID, **fields):
|
||||
from admin.mpr.media_assets.models import DetectJob
|
||||
DetectJob.objects.filter(id=job_id).update(**fields)
|
||||
def update_detect_job(job_id: UUID, **fields) -> None:
|
||||
with get_session() as session:
|
||||
job = session.get(DetectJob, job_id)
|
||||
if not job:
|
||||
return
|
||||
for k, v in fields.items():
|
||||
setattr(job, k, v)
|
||||
session.commit()
|
||||
|
||||
|
||||
def list_detect_jobs(
|
||||
parent_job_id: Optional[UUID] = None,
|
||||
status: Optional[str] = None,
|
||||
):
|
||||
from admin.mpr.media_assets.models import DetectJob
|
||||
|
||||
qs = DetectJob.objects.all()
|
||||
if parent_job_id:
|
||||
qs = qs.filter(parent_job_id=parent_job_id)
|
||||
if status:
|
||||
qs = qs.filter(status=status)
|
||||
return list(qs)
|
||||
) -> list[DetectJob]:
|
||||
with get_session() as session:
|
||||
stmt = select(DetectJob)
|
||||
if parent_job_id:
|
||||
stmt = stmt.where(DetectJob.parent_job_id == parent_job_id)
|
||||
if status:
|
||||
stmt = stmt.where(DetectJob.status == status)
|
||||
return list(session.exec(stmt).all())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# StageCheckpoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def save_stage_checkpoint(**fields):
|
||||
from admin.mpr.media_assets.models import StageCheckpoint
|
||||
return StageCheckpoint.objects.create(**fields)
|
||||
def save_stage_checkpoint(**fields) -> StageCheckpoint:
|
||||
with get_session() as session:
|
||||
# Upsert: replace if same job_id + stage
|
||||
job_id = fields.get("job_id")
|
||||
stage = fields.get("stage")
|
||||
if job_id and stage:
|
||||
stmt = select(StageCheckpoint).where(
|
||||
StageCheckpoint.job_id == job_id,
|
||||
StageCheckpoint.stage == stage,
|
||||
)
|
||||
existing = session.exec(stmt).first()
|
||||
if existing:
|
||||
for k, v in fields.items():
|
||||
setattr(existing, k, v)
|
||||
session.commit()
|
||||
session.refresh(existing)
|
||||
return existing
|
||||
|
||||
checkpoint = StageCheckpoint(**fields)
|
||||
session.add(checkpoint)
|
||||
session.commit()
|
||||
session.refresh(checkpoint)
|
||||
return checkpoint
|
||||
|
||||
|
||||
def get_stage_checkpoint(job_id: UUID, stage: str):
|
||||
from admin.mpr.media_assets.models import StageCheckpoint
|
||||
return StageCheckpoint.objects.get(job_id=job_id, stage=stage)
|
||||
def get_stage_checkpoint(job_id: UUID, stage: str) -> StageCheckpoint | None:
|
||||
with get_session() as session:
|
||||
stmt = select(StageCheckpoint).where(
|
||||
StageCheckpoint.job_id == job_id,
|
||||
StageCheckpoint.stage == stage,
|
||||
)
|
||||
return session.exec(stmt).first()
|
||||
|
||||
|
||||
def list_stage_checkpoints(job_id: UUID) -> list[str]:
|
||||
from admin.mpr.media_assets.models import StageCheckpoint
|
||||
|
||||
stages = (
|
||||
StageCheckpoint.objects
|
||||
.filter(job_id=job_id)
|
||||
.order_by("stage_index")
|
||||
.values_list("stage", flat=True)
|
||||
)
|
||||
return list(stages)
|
||||
with get_session() as session:
|
||||
stmt = (
|
||||
select(StageCheckpoint.stage)
|
||||
.where(StageCheckpoint.job_id == job_id)
|
||||
.order_by(StageCheckpoint.stage_index)
|
||||
)
|
||||
return list(session.exec(stmt).all())
|
||||
|
||||
|
||||
def delete_stage_checkpoints(job_id: UUID):
|
||||
from admin.mpr.media_assets.models import StageCheckpoint
|
||||
StageCheckpoint.objects.filter(job_id=job_id).delete()
|
||||
def delete_stage_checkpoints(job_id: UUID) -> None:
|
||||
with get_session() as session:
|
||||
stmt = select(StageCheckpoint).where(StageCheckpoint.job_id == job_id)
|
||||
for cp in session.exec(stmt).all():
|
||||
session.delete(cp)
|
||||
session.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# KnownBrand
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def get_or_create_brand(canonical_name: str, aliases: list[str] | None = None,
|
||||
source: str = "ocr") -> tuple:
|
||||
"""Get existing brand or create new one. Returns (brand, created)."""
|
||||
from admin.mpr.media_assets.models import KnownBrand
|
||||
import uuid
|
||||
|
||||
def get_or_create_brand(canonical_name: str, aliases: Optional[list[str]] = None,
|
||||
source: str = "ocr") -> tuple[KnownBrand, bool]:
|
||||
normalized = canonical_name.strip()
|
||||
brand = KnownBrand.objects.filter(canonical_name__iexact=normalized).first()
|
||||
if brand:
|
||||
return brand, False
|
||||
with get_session() as session:
|
||||
stmt = select(KnownBrand).where(KnownBrand.canonical_name.ilike(normalized))
|
||||
brand = session.exec(stmt).first()
|
||||
if brand:
|
||||
return brand, False
|
||||
|
||||
# Check aliases of existing brands
|
||||
for existing in KnownBrand.objects.all():
|
||||
existing_aliases = [a.lower() for a in (existing.aliases or [])]
|
||||
if normalized.lower() in existing_aliases:
|
||||
return existing, False
|
||||
|
||||
brand = KnownBrand.objects.create(
|
||||
id=uuid.uuid4(),
|
||||
canonical_name=normalized,
|
||||
aliases=aliases or [],
|
||||
first_source=source,
|
||||
)
|
||||
return brand, True
|
||||
brand = KnownBrand(
|
||||
canonical_name=normalized,
|
||||
aliases=aliases or [],
|
||||
first_source=source,
|
||||
)
|
||||
session.add(brand)
|
||||
session.commit()
|
||||
session.refresh(brand)
|
||||
return brand, True
|
||||
|
||||
|
||||
def find_brand_by_text(text: str) -> Optional[object]:
|
||||
"""Find a known brand by canonical name or alias (case-insensitive)."""
|
||||
from admin.mpr.media_assets.models import KnownBrand
|
||||
|
||||
def find_brand_by_text(text: str) -> KnownBrand | None:
|
||||
normalized = text.strip().lower()
|
||||
|
||||
# Exact canonical match
|
||||
brand = KnownBrand.objects.filter(canonical_name__iexact=normalized).first()
|
||||
if brand:
|
||||
return brand
|
||||
|
||||
# Search aliases (jsonb contains)
|
||||
for brand in KnownBrand.objects.all():
|
||||
brand_aliases = [a.lower() for a in (brand.aliases or [])]
|
||||
if normalized in brand_aliases:
|
||||
with get_session() as session:
|
||||
stmt = select(KnownBrand).where(KnownBrand.canonical_name.ilike(normalized))
|
||||
brand = session.exec(stmt).first()
|
||||
if brand:
|
||||
return brand
|
||||
|
||||
return None
|
||||
# Alias search — check if normalized is in any brand's aliases
|
||||
all_brands = session.exec(select(KnownBrand)).all()
|
||||
for b in all_brands:
|
||||
if normalized in [a.lower() for a in (b.aliases or [])]:
|
||||
return b
|
||||
return None
|
||||
|
||||
|
||||
def list_all_brands() -> list:
|
||||
from admin.mpr.media_assets.models import KnownBrand
|
||||
return list(KnownBrand.objects.all().order_by("canonical_name"))
|
||||
def list_all_brands() -> list[KnownBrand]:
|
||||
with get_session() as session:
|
||||
return list(session.exec(select(KnownBrand).order_by(KnownBrand.canonical_name)).all())
|
||||
|
||||
|
||||
def update_brand(brand_id: UUID, **fields):
|
||||
from admin.mpr.media_assets.models import KnownBrand
|
||||
KnownBrand.objects.filter(id=brand_id).update(**fields)
|
||||
def update_brand(brand_id: UUID, **fields) -> None:
|
||||
with get_session() as session:
|
||||
brand = session.get(KnownBrand, brand_id)
|
||||
if not brand:
|
||||
return
|
||||
for k, v in fields.items():
|
||||
setattr(brand, k, v)
|
||||
session.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SourceBrandSighting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def get_source_sightings(source_asset_id: UUID) -> list:
|
||||
"""Get all brand sightings for a specific source video."""
|
||||
from admin.mpr.media_assets.models import SourceBrandSighting
|
||||
return list(
|
||||
SourceBrandSighting.objects
|
||||
.filter(source_asset_id=source_asset_id)
|
||||
.order_by("-occurrences")
|
||||
)
|
||||
def get_source_sightings(source_asset_id: UUID) -> list[SourceBrandSighting]:
|
||||
with get_session() as session:
|
||||
stmt = (
|
||||
select(SourceBrandSighting)
|
||||
.where(SourceBrandSighting.source_asset_id == source_asset_id)
|
||||
.order_by(SourceBrandSighting.occurrences.desc())
|
||||
)
|
||||
return list(session.exec(stmt).all())
|
||||
|
||||
|
||||
def record_sighting(source_asset_id: UUID, brand_id: UUID, brand_name: str,
|
||||
timestamp: float, confidence: float, source: str = "ocr"):
|
||||
"""Record or update a brand sighting for a source."""
|
||||
from admin.mpr.media_assets.models import SourceBrandSighting
|
||||
import uuid
|
||||
timestamp: float, confidence: float, source: str = "ocr") -> SourceBrandSighting:
|
||||
with get_session() as session:
|
||||
stmt = select(SourceBrandSighting).where(
|
||||
SourceBrandSighting.source_asset_id == source_asset_id,
|
||||
SourceBrandSighting.brand_id == brand_id,
|
||||
)
|
||||
sighting = session.exec(stmt).first()
|
||||
|
||||
sighting = SourceBrandSighting.objects.filter(
|
||||
source_asset_id=source_asset_id,
|
||||
brand_id=brand_id,
|
||||
).first()
|
||||
if sighting:
|
||||
total_conf = sighting.avg_confidence * sighting.occurrences + confidence
|
||||
sighting.occurrences += 1
|
||||
sighting.last_seen_timestamp = timestamp
|
||||
sighting.avg_confidence = total_conf / sighting.occurrences
|
||||
session.commit()
|
||||
session.refresh(sighting)
|
||||
return sighting
|
||||
|
||||
if sighting:
|
||||
sighting.occurrences += 1
|
||||
sighting.last_seen_timestamp = timestamp
|
||||
total_conf = sighting.avg_confidence * (sighting.occurrences - 1) + confidence
|
||||
sighting.avg_confidence = total_conf / sighting.occurrences
|
||||
sighting.save()
|
||||
sighting = SourceBrandSighting(
|
||||
source_asset_id=source_asset_id,
|
||||
brand_id=brand_id,
|
||||
brand_name=brand_name,
|
||||
first_seen_timestamp=timestamp,
|
||||
last_seen_timestamp=timestamp,
|
||||
occurrences=1,
|
||||
detection_source=source,
|
||||
avg_confidence=confidence,
|
||||
)
|
||||
session.add(sighting)
|
||||
session.commit()
|
||||
session.refresh(sighting)
|
||||
return sighting
|
||||
|
||||
sighting = SourceBrandSighting.objects.create(
|
||||
id=uuid.uuid4(),
|
||||
source_asset_id=source_asset_id,
|
||||
brand_id=brand_id,
|
||||
brand_name=brand_name,
|
||||
first_seen_timestamp=timestamp,
|
||||
last_seen_timestamp=timestamp,
|
||||
occurrences=1,
|
||||
detection_source=source,
|
||||
avg_confidence=confidence,
|
||||
)
|
||||
return sighting
|
||||
|
||||
Reference in New Issue
Block a user