From 5f6e09102f828cce522106ab67b8638c8a648dc6 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 25 Jun 2025 13:35:30 +0200 Subject: [PATCH] DAQ: SSE work in progress --- daq/src/aaredaq/aaredb.py | 31 ---- daq/src/aaredaq/daq.py | 2 +- daq/src/aaredaq/server.py | 40 +++++- gui/src/aaregui/threads/sse_client.py | 199 ++++++++++++++++++++++++++ 4 files changed, 234 insertions(+), 38 deletions(-) create mode 100644 gui/src/aaregui/threads/sse_client.py diff --git a/daq/src/aaredaq/aaredb.py b/daq/src/aaredaq/aaredb.py index 0af75885..848cb1eb 100644 --- a/daq/src/aaredaq/aaredb.py +++ b/daq/src/aaredaq/aaredb.py @@ -67,37 +67,6 @@ class AareWrapper: ) print(ret) - def get_sample_info(self) -> SampleShortInfoList: - sample_list = ( - self.__tell_api.get_pucks_with_tell_position(tell=self.__bl.value.upper()) - ) - - ret = [] - for p in sample_list: - for s in p.samples: - dewar_address = None - if p.tell_position is not None and len(p.tell_position) == 2: - dewar_address = DewarAddress( - segment=p.tell_position[0], pos=p.tell_position[1] - ) - - ret.append( - SampleShortInfo( - db_id=s.id, - puck_name=p.puck_name, - dewar_name="Unknown" if p.dewar_name is None else p.dewar_name, - sample_name=s.sample_name, - user=p.pgroup, - pin=s.position, - location=dewar_address, - priority=1.0 if s.priority is None else s.priority, - comment=s.comments, - mount_count=s.mount_count or 0, - aaredb_params=s.data_collection_parameters - ) - ) - return SampleShortInfoList(s=ret) - def sample_mounted(self, s: Optional[SampleShortInfo]): if s is not None: try: diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index 2638e53d..1f4abe9c 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -607,7 +607,7 @@ class AareDAQ: @property def sample_spreadsheet(self) -> SampleShortInfoList: - return self.__aare.get_sample_info() + return self.__cfg.spreadsheet def sample_spreadsheet_user(self, pgroup: str) -> SampleShortInfoList: sample = self.sample_spreadsheet diff --git a/daq/src/aaredaq/server.py b/daq/src/aaredaq/server.py index c9d4f1b1..97097d88 100644 --- a/daq/src/aaredaq/server.py +++ b/daq/src/aaredaq/server.py @@ -1,5 +1,6 @@ +import asyncio import io -from typing import Tuple +from typing import Tuple, AsyncGenerator import cv2 import urllib3 @@ -210,13 +211,40 @@ async def unmount(token: str = Depends(oauth2_scheme)): return "OK" -@app.get("/tell/spreadsheet") -async def spreadsheet(token: str = Depends(oauth2_scheme)): - data = auth.parse_token(token) +def get_spreadsheet(data: TokenData) -> SampleShortInfoList: if data.staff: - return daq.sample_spreadsheet + return cfg.spreadsheet else: - return daq.sample_spreadsheet_user(data.group) + return cfg.spreadsheet_pgroup(data.group) + + +async def spreadsheet_event_stream(data: TokenData) -> AsyncGenerator[str, None]: + try: + while True: + yield get_spreadsheet(data).model_dump_json() + await asyncio.sleep(10) + except asyncio.CancelledError: + return + + +@app.get("/sse/spreadsheet") +async def spreadsheet_sse(token: str = Depends(oauth2_scheme)): + data = auth.parse_token(token) + return StreamingResponse( + spreadsheet_event_stream(data), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Headers": "Cache-Control" + } + ) + + +@app.get("/tell/spreadsheet") +async def spreadsheet(token: str = Depends(oauth2_scheme)) -> SampleShortInfoList: + return get_spreadsheet(auth.parse_token(token)) # State transitions diff --git a/gui/src/aaregui/threads/sse_client.py b/gui/src/aaregui/threads/sse_client.py new file mode 100644 index 00000000..50d9124f --- /dev/null +++ b/gui/src/aaregui/threads/sse_client.py @@ -0,0 +1,199 @@ +from PySide6.QtCore import QObject, Signal, Slot + +from aaredaqlib.models import SampleShortInfoList + + +class SSEClient(QObject): + # Signals + connected = Signal() + disconnected = Signal() + message_received = Signal(str) # Raw SSE data + event_received = Signal(str, str) # event_type, data + error_occurred = Signal(str) + + def __init__(self, parent=None): + super().__init__(parent) + self._network_manager = QNetworkAccessManager(self) + self._reply: Optional[QNetworkReply] = None + self._reconnect_timer = QTimer(self) + self._reconnect_timer.setSingleShot(True) + self._reconnect_timer.timeout.connect(self._attempt_reconnect) + + self._url = QUrl() + self._headers: Dict[str, str] = {} + self._buffer = QByteArray() + self._connected = False + self._reconnect_delay = 1000 # Start with 1 second + self._max_reconnect_delay = 30000 # Max 30 seconds + + # Current SSE event data + self._current_event = "" + self._current_data = "" + self._current_id = "" + + def connect_to_sse(self, url: str, headers: Optional[Dict[str, str]] = None): + """Connect to SSE endpoint""" + self._url = QUrl(url) + self._headers = headers or {} + + if self._reply: + self._reply.abort() + self._reply.deleteLater() + self._reply = None + + self._start_connection() + + def disconnect_from_sse(self): + """Disconnect from SSE""" + self._reconnect_timer.stop() + + if self._reply: + self._reply.abort() + self._reply.deleteLater() + self._reply = None + + if self._connected: + self._connected = False + self.disconnected.emit() + + def is_connected(self) -> bool: + """Check if connected to SSE""" + return self._connected and self._reply and self._reply.isOpen() + + def _start_connection(self): + """Start SSE connection""" + request = QNetworkRequest(self._url) + + # Set SSE-specific headers + request.setRawHeader(b"Accept", b"text/event-stream") + request.setRawHeader(b"Cache-Control", b"no-cache") + request.setRawHeader(b"Connection", b"keep-alive") + + # Add custom headers + for key, value in self._headers.items(): + request.setRawHeader(key.encode(), value.encode()) + + print(f"Connecting to SSE: {self._url.toString()}") + + self._reply = self._network_manager.get(request) + self._reply.readyRead.connect(self._on_ready_read) + self._reply.finished.connect(self._on_finished) + self._reply.errorOccurred.connect(self._on_error) + + @Slot() + def _on_ready_read(self): + """Handle incoming SSE data""" + if not self._reply: + return + + # First successful read means we're connected + if not self._connected: + self._connected = True + self._reconnect_delay = 1000 # Reset reconnect delay + self.connected.emit() + + data = self._reply.readAll() + self._buffer.append(data) + self._process_buffer() + + def _process_buffer(self): + """Process SSE data buffer""" + while b'\n' in self._buffer: + newline_index = self._buffer.indexOf(b'\n') + line_data = self._buffer.left(newline_index) + self._buffer.remove(0, newline_index + 1) + + # Remove carriage return if present + if line_data.endsWith(b'\r'): + line_data.chop(1) + + line = line_data.data().decode('utf-8', errors='ignore') + self._parse_sse_line(line) + + def _parse_sse_line(self, line: str): + """Parse individual SSE line""" + if not line: + # Empty line indicates end of event + if self._current_data: + if self._current_event: + self.event_received.emit(self._current_event, self._current_data) + else: + self.message_received.emit(self._current_data) + + # Clear current event data + self._current_event = "" + self._current_data = "" + self._current_id = "" + return + + # Skip comments + if line.startswith(':'): + return + + # Parse field and value + colon_index = line.find(':') + if colon_index == -1: + # Field with no value + field = line + if field == "data": + self._current_data += '\n' + return + + field = line[:colon_index] + value = line[colon_index + 1:] + + # Remove leading space from value + if value.startswith(' '): + value = value[1:] + + if field == "data": + if self._current_data: + self._current_data += '\n' + self._current_data += value + elif field == "event": + self._current_event = value + elif field == "id": + self._current_id = value + elif field == "retry": + try: + retry_time = int(value) + if retry_time > 0: + self._reconnect_delay = retry_time + except ValueError: + pass + + @Slot() + def _on_finished(self): + """Handle connection finished""" + if self._connected: + self._connected = False + self.disconnected.emit() + + if self._reply: + self._reply.deleteLater() + self._reply = None + + # Auto-reconnect after delay + print(f"SSE connection finished, reconnecting in {self._reconnect_delay}ms") + self._reconnect_timer.start(self._reconnect_delay) + + # Exponential backoff + self._reconnect_delay = min(self._reconnect_delay * 2, self._max_reconnect_delay) + + @Slot(QNetworkReply.NetworkError) + def _on_error(self, error: QNetworkReply.NetworkError): + """Handle network error""" + error_string = self._reply.errorString() if self._reply else "Unknown network error" + print(f"SSE Network error: {error} - {error_string}") + + self.error_occurred.emit(error_string) + + if self._connected: + self._connected = False + self.disconnected.emit() + + @Slot() + def _attempt_reconnect(self): + """Attempt to reconnect""" + print("Attempting to reconnect to SSE...") + self._start_connection()