"""Async Galaxy Repository client wrapper. Wraps the read-only ``GalaxyRepository`` gRPC service exposed by the MxAccess Gateway. The service lets callers test connectivity to the AVEVA System Platform Galaxy Repository (ZB SQL database), read the last deployment timestamp, and enumerate the deployed object hierarchy plus the attributes on each object. """ from __future__ import annotations import asyncio from collections.abc import AsyncIterator, Sequence from datetime import datetime from typing import Any import grpc from google.protobuf.timestamp_pb2 import Timestamp from .auth import merge_metadata from .errors import MxGatewayError, map_rpc_error from .generated import galaxy_repository_pb2 as galaxy_pb from .generated import galaxy_repository_pb2_grpc as galaxy_pb_grpc from .options import ClientOptions, create_channel _DISCOVER_HIERARCHY_PAGE_SIZE = 5000 class GalaxyRepositoryClient: """Async client for the Galaxy Repository gRPC service.""" def __init__( self, *, options: ClientOptions, stub: Any, channel: grpc.aio.Channel | None = None, ) -> None: """Initialize the client with resolved options and a gRPC stub.""" self.options = options self.raw_stub = stub self._channel = channel self._closed = False @classmethod async def connect( cls, options: ClientOptions | None = None, *, endpoint: str | None = None, api_key: str | None = None, plaintext: bool = False, ca_file: str | None = None, server_name_override: str | None = None, stub: Any | None = None, ) -> "GalaxyRepositoryClient": """Create a client with either a real async channel or a supplied fake stub.""" resolved = options or ClientOptions( endpoint=endpoint or "", api_key=api_key, plaintext=plaintext, ca_file=ca_file, server_name_override=server_name_override, ) if stub is not None: return cls(options=resolved, stub=stub) channel = create_channel(resolved) return cls( options=resolved, stub=galaxy_pb_grpc.GalaxyRepositoryStub(channel), channel=channel, ) async def __aenter__(self) -> "GalaxyRepositoryClient": """Return self to support ``async with`` usage.""" return self async def __aexit__(self, *_exc_info: object) -> None: """Close the client when leaving an ``async with`` block.""" await self.close() async def close(self) -> None: """Close the owned gRPC channel.""" if self._closed: return if self._channel is not None: await self._channel.close() self._closed = True async def test_connection(self) -> bool: """Return ``True`` when the gateway can reach the Galaxy Repository DB.""" reply = await self._unary( "test connection", self.raw_stub.TestConnection, galaxy_pb.TestConnectionRequest(), ) return bool(reply.ok) async def get_last_deploy_time(self) -> datetime | None: """Return the last Galaxy deploy timestamp or ``None`` when unset.""" reply = await self._unary( "get last deploy time", self.raw_stub.GetLastDeployTime, galaxy_pb.GetLastDeployTimeRequest(), ) if not reply.present: return None return reply.time_of_last_deploy.ToDatetime() async def discover_hierarchy(self) -> list[galaxy_pb.GalaxyObject]: """Return the deployed Galaxy object hierarchy as raw proto messages.""" objects: list[galaxy_pb.GalaxyObject] = [] seen_page_tokens: set[str] = set() page_token = "" while True: reply = await self._unary( "discover hierarchy", self.raw_stub.DiscoverHierarchy, galaxy_pb.DiscoverHierarchyRequest( page_size=_DISCOVER_HIERARCHY_PAGE_SIZE, page_token=page_token, ), ) objects.extend(reply.objects) page_token = reply.next_page_token if not page_token: return objects if page_token in seen_page_tokens: raise MxGatewayError( f"galaxy discover hierarchy returned repeated page token {page_token!r}" ) seen_page_tokens.add(page_token) def watch_deploy_events( self, last_seen_deploy_time: datetime | None = None, *, metadata: Sequence[tuple[str, str]] | None = None, ) -> AsyncIterator[galaxy_pb.DeployEvent]: """Stream Galaxy deploy events. On subscribe the gateway emits the current cached state and then one event per new deploy time. ``sequence`` is monotonic per server start; gaps mean events were dropped from the per-subscriber buffer. When ``last_seen_deploy_time`` is supplied and matches the current cached deploy time the bootstrap event is suppressed. """ request = galaxy_pb.WatchDeployEventsRequest() if last_seen_deploy_time is not None: timestamp = Timestamp() timestamp.FromDatetime(last_seen_deploy_time) request.last_seen_deploy_time.CopyFrom(timestamp) kwargs: dict[str, Any] = {"metadata": merge_metadata(self.options.api_key, metadata)} if self.options.stream_timeout is not None: kwargs["timeout"] = self.options.stream_timeout try: call = self.raw_stub.WatchDeployEvents(request, **kwargs) except TypeError as error: if "timeout" not in kwargs or "unexpected keyword argument 'timeout'" not in str(error): raise kwargs.pop("timeout") call = self.raw_stub.WatchDeployEvents(request, **kwargs) return _canceling_iterator(call) async def _unary( self, operation: str, method: Any, request: Any, *, metadata: Sequence[tuple[str, str]] | None = None, ) -> Any: kwargs: dict[str, Any] = {"metadata": merge_metadata(self.options.api_key, metadata)} if self.options.call_timeout is not None: kwargs["timeout"] = self.options.call_timeout try: call = method(request, **kwargs) except TypeError as error: if "timeout" not in kwargs or "unexpected keyword argument 'timeout'" not in str(error): raise kwargs.pop("timeout") call = method(request, **kwargs) try: return await call except asyncio.CancelledError: cancel = getattr(call, "cancel", None) if cancel is not None: cancel() raise except grpc.RpcError as error: raise map_rpc_error(operation, error) from error async def _canceling_iterator(call: Any) -> AsyncIterator[galaxy_pb.DeployEvent]: try: async for event in call: yield event except asyncio.CancelledError: cancel = getattr(call, "cancel", None) if cancel is not None: cancel() raise except grpc.RpcError as error: raise map_rpc_error("watch deploy events", error) from error finally: cancel = getattr(call, "cancel", None) if cancel is not None: cancel()