server: adding basic tcp server

Author: Marius Isken

src/qseek/apps/qseek.py

@@ -255,41 +255,28 @@ def main() -> None:
         case "search":
             nest_asyncio.apply()
             from qseek.search import Search
-            from qseek.server import WebServer
 
             search = Search.from_config(args.config)
 
-            webserver = WebServer(search)
-
-            async def run() -> None:
-                http = asyncio.create_task(webserver.start())
-                await search.start(
+            asyncio.run(
+                search.start(
                     force_rundir=args.force,
                     create_backup=not args.no_backup,
-                )
-                await http
-
-            asyncio.run(run(), debug=loop_debug)
+                ),
+                debug=loop_debug,
+            )
 
         case "continue":
             nest_asyncio.apply()
             from qseek.search import Search
-            from qseek.server import WebServer
 
             search = Search.load_rundir(args.rundir)
             if search._progress.time_progress:
                 console.rule(f"Continuing search from {search._progress.time_progress}")
             else:
                 console.rule("Starting search from scratch")
 
-            webserver = WebServer(search)
-
-            async def run() -> None:
-                http = asyncio.create_task(webserver.start())
-                await search.start()
-                await http
-
-            asyncio.run(run(), debug=loop_debug)
+            asyncio.run(search.start(), debug=loop_debug)
 
         case "snuffler":
             from qseek.search import Search
@@ -380,18 +367,6 @@ async def worker() -> None:
 
             start_ui(args.uris, reload=False)
 
-        case "serve":
-            from qseek.search import Search
-            from qseek.server import WebServer
-
-            search = Search.load_rundir(args.rundir)
-            webserver = WebServer(search)
-
-            async def start() -> None:
-                await webserver.start()
-
-            asyncio.run(start(), debug=loop_debug)
-
         case "clear-cache":
             logger.info("clearing cache directory %s", CACHE_DIR)
             shutil.rmtree(CACHE_DIR)

src/qseek/search.py

@@ -38,6 +38,7 @@
 from qseek.pre_processing.frequency_filters import Bandpass
 from qseek.pre_processing.module import Downsample, PreProcessing
 from qseek.reduce import DelaySumReduce
+from qseek.server import WebServer
 from qseek.signals import Signal
 from qseek.stats import RuntimeStats, Stats
 from qseek.tracers.tracers import RayTracer, RayTracers
@@ -284,6 +285,10 @@ class Search(Model):
         default_factory=list,
         description="Event features to extract.",
     )
