Lab Accession

1"""Lab Accession API built with Pixeltable.
2
3    pxt schema update app.py lab
4    pxt service run app.py lab
5"""
6import pixeltable as pxt
7import pixeltable.functions as pxtf
8from pixeltable.serving import FastAPIRouter
9
10from udfs import priority_name, station_label, urgency_band
11
12# ---- tables ----
13TableModel = pxt.model_base()
14
15
16class Stations(TableModel, name='stations'):
17    station_id = pxt.Column(type=pxt.String, primary_key=True)
18    name: pxt.String
19    capacity: pxt.Int
20
21    label = station_label(name, capacity)
22
23
24class Samples(TableModel, name='samples'):
25    accession_id = pxt.Column(type=pxt.String, primary_key=True)
26    specimen: pxt.String
27    priority: pxt.Int
28    received_at: pxt.String
29    batch_code: pxt.String | None
30
31    priority_label = priority_name(priority)
32
33
34class Assignments(TableModel, name='assignments', has_default_idxs=False):
35    id = pxt.Column(value=pxtf.uuid.uuid7(), primary_key=True)
36    accession_id: pxt.String
37    station_id: pxt.String
38    priority: pxt.Int
39    delayed_min: pxt.Int
40    status: pxt.String
41    claimed_by: pxt.String | None
42
43    urgency = urgency_band(priority, delayed_min)
44
45    __indexes__ = [pxt.BtreeIndex(station_id), pxt.BtreeIndex(status)]
46
47
48# ---- queries ----
49@pxt.query
50def station_queue(station_id: str):
51    """Work for one station, most urgent first."""
52    return Assignments.where(Assignments.station_id == station_id).select(
53        Assignments.id, Assignments.accession_id, Assignments.urgency, Assignments.status, Assignments.claimed_by
54    ).order_by(Assignments.priority, asc=False)
55
56
57@pxt.query
58def batch(batch_code: str):
59    """Samples received in one batch."""
60    return Samples.where(Samples.batch_code == batch_code).select(
61        Samples.accession_id, Samples.specimen, Samples.priority_label
62    ).order_by(Samples.accession_id)
63
64
65# ---- routes ----
66lab_api = FastAPIRouter(name='lab_api')
67lab_api.add_insert_route(Stations, path='/stations', inputs=[Stations.station_id, Stations.name, Stations.capacity],
68                         outputs=[Stations.station_id, Stations.label])
69lab_api.add_insert_route(Samples, path='/samples',
70                         inputs=[Samples.accession_id, Samples.specimen, Samples.priority, Samples.received_at,
71                                 Samples.batch_code],
72                         outputs=[Samples.accession_id, Samples.priority_label])
73lab_api.add_insert_route(Assignments, path='/assignments',
74                         inputs=[Assignments.accession_id, Assignments.station_id, Assignments.priority,
75                                 Assignments.delayed_min, Assignments.status, Assignments.claimed_by],
76                         outputs=[Assignments.id, Assignments.urgency])
77lab_api.add_update_route(Assignments, path='/assignments/claim', inputs=[Assignments.status, Assignments.claimed_by],
78                         outputs=[Assignments.id, Assignments.status, Assignments.claimed_by])
79lab_api.add_compute_route(Assignments, path='/urgency', inputs=[Assignments.priority, Assignments.delayed_min],
80                          outputs=[Assignments.urgency])
81lab_api.add_query_route(path='/stations/queue', query=station_queue, method='get')
82lab_api.add_query_route(path='/batches', query=batch, method='get')