|
| 1 | +"""S3-backed storage for allowed acquisition types and scheduled acquisitions. |
| 2 | +
|
| 3 | +Data is stored as JSON objects in S3 (bucket: ``aind-scratch-data``, prefix: |
| 4 | +``aind-metadata-viz-data/``). |
| 5 | +
|
| 6 | +Object layout:: |
| 7 | +
|
| 8 | + aind-metadata-viz-data/allowed_acquisition_types.json |
| 9 | + aind-metadata-viz-data/scheduled_acquisitions.json |
| 10 | +
|
| 11 | +* ``add_acquisition_type`` appends a new (platform, acquisition_type) pair. |
| 12 | +* ``get_allowed_types`` returns all allowed (platform, acquisition_type) pairs. |
| 13 | +* ``add_scheduled_acquisition`` validates acquisition_type, generates a uuid, |
| 14 | + and stores the record. |
| 15 | +* ``get_scheduled_acquisitions`` returns all scheduled acquisitions, optionally |
| 16 | + filtered to future-or-today only. |
| 17 | +* ``get_scheduled_acquisition`` returns a single record by uuid. |
| 18 | +""" |
| 19 | + |
| 20 | +import json |
| 21 | +import uuid |
| 22 | +from datetime import date |
| 23 | +from typing import List, Optional |
| 24 | + |
| 25 | +import boto3 |
| 26 | +from botocore.exceptions import ClientError |
| 27 | + |
| 28 | +from .models import ALLOWED_PLATFORMS |
| 29 | + |
| 30 | +_S3_BUCKET = "aind-scratch-data" |
| 31 | +_S3_PREFIX = "aind-metadata-viz-data" |
| 32 | + |
| 33 | +_ALLOWED_TYPES_KEY = f"{_S3_PREFIX}/allowed_acquisition_types.json" |
| 34 | +_SCHEDULED_ACQUISITIONS_KEY = f"{_S3_PREFIX}/scheduled_acquisitions.json" |
| 35 | + |
| 36 | + |
| 37 | +def _s3(): |
| 38 | + return boto3.client("s3") |
| 39 | + |
| 40 | + |
| 41 | +def _put_json(key: str, obj) -> None: |
| 42 | + _s3().put_object( |
| 43 | + Bucket=_S3_BUCKET, |
| 44 | + Key=key, |
| 45 | + Body=json.dumps(obj).encode(), |
| 46 | + ContentType="application/json", |
| 47 | + ) |
| 48 | + |
| 49 | + |
| 50 | +def _get_json(key: str): |
| 51 | + try: |
| 52 | + response = _s3().get_object(Bucket=_S3_BUCKET, Key=key) |
| 53 | + return json.loads(response["Body"].read().decode()) |
| 54 | + except ClientError as exc: |
| 55 | + if exc.response["Error"]["Code"] in ("NoSuchKey", "404"): |
| 56 | + return None |
| 57 | + raise |
| 58 | + |
| 59 | + |
| 60 | +def get_allowed_types() -> List[dict]: |
| 61 | + """Return all allowed (platform, acquisition_type) pairs.""" |
| 62 | + return _get_json(_ALLOWED_TYPES_KEY) or [] |
| 63 | + |
| 64 | + |
| 65 | +def add_acquisition_type(platform: str, acquisition_type: str) -> dict: |
| 66 | + """Add a new allowed (platform, acquisition_type) pair. |
| 67 | +
|
| 68 | + Raises ``ValueError`` if ``platform`` is set but not in ``ALLOWED_PLATFORMS`` |
| 69 | + (only enforced once that list has been populated). Deduplicates on the |
| 70 | + exact (platform, acquisition_type) pair. |
| 71 | + """ |
| 72 | + if ALLOWED_PLATFORMS and platform not in ALLOWED_PLATFORMS: |
| 73 | + raise ValueError(f"platform '{platform}' is not one of the allowed platforms: {ALLOWED_PLATFORMS}") |
| 74 | + |
| 75 | + entries = get_allowed_types() |
| 76 | + entry = {"platform": platform, "acquisition_type": acquisition_type} |
| 77 | + if entry not in entries: |
| 78 | + entries.append(entry) |
| 79 | + _put_json(_ALLOWED_TYPES_KEY, entries) |
| 80 | + return entry |
| 81 | + |
| 82 | + |
| 83 | +def _find_platform_for_type(acquisition_type: str) -> Optional[str]: |
| 84 | + """Return the platform registered for ``acquisition_type``, or None if not found.""" |
| 85 | + for entry in get_allowed_types(): |
| 86 | + if entry["acquisition_type"] == acquisition_type: |
| 87 | + return entry["platform"] |
| 88 | + return None |
| 89 | + |
| 90 | + |
| 91 | +def add_scheduled_acquisition(subject_id: str, acquisition_date: date, acquisition_type: str) -> str: |
| 92 | + """Validate ``acquisition_type`` against the allowed types and store a new scheduled acquisition. |
| 93 | +
|
| 94 | + Returns the generated uuid for the new record. Raises ``ValueError`` if |
| 95 | + ``acquisition_type`` is not in the allowed types list. |
| 96 | + """ |
| 97 | + platform = _find_platform_for_type(acquisition_type) |
| 98 | + if platform is None: |
| 99 | + raise ValueError(f"acquisition_type '{acquisition_type}' is not an allowed acquisition type") |
| 100 | + |
| 101 | + acquisition_uuid = str(uuid.uuid4()) |
| 102 | + records = _get_json(_SCHEDULED_ACQUISITIONS_KEY) or {} |
| 103 | + records[acquisition_uuid] = { |
| 104 | + "subject_id": subject_id, |
| 105 | + "date": acquisition_date.isoformat() if isinstance(acquisition_date, date) else acquisition_date, |
| 106 | + "acquisition_type": acquisition_type, |
| 107 | + "platform": platform, |
| 108 | + } |
| 109 | + _put_json(_SCHEDULED_ACQUISITIONS_KEY, records) |
| 110 | + return acquisition_uuid |
| 111 | + |
| 112 | + |
| 113 | +def get_scheduled_acquisitions(include_past: bool = False) -> List[dict]: |
| 114 | + """Return all scheduled acquisitions, each including its ``uuid``. |
| 115 | +
|
| 116 | + If ``include_past`` is False (default), only acquisitions scheduled for |
| 117 | + today or later are returned. |
| 118 | + """ |
| 119 | + records = _get_json(_SCHEDULED_ACQUISITIONS_KEY) or {} |
| 120 | + today = date.today() |
| 121 | + results = [] |
| 122 | + for acquisition_uuid, record in records.items(): |
| 123 | + if not include_past and date.fromisoformat(record["date"]) < today: |
| 124 | + continue |
| 125 | + results.append({"uuid": acquisition_uuid, **record}) |
| 126 | + return results |
| 127 | + |
| 128 | + |
| 129 | +def get_scheduled_acquisition(acquisition_uuid: str) -> Optional[dict]: |
| 130 | + """Return a single scheduled acquisition record by uuid, or None if not found.""" |
| 131 | + records = _get_json(_SCHEDULED_ACQUISITIONS_KEY) or {} |
| 132 | + return records.get(acquisition_uuid) |
0 commit comments