from database import db

watermarks_collection = db["watermarks"]


def get_watermark(source: str) -> str | None:
    """Return the ISO timestamp of the last successful fetch for a source."""
    doc = watermarks_collection.find_one({"source": source})
    return doc.get("last_fetch_time") if doc else None


def set_watermark(source: str, timestamp: str):
    """Persist the latest successful fetch timestamp for a source."""
    watermarks_collection.update_one(
        {"source": source},
        {"$set": {"last_fetch_time": timestamp}},
        upsert=True,
    )