feat: stream on-demand inventory queries

src/archive_clients/daemon.py

@@ -13,6 +13,7 @@ from websockets.asyncio.client import connect
 
 from archive_clients.backup import SQLiteBackupManager
 from archive_clients.config import ClientConfig
+from archive_clients.inventory import InventoryService
 from archive_clients.locking import DatabaseLease
 from archive_clients.probes import FilesystemProbe
 from archive_clients.protocol import (
@@ -22,9 +23,12 @@ from archive_clients.protocol import (
     encode_message,
     new_envelope,
 )
+from archive_clients.qbittorrent import QBittorrentReader
 from archive_clients.services import ServiceProbe
 from archive_clients.state import ClientStore, CommandConflict
-from archive_control.v1 import client_pb2, common_pb2, control_pb2, job_pb2
+from archive_control.v1 import (
+    client_pb2, common_pb2, control_pb2, inventory_pb2, job_pb2,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -36,6 +40,7 @@ class ArchiveClientDaemon:
         config: ClientConfig,
         probes: list[FilesystemProbe],
         service_probes: list[ServiceProbe],
+        resource_reader: QBittorrentReader | None = None,
     ):
         if len(probes) != 2:
             raise ValueError(
@@ -44,6 +49,10 @@ class ArchiveClientDaemon:
         self.config = config
         self.probes = probes
         self.service_probes = service_probes
+        self.inventory = (
+            InventoryService(resource_reader, config.client_id)
+            if resource_reader is not None else None
+        )
         self.store = ClientStore(config.state_db)
         self.backups = SQLiteBackupManager(
             config.state_db, config.backup_dir, config.backup
@@ -137,12 +146,17 @@ class ArchiveClientDaemon:
             logger.info("control_connection_registered")
             outbound: asyncio.Queue[str] = asyncio.Queue(maxsize=100)
             writer = asyncio.create_task(self._writer(websocket, outbound))
+            command_tasks: set[asyncio.Task[None]] = set()
             try:
                 async for frame in websocket:
-                    await self._handle(decode(frame), outbound)
+                    await self._handle(decode(frame), outbound, command_tasks)
             finally:
                 writer.cancel()
-                await asyncio.gather(writer, return_exceptions=True)
+                for task in command_tasks:
+                    task.cancel()
+                await asyncio.gather(
+                    writer, *command_tasks, return_exceptions=True
+                )
 
     def _registration(self):
         envelope = new_envelope()
@@ -195,6 +209,11 @@ class ArchiveClientDaemon:
         if all(probe.sparse_files for probe in self.probes):
             request.capabilities.features.append(client_pb2.CLIENT_FEATURE_SPARSE_FILES)
         request.capabilities.features.append(client_pb2.CLIENT_FEATURE_DB_BACKUP)
+        if self.inventory is not None:
+            request.capabilities.features.extend((
+                client_pb2.CLIENT_FEATURE_INVENTORY_CHUNKS,
+                client_pb2.CLIENT_FEATURE_CONTENT_TREE,
+            ))
         for cursor in self.store.list_active_job_cursors():
             active = request.active_jobs.add()
             active.job_id = str(cursor["job_id"])
@@ -211,7 +230,10 @@ class ArchiveClientDaemon:
             await websocket.send(await outbound.get())
 
     async def _handle(
-        self, envelope: Any, outbound: asyncio.Queue[str]
+        self,
+        envelope: Any,
+        outbound: asyncio.Queue[str],
+        command_tasks: set[asyncio.Task[None]],
     ) -> None:
         payload = envelope.WhichOneof("payload")
         if payload == "heartbeat":
@@ -220,7 +242,7 @@ class ArchiveClientDaemon:
             response.heartbeat_ack.sequence = envelope.heartbeat.sequence
             await outbound.put(encode(response))
         elif payload == "command":
-            await self._accept_command(envelope, outbound)
+            await self._accept_command(envelope, outbound, command_tasks)
         elif payload == "protocol_error":
             logger.warning(
                 "control_reported_protocol_error",
@@ -228,7 +250,10 @@ class ArchiveClientDaemon:
             )
 
     async def _accept_command(
-        self, envelope: Any, outbound: asyncio.Queue[str]
+        self,
+        envelope: Any,
+        outbound: asyncio.Queue[str],
+        command_tasks: set[asyncio.Task[None]],
     ) -> None:
         command = envelope.command
         snapshot_rows: list[dict[str, object]] = []
@@ -301,9 +326,55 @@ class ArchiveClientDaemon:
                     row["last_event_sequence"]
                 )
                 await outbound.put(encode(snapshot))
+        elif (
+            accepted is not None
+            and accepted_for_execution
+            and command.WhichOneof("payload") == "inventory_query"
+            and self.inventory is not None
+        ):
+            task = asyncio.create_task(
+                self._send_inventory(
+                    command.inventory_query, envelope.message_id, outbound
+                ),
+                name=f"inventory-{command.inventory_query.query_id}",
+            )
+            command_tasks.add(task)
+            task.add_done_callback(
+                lambda completed: self._command_finished(
+                    completed, command_tasks
+                )
+            )
+
+    async def _send_inventory(
+        self,
+        query: Any,
+        correlation_id: str,
+        outbound: asyncio.Queue[str],
+    ) -> None:
+        assert self.inventory is not None
+        chunks = await asyncio.to_thread(self.inventory.execute, query)
+        for chunk in chunks:
+            response = new_envelope()
+            response.correlation_id = correlation_id
+            response.inventory_chunk.CopyFrom(chunk)
+            await outbound.put(encode(response))
 
     @staticmethod
+    def _command_finished(
+        task: asyncio.Task[None], command_tasks: set[asyncio.Task[None]]
+    ) -> None:
+        command_tasks.discard(task)
+        if task.cancelled():
+            return
+        error = task.exception()
+        if error is not None:
+            logger.error(
+                "background_command_failed",
+                extra={"error_type": type(error).__name__},
+            )
+
     def _initial_acknowledgement(
+        self,
         command: Any,
         missing_snapshot_jobs: set[str],
         has_snapshot_jobs: bool,
@@ -324,6 +395,22 @@ class ArchiveClientDaemon:
                 acknowledgement.error.message = "requested client job is not found"
             else:
                 acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_ACCEPTED
+        elif command.WhichOneof("payload") == "inventory_query":
+            scope = command.inventory_query.scope
+            if self.inventory is None:
+                acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_REJECTED
+                acknowledgement.error.code = common_pb2.ERROR_CODE_UNAVAILABLE
+                acknowledgement.error.message = "inventory adapter is unavailable"
+            elif scope not in {
+                inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES,
+                inventory_pb2.INVENTORY_SCOPE_RESOURCE_LOOKUP,
+                inventory_pb2.INVENTORY_SCOPE_CONTENT_TREE,
+            }:
+                acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_REJECTED
+                acknowledgement.error.code = common_pb2.ERROR_CODE_UNSUPPORTED
+                acknowledgement.error.message = "inventory scope is unsupported"
+            else:
+                acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_ACCEPTED
         else:
             acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_REJECTED
             acknowledgement.error.code = common_pb2.ERROR_CODE_UNSUPPORTED