diff --git a/db_lib.py b/db_lib.py index 85181422..a0321618 100755 --- a/db_lib.py +++ b/db_lib.py @@ -1,16 +1,15 @@ +import logging import os - import time - -import six - import uuid +from collections import defaultdict import amostra.client.commands as acc import conftrak.client.commands as ccc -from analysisstore.client.commands import AnalysisClient import conftrak.exceptions -import logging +import six +from analysisstore.client.commands import AnalysisClient + logger = logging.getLogger(__name__) #12/19 - Skinner inherited this from Hugo, who inherited it from Matt. Arman wrote the underlying DB and left BNL in 2018. @@ -486,6 +485,43 @@ def getContainerByID(container_id): c = getContainers(filters={'uid': container_id})[0] return c +def get_dewar_tree_data(dewar_name, beamline): + """ + returns all data required to show dewar tree data with minimum number of database accesses + """ + dewar_data = getContainers(filters={"name": dewar_name, "owner": beamline})[0] + + puck_ids = [ + pid for pid in dewar_data.get("content", []) if pid + ] # removes blank ids + pucks = getContainers(filters={"uid": {"$in": puck_ids}}) + + # Create a mega list of sample ids from puck information + sample_ids = [ + sample_id + for puck in pucks + for sample_id in puck.get("content", []) + if sample_id + ] + + # Get all sample info in one call + params = {"uid": {"$in": sample_ids}} + samples = sample_ref.find(as_document=False, **params) + + # Get all request info in one call + params = {"sample": {"$in": sample_ids}, "state": "active"} + reqs = list(request_ref.find(**params)) + + # Assemble data into dictionaries + puck_data = {puck["uid"]: puck for puck in pucks} + sample_data = {sample["uid"]: sample for sample in samples} + request_data = defaultdict(list) + for req in reqs: + request_data[req["sample"]].append(req) + + return dewar_data, puck_data, sample_data, request_data + + def getQueue(beamlineName): """ diff --git a/gui/dewar_tree.py b/gui/dewar_tree.py index 355dbc53..02fbc03c 100644 --- a/gui/dewar_tree.py +++ b/gui/dewar_tree.py @@ -1,16 +1,19 @@ +import getpass import logging +import os import typing -import os, getpass + +import requests from qtpy import QtCore, QtGui, QtWidgets from qtpy.QtCore import Qt -import requests + import daq_utils import db_lib from config_params import ( DEWAR_SECTORS, + IS_STAFF, PUCKS_PER_DEWAR_SECTOR, SAMPLE_TIMER_DELAY, - IS_STAFF, ) if typing.TYPE_CHECKING: @@ -109,20 +112,17 @@ def set_mounted_sample(self, item): def refreshTreeDewarView(self): puck = "" self.model.clear() - dewarContents = db_lib.getContainerByName( + dewar_data, puck_data, sample_data, request_data = db_lib.get_dewar_tree_data( daq_utils.primaryDewarName, daq_utils.beamline - )["content"] + ) parentItem = self.model.invisibleRootItem() for i, puck_id in enumerate( - dewarContents + dewar_data["content"] ): # dewar contents is the list of puck IDs puck = "" puckName = "" if puck_id: - puck = containerDict.get( - puck_id, - containerDict.setdefault(puck_id, db_lib.getContainerByID(puck_id)), - ) + puck = puck_data[puck_id] puckName = puck["name"] sector, puck_pos = divmod(i, self.pucksPerDewarSector) index_s = f"{sector+1}{chr(puck_pos + ord('A'))}" @@ -132,7 +132,9 @@ def refreshTreeDewarView(self): parentItem.appendRow(item) if puck != "" and puckName != "private": puckContents = puck.get("content", []) - self.add_samples_to_puck_tree(puckContents, item, index_s) + self.add_samples_to_puck_tree( + puckContents, item, index_s, sample_data, request_data + ) self.setModel(self.model) self.model.itemChanged.connect(self.queueSelectedSample) if self.isExpanded: @@ -142,7 +144,12 @@ def refreshTreeDewarView(self): self.scrollTo(self.currentIndex(), QtWidgets.QAbstractItemView.PositionAtCenter) def add_samples_to_puck_tree( - self, puckContents, parentItem: QtGui.QStandardItem, index_label + self, + puckContents, + parentItem: QtGui.QStandardItem, + index_label, + sample_data, + request_data, ): # Method will attempt to add samples to the puck. If you don't belong to the proposal, # it will not add samples and clear the puck information @@ -159,10 +166,7 @@ def add_samples_to_puck_tree( parentItem.appendRow(item) continue - sample = sampleDict.get( - sample_id, - sampleDict.setdefault(sample_id, db_lib.getSampleByID(sample_id)), - ) + sample = sample_data[sample_id] if not IS_STAFF and not self.is_proposal_member(sample["proposalID"]): # If the user is not part of the proposal and is not staff, don't fill tree @@ -196,7 +200,7 @@ def add_samples_to_puck_tree( if sample_id == self.parent.selectedSampleID: logger.info("found " + str(self.parent.SelectedItemData)) selectedSampleIndex = self.model.indexFromItem(item) - sampleRequestList = db_lib.getRequestsBySampleID(sample_id) + sampleRequestList = request_data[sample_id] for request in sampleRequestList: if not ("protocol" in request["request_obj"]): continue