app.bettersight.io/backend/repositories/scrape_repository.py

71 lines
2.4 KiB
Python

import uuid
from repositories.pocketbase_client import pb_client
from repositories.repo_config import USE_MOCK_REPOSITORIES
COLLECTION = 'scrape_runs'
class ScrapeRepository:
"""Data access for the `scrape_runs` collection — also the RQ job_id record."""
def create(self, data):
return pb_client.create(COLLECTION, data)
def update(self, job_id, data):
return pb_client.update(COLLECTION, job_id, data)
def get_by_id(self, job_id):
return pb_client.get_one(COLLECTION, job_id)
def list_recent_for_tenant(self, tenant_id, limit=5):
"""Returns a tenant's most recent runs, newest first — backs GET /research/history."""
return pb_client.list(
COLLECTION, filter_str=f'tenant_id = "{tenant_id}"', sort='-started_at', per_page=limit
)
def get_latest_complete_for_tenant(self, tenant_id):
"""Returns the tenant's most recent completed run, or None — backs GET /research/history/latest."""
items = pb_client.list(
COLLECTION, filter_str=f'tenant_id = "{tenant_id}" && status = "complete"',
sort='-started_at', per_page=1
)
return items[0] if items else None
class MockScrapeRepository:
"""In-memory stand-in for ScrapeRepository."""
def __init__(self):
self._records = {}
def create(self, data):
record_id = data.get('id') or str(uuid.uuid4())
record = {'id': record_id, **data}
self._records[record_id] = record
return record
def update(self, job_id, data):
if job_id not in self._records:
return None
self._records[job_id] = {**self._records[job_id], **data}
return self._records[job_id]
def get_by_id(self, job_id):
return self._records.get(job_id)
def list_recent_for_tenant(self, tenant_id, limit=5):
items = [r for r in self._records.values() if r.get('tenant_id') == tenant_id]
items.sort(key=lambda r: r.get('started_at') or '', reverse=True)
return items[:limit]
def get_latest_complete_for_tenant(self, tenant_id):
items = [
r for r in self._records.values()
if r.get('tenant_id') == tenant_id and r.get('status') == 'complete'
]
items.sort(key=lambda r: r.get('started_at') or '', reverse=True)
return items[0] if items else None
scrape_repository = MockScrapeRepository() if USE_MOCK_REPOSITORIES else ScrapeRepository()