aaredb/backend/app/models.py
GotthardG 1606e80f81 Add Image models and clean up test code structure
Introduced `ImageCreate` and `Image` models to handle image-related data in the backend. Improved the organization and readability of the testing notebook by consolidating and formatting code into distinct sections with markdown cells.
2025-02-26 13:33:23 +01:00

284 lines
10 KiB
Python

from sqlalchemy import (
Column,
Integer,
String,
Date,
ForeignKey,
JSON,
DateTime,
Boolean,
)
from sqlalchemy.orm import relationship
from .database import Base
from datetime import datetime
class Shipment(Base):
__tablename__ = "shipments"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
pgroups = Column(String(255), nullable=False)
shipment_name = Column(String(255), index=True)
shipment_date = Column(Date, nullable=True)
shipment_status = Column(String(255), nullable=True)
comments = Column(String(200), nullable=True)
contact_id = Column(Integer, ForeignKey("contacts.id"), nullable=False)
return_address_id = Column(Integer, ForeignKey("addresses.id"), nullable=False)
proposal_id = Column(Integer, ForeignKey("proposals.id"), nullable=True)
contact = relationship("Contact", back_populates="shipments")
return_address = relationship("Address", back_populates="shipments")
proposal = relationship("Proposal", back_populates="shipments")
dewars = relationship("Dewar", back_populates="shipment")
class Contact(Base):
__tablename__ = "contacts"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
status = Column(String(255), default="active")
pgroups = Column(String(255), nullable=False)
firstname = Column(String(255), nullable=False)
lastname = Column(String(255))
phone_number = Column(String(255))
email = Column(String(255))
shipments = relationship("Shipment", back_populates="contact")
class LocalContact(Base):
__tablename__ = "local_contacts"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
status = Column(String(255), default="active")
firstname = Column(String(255), nullable=False)
lastname = Column(String(255), nullable=False)
phone_number = Column(String(255), nullable=False)
email = Column(String(255), nullable=False)
class Address(Base):
__tablename__ = "addresses"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
status = Column(String(255), default="active")
pgroups = Column(String(255), nullable=False)
street = Column(String(255), nullable=False)
house_number = Column(String(255), nullable=True)
city = Column(String(255), nullable=False)
state = Column(String(255), nullable=True)
zipcode = Column(String(255), nullable=False)
country = Column(String(255), nullable=False)
shipments = relationship("Shipment", back_populates="return_address")
class DewarType(Base):
__tablename__ = "dewar_types"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
dewar_type = Column(String(255), unique=True, index=True)
serial_numbers = relationship("DewarSerialNumber", back_populates="dewar_type")
class DewarSerialNumber(Base):
__tablename__ = "dewar_serial_numbers"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
serial_number = Column(String(255), index=True)
dewar_type_id = Column(Integer, ForeignKey("dewar_types.id"))
dewar_type = relationship("DewarType", back_populates="serial_numbers")
class Dewar(Base):
__tablename__ = "dewars"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
pgroups = Column(String(255), nullable=False)
dewar_name = Column(String(255), nullable=False)
dewar_type_id = Column(Integer, ForeignKey("dewar_types.id"), nullable=True)
dewar_serial_number_id = Column(
Integer, ForeignKey("dewar_serial_numbers.id"), nullable=True
)
status = Column(String(255), nullable=True)
unique_id = Column(String(255), unique=True, index=True, nullable=True)
tracking_number = Column(String(255), nullable=True)
shipment_id = Column(Integer, ForeignKey("shipments.id"))
return_address_id = Column(Integer, ForeignKey("addresses.id"))
contact_id = Column(Integer, ForeignKey("contacts.id"))
shipment = relationship("Shipment", back_populates="dewars")
return_address = relationship("Address")
contact = relationship("Contact")
pucks = relationship("Puck", back_populates="dewar")
dewar_type = relationship("DewarType")
dewar_serial_number = relationship("DewarSerialNumber")
slot = relationship("Slot", back_populates="dewar")
events = relationship("LogisticsEvent", back_populates="dewar")
beamline_location = None
local_contact_id = Column(Integer, ForeignKey("local_contacts.id"), nullable=True)
local_contact = relationship("LocalContact")
beamtime = relationship("Beamtime", back_populates="dewars")
beamtime_id = Column(Integer, ForeignKey("beamtimes.id"), nullable=True)
@property
def number_of_pucks(self) -> int:
return len(self.pucks) if self.pucks else 0
@property
def number_of_samples(self) -> int:
if not self.pucks:
return 0
return sum(len(puck.samples) for puck in self.pucks)
class Proposal(Base):
__tablename__ = "proposals"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
number = Column(String(255))
shipments = relationship("Shipment", back_populates="proposal")
class Puck(Base):
__tablename__ = "pucks"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
puck_name = Column(String(255), index=True)
puck_type = Column(String(255))
puck_location_in_dewar = Column(Integer)
# Foreign keys and relationships
dewar_id = Column(Integer, ForeignKey("dewars.id"))
dewar = relationship("Dewar", back_populates="pucks")
samples = relationship("Sample", back_populates="puck")
events = relationship("PuckEvent", back_populates="puck")
class Sample(Base):
__tablename__ = "samples"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
sample_name = Column(String(255), index=True)
proteinname = Column(String(255), index=True)
position = Column(Integer)
priority = Column(Integer)
comments = Column(String(255))
data_collection_parameters = Column(JSON, nullable=True)
# Foreign keys and relationships
dewar_id = Column(Integer, ForeignKey("dewars.id"))
puck_id = Column(Integer, ForeignKey("pucks.id"))
puck = relationship("Puck", back_populates="samples")
events = relationship("SampleEvent", back_populates="sample", lazy="joined")
@property
def mount_count(self) -> int:
# Dynamically calculate mount_count
return len([event for event in self.events if event.event_type == "Mounted"])
@property
def unmount_count(self) -> int:
# Dynamically calculate unmount_count
return len([event for event in self.events if event.event_type == "Unmounted"])
class Slot(Base):
__tablename__ = "slots"
id = Column(Integer, primary_key=True, index=True)
qr_code = Column(String(255), unique=True, index=True)
label = Column(String(255))
qr_base = Column(String(255), nullable=True)
occupied = Column(Boolean, default=False)
needs_refill = Column(Boolean, default=False)
dewar_unique_id = Column(String(255), ForeignKey("dewars.unique_id"), nullable=True)
dewar = relationship("Dewar", back_populates="slot")
events = relationship("LogisticsEvent", back_populates="slot")
class LogisticsEvent(Base):
__tablename__ = "logistics_events"
id = Column(Integer, primary_key=True, index=True)
dewar_id = Column(Integer, ForeignKey("dewars.id"))
slot_id = Column(Integer, ForeignKey("slots.id"))
event_type = Column(String(255), index=True)
timestamp = Column(DateTime, default=datetime.now)
dewar = relationship("Dewar", back_populates="events")
slot = relationship("Slot", back_populates="events")
class SampleEvent(Base):
__tablename__ = "sample_events"
id = Column(Integer, primary_key=True, index=True)
sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
event_type = Column(String(255), nullable=False)
timestamp = Column(DateTime, default=datetime.now)
sample = relationship("Sample", back_populates="events")
class PuckEvent(Base):
__tablename__ = "puck_events"
id = Column(Integer, primary_key=True, index=True)
puck_id = Column(Integer, ForeignKey("pucks.id"))
tell = Column(String(255), nullable=True)
tell_position = Column(String(255), nullable=True)
event_type = Column(String(255), index=True)
timestamp = Column(DateTime, default=datetime.now)
puck = relationship("Puck", back_populates="events")
class Beamtime(Base):
__tablename__ = "beamtimes"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
pgroups = Column(String(255), nullable=False)
beamtime_name = Column(String(255), index=True)
beamline = Column(String(255), nullable=True)
start_date = Column(Date, nullable=True)
end_date = Column(Date, nullable=True)
status = Column(String(255), nullable=True)
comments = Column(String(200), nullable=True)
proposal_id = Column(Integer, ForeignKey("proposals.id"), nullable=True)
local_contact_id = Column(Integer, ForeignKey("local_contacts.id"), nullable=False)
local_contact = relationship("LocalContact")
dewars = relationship("Dewar", back_populates="beamtime")
class Image(Base):
__tablename__ = "images"
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
pgroup = Column(String(255), nullable=False)
comment = Column(String(200), nullable=True)
filepath = Column(String(255), nullable=False)
status = Column(String(255), nullable=True)
sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
# class Results(Base):
# __tablename__ = "results"
#
# id = Column(Integer, primary_key=True, index=True, autoincrement=True)
# pgroup = Column(String(255), nullable=False)
# sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
# method = Column(String(255), nullable=False)
# #resolution: Column(Float(255), nullable=False)
# unit_cell: str
# spacegroup: str
# rmerge: float
# rmeas: float
# isig: float
# cc: float
# cchalf: float
# completeness: float
# multiplicity: float
# nobs: int
# total_refl: int
# unique_refl: int
# #comments: Optional[constr(max_length=200)] = None