Joseph Doherty
365 lines
12 KiB
Python
365 lines
12 KiB
Python
"""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 BrowseChildrenOptions, ClientOptions, create_channel
|
|
|
|
_DISCOVER_HIERARCHY_PAGE_SIZE = 5000
|
|
_BROWSE_CHILDREN_PAGE_SIZE = 500
|
|
|
|
|
|
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)
|
|
|
|
async def browse_children_raw(
|
|
self, request: galaxy_pb.BrowseChildrenRequest
|
|
) -> galaxy_pb.BrowseChildrenReply:
|
|
"""Issue one BrowseChildren RPC and return the raw reply.
|
|
|
|
Lower-level escape hatch for callers that need direct page-token control
|
|
or do not want LazyBrowseNode wrapping. Most callers should use
|
|
:py:meth:`browse` and :py:meth:`LazyBrowseNode.expand` instead.
|
|
"""
|
|
|
|
return await self._unary(
|
|
"browse children",
|
|
self.raw_stub.BrowseChildren,
|
|
request,
|
|
)
|
|
|
|
async def browse(
|
|
self,
|
|
options: BrowseChildrenOptions | None = None,
|
|
) -> list["LazyBrowseNode"]:
|
|
"""Return the root browse nodes for lazy hierarchy traversal.
|
|
|
|
Each returned ``LazyBrowseNode`` wraps a Galaxy object whose direct
|
|
children can be loaded on demand by ``await node.expand()``.
|
|
"""
|
|
|
|
effective = options or BrowseChildrenOptions()
|
|
return [
|
|
node
|
|
async for node in self._iter_browse_children(
|
|
parent_gobject_id=None,
|
|
options=effective,
|
|
)
|
|
]
|
|
|
|
async def _iter_browse_children(
|
|
self,
|
|
*,
|
|
parent_gobject_id: int | None,
|
|
options: BrowseChildrenOptions,
|
|
) -> AsyncIterator["LazyBrowseNode"]:
|
|
page_token = ""
|
|
seen_page_tokens: set[str] = set()
|
|
while True:
|
|
request = galaxy_pb.BrowseChildrenRequest(
|
|
page_size=_BROWSE_CHILDREN_PAGE_SIZE,
|
|
page_token=page_token,
|
|
alarm_bearing_only=options.alarm_bearing_only,
|
|
historized_only=options.historized_only,
|
|
)
|
|
if parent_gobject_id is not None:
|
|
request.parent_gobject_id = parent_gobject_id
|
|
if options.category_ids:
|
|
request.category_ids.extend(options.category_ids)
|
|
if options.template_chain_contains:
|
|
request.template_chain_contains.extend(options.template_chain_contains)
|
|
if options.tag_name_glob:
|
|
request.tag_name_glob = options.tag_name_glob
|
|
if options.include_attributes is not None:
|
|
request.include_attributes = options.include_attributes
|
|
|
|
reply = await self._unary(
|
|
"browse children",
|
|
self.raw_stub.BrowseChildren,
|
|
request,
|
|
)
|
|
|
|
for index, obj in enumerate(reply.children):
|
|
hint = (
|
|
index < len(reply.child_has_children)
|
|
and bool(reply.child_has_children[index])
|
|
)
|
|
yield LazyBrowseNode(self, obj, hint, options)
|
|
|
|
page_token = reply.next_page_token
|
|
if not page_token:
|
|
return
|
|
if page_token in seen_page_tokens:
|
|
raise MxGatewayError(
|
|
f"galaxy browse children 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
|
|
|
|
|
|
class LazyBrowseNode:
|
|
"""One node in a lazy-loaded Galaxy browse tree.
|
|
|
|
Calling ``expand`` once fetches direct children (paginating as needed)
|
|
and populates ``children``. Subsequent calls are no-ops so callers can
|
|
drive UI expand toggles without de-duping.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
client: "GalaxyRepositoryClient",
|
|
obj: galaxy_pb.GalaxyObject,
|
|
has_children_hint: bool,
|
|
options: BrowseChildrenOptions,
|
|
) -> None:
|
|
"""Initialize a node bound to its owning client and filter set."""
|
|
self._client = client
|
|
self._object = obj
|
|
self._has_children_hint = has_children_hint
|
|
self._options = options
|
|
self._children: list[LazyBrowseNode] = []
|
|
self._is_expanded = False
|
|
self._expand_lock = asyncio.Lock()
|
|
|
|
@property
|
|
def object(self) -> galaxy_pb.GalaxyObject:
|
|
"""Return the underlying ``GalaxyObject`` proto for this node."""
|
|
return self._object
|
|
|
|
@property
|
|
def has_children_hint(self) -> bool:
|
|
"""Return the server hint about whether this node has children."""
|
|
return self._has_children_hint
|
|
|
|
@property
|
|
def children(self) -> list["LazyBrowseNode"]:
|
|
"""Return a copy of the loaded child nodes (empty until expanded)."""
|
|
return list(self._children)
|
|
|
|
@property
|
|
def is_expanded(self) -> bool:
|
|
"""Return whether ``expand`` has already populated ``children``."""
|
|
return self._is_expanded
|
|
|
|
async def expand(self) -> None:
|
|
"""Fetch direct children of this node; no-op on subsequent calls."""
|
|
if self._is_expanded:
|
|
return
|
|
async with self._expand_lock:
|
|
if self._is_expanded:
|
|
return
|
|
new_children: list[LazyBrowseNode] = []
|
|
async for child in self._client._iter_browse_children(
|
|
parent_gobject_id=self._object.gobject_id,
|
|
options=self._options,
|
|
):
|
|
new_children.append(child)
|
|
self._children.extend(new_children)
|
|
self._is_expanded = True
|
|
|
|
|
|
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()
|