+    webserver: WebServer | None = Field(
+        default_factory=WebServer,
+        description="Web server for serving search results and monitoring.",
+    )
 
     semblance_sampling_rate: SamplingRate = Field(
         default=100,
@@ -512,6 +517,10 @@ async def prepare(self) -> None:
         for magnitude in self.magnitudes:
             await magnitude.prepare(self.octree, self.stations)
 
+        if self.webserver:
+            await self.webserver.prepare(search=self)
+            BackgroundTasks.create_task(self.webserver.start())
+
         csv_header = EventDetection.csv_header()
         for m in self.magnitudes:
             csv_header.extend(m.csv_header())
@@ -605,6 +614,8 @@ async def start(
         await self._catalog.export_detections(
             jitter_location=self.octree.smallest_node_size()
         )
+        if self.webserver:
+            self.webserver.stop()
         console.cancel()
         logger.info("finished search in %s", datetime_now() - processing_start)
         logger.info("detected %d events", self._catalog.n_events)
@@ -627,6 +638,9 @@ async def new_detections(self, detections: list[EventDetection]) -> None:
             )
             await self._new_detection.emit(detection)
 
+        if self.webserver:
+            BackgroundTasks.create_task(self.webserver.new_detections(detections))
+
         if not catalog.n_events:
             return
 

src/qseek/server.py

@@ -2,100 +2,59 @@
 
 import logging
 from importlib.metadata import version
-from typing import TYPE_CHECKING, AbstractSet, Any, Mapping, Type, TypeVar
-from uuid import UUID
+from ipaddress import IPv4Address
+from typing import TYPE_CHECKING, Literal
+from weakref import WeakSet
 
+import aiofiles
 import aiohttp_cors
 from aiohttp import web
-from pydantic import BaseModel, ValidationError
+from pydantic import Field, PrivateAttr
+
+from qseek.base import Model
 
 if TYPE_CHECKING:
+    from qseek.models.detection import EventDetection
     from qseek.search import Search
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
-MAX_TRIES = 5
 
-TBaseModel = TypeVar("TBaseModel", bound=BaseModel)
-AbstractSetIntStr = AbstractSet[int | str]
-MappingIntStrAny = Mapping[int | str, Any]
-
-
-class DetectionRequest(BaseModel):
-    uid: UUID
-
-
-async def parse_pydantic_request(
-    request: web.Request, model: Type[TBaseModel]
-) -> TBaseModel:
-    """Parse a aiohttp request and return a pydantic model.
-
-    Raises:
-        web.HTTPMethodNotAllowed: If method is not GET/POST/PUT
-        web.HTTPBadRequest: _description_
-        web.HTTPError: _description_
-
-    Returns:
-        TBaseModel: Parsed pydantic model.
-    """
-    try:
-        if request.method == "GET":
-            query = {}
-            for key, value in request.query.items():
-                if key.endswith("[]"):
-                    key = key.strip("[]")
-                    if key not in query:
-                        query[key] = []
-                    query[key].append(value)
-                else:
-                    query[key] = value
-
-            data = model.model_validate(query)
-        elif request.method in ("POST", "PUT"):
-            data = model.model_validate(await request.post())
-        else:
-            raise web.HTTPMethodNotAllowed(
-                request.method, allowed_methods=("GET", "POST", "PUT")
-            )
-    except ValidationError as exc:
-        raise web.HTTPBadRequest(
-            text=exc.json(), content_type="application/json"
-        ) from exc
-    except Exception as exc:
-        raise web.HTTPError(text="Cannot parse request") from exc
-    return data
-
-
-def pydantic_response(
-    model: BaseModel,
-    headers: dict[str, Any] | None = None,
-    exclude: dict | None = None,
-) -> web.Response:
-    try:
-        return web.json_response(
-            text=model.model_dump_json(exclude=exclude).replace("NaN", "null"),
-            headers=headers,
-        )
-    except TypeError as exc:
-        return web.HTTPServerError(reason=str(exc))
+MAX_TRIES = 5
 
 
-class WebServer:
-    def __init__(self, search: Search) -> None:
-        self.search = search
-        self.app = web.Application(logger=logger)
+class WebServer(Model):
+    host: IPv4Address = Field(
+        default=IPv4Address("0.0.0.0"),
+        description="Host to bind the server to. Default is all interfaces.",
+    )
+    port: int | Literal["auto"] = Field(
+        default="auto",
+        description="Port to bind the server to. Default is 'auto', which tries to "
+        "bind to port 19000 and increments if the port is already in use.",
+    )
+
+    _search: Search = PrivateAttr()
+    _app: web.Application = PrivateAttr()
+    _site: web.TCPSite | None = PrivateAttr(default=None)
+    _open_websockets: WeakSet[web.WebSocketResponse] = PrivateAttr(
+        default_factory=WeakSet
+    )
+
+    def prepare(self, search: Search) -> None:
+        self._search = search
+        self._app = web.Application(logger=logger)
         self.add_routes()
 
     def add_routes(self) -> None:
-        router = self.app.router
+        router = self._app.router
         router.add_get("/", self.get_index)
         router.add_get("/api/v1/detections", self.get_detections)
+        router.add_get("/api/v1/receivers", self.get_receivers)
         router.add_get("/api/v1/search", self.get_search)
-        router.add_get(r"/api/v1/detection/{uid}", self.get_detection)
 
         # Configure default CORS settings.
         cors = aiohttp_cors.setup(
-            self.app,
+            self._app,
             defaults={
                 "*": aiohttp_cors.ResourceOptions(
                     allow_credentials=True,
@@ -117,34 +76,80 @@ async def get_index(self, request: web.Request) -> web.Response:
         qseek_version = version("qseek")
         return web.Response(text=f"{qseek_version} HTTP server running...")
 
-    async def get_detections(self, request: web.Request) -> web.Response:
-        return pydantic_response(
-            self.search._catalog,
-            exclude={"detections": {"__all__": {"octree", "stations"}}},
+    async def get_detections(
+        self, request: web.Request, n_detections: int = -1
+    ) -> web.Response:
+        async with aiofiles.open(self._search._catalog.detections_file, "r") as f:
+            lines = await f.readlines()
+        data = "{" + lines + "}"
+        await web.Response(text=data, content_type="application/json")
+
+    async def get_receivers(
+        self, request: web.Request, n_detections: int = -1
+    ) -> web.Response:
+        async with aiofiles.open(self._search._catalog.receivers_file, "r") as f:
+            lines = await f.readlines()
+        data = "{" + lines + "}"
+        await web.Response(text=data, content_type="application/json")
+
+    async def get_search(self, request: web.Request) -> web.Response:
+        return web.Response(
+            self._search.model_dump_json(),
+            content_type="application/json",
         )
 
-    async def get_detection(self, request: web.Request) -> web.Response:
-        req = await parse_pydantic_request(request, model=DetectionRequest)
-        try:
-            detection = self.search._catalog.get(req.uid)
-        except KeyError:
-            return web.HTTPNotFound()
-        return pydantic_response(detection)
+    async def websocket_handler(self, request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        self._open_websockets.add(ws)
+        logger.info("New WebSocket connection established: %s", request.remote)
+
+        for msg in ws:
+            if msg.type == web.WSMsgType.TEXT:
+                # Handle incoming messages from the client if needed
+                pass
+            elif msg.type == web.WSMsgType.ERROR:
+                logger.error(
+                    "WebSocket connection closed with exception %s", ws.exception()
+                )
 
-    async def get_search(self, request: web.Request) -> web.Response:
-        return pydantic_response(self.search)
+    async def new_detections(self, detections: list[EventDetection]) -> None:
+        if not self._open_websockets:
+            return
+
+        data = {
+            "type": "new_detections",
+            "detections": [d.model_dump(exclude={"receivers"}) for d in detections],
+            "receivers": [d.receivers.model_dump() for d in detections],
+        }
+
+        for ws in self._open_websockets:
+            await ws.send_json(data)
 
     async def start(self) -> None:
-        runner = web.AppRunner(self.app)
+        runner = web.AppRunner(self._app)
         await runner.setup()
 
-        port = 3445
+        port = 19000 if self.port == "auto" else self.port
 
         for _ in range(MAX_TRIES):
             try:
-                site = web.TCPSite(runner, "0.0.0.0", port)
+                site = web.TCPSite(runner, str(self.host), port)
                 await site.start()
-                logger.info("started on http://0.0.0.0:%d", port)
+                logger.info("started on http://%s:%d", self.host, port)
+                self._site = site
                 break
             except OSError:
+                if self.port != "auto":
+                    raise
+                logger.warning("port %d is in use, trying next port...", port)
                 port += 1
+        else:
+            raise OSError(f"Could not bind to a port after {MAX_TRIES} tries")
+
+    async def stop(self) -> None:
+        logger.info("shutting down HTTP server...")
+        if self._site:
+            await self._site.stop()
+        await self._app.shutdown()
+        await self._app.cleanup()