From 9b2280d50cc6b075c91e5b969a4e9aa7cb8d2a13 Mon Sep 17 00:00:00 2001 From: Craig McChesney Date: Wed, 9 Sep 2026 13:49:47 -0600 Subject: [PATCH 1/3] feat: DataSet, Annotation, and Export clients (issue #6, Phase 1) Wraps the modernized DataSet / Annotation / Calculations / Export area of DpAnnotationService in the house client pattern, as three feature clients on the existing annotation facade (plan D1): - client.annotation.datasets (DataSetClient): save/get/query/iter/delete plus the get_datasets() batch fetch (D9), the DataSetQuery criterion helpers, and the data_block() builder. - client.annotation.annotations (AnnotationsClient): save/get/query/iter/delete plus get_calculations(), the AnnotationQuery helpers, and calculations(). - client.annotation.export (ExportClient): export_data(), the ExportFormat enum, and calculations_spec(). Every unary sender goes through ServiceApiClientBase._dispatch() (#14). patchDataSet / patchAnnotation are not wrapped: reserved placeholders. Two criterion-helper differences from the older #5/#9 helpers, both following the proto: attributes() accepts an absent values list as a key-only existence search, and criteria is optional because the server treats an empty list as match-all. Back-porting both to the existing helpers is #40 / #41. Client-side validation covers the server's shape rules only (D11): at most one TextCriterion, at least one export source, a non-empty pv_names, and begin < end on a data block. Plan updated with four triage findings made before implementation: - delete-not-found is a REJECT business error, not a silent success - the server never checks begin < end on a DataBlock, so data_block()'s check is the only one there is - a calculations SamplingClock is rejected when startTime.epochSeconds is 0 - the redundant sampling_clock() count is kept, matching SampleStatusFrame (D6); deriving it would fork the axis API by caller 149 new unit tests (570 total, all passing); ruff lint and format clean. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019he3UCsAnqTDE2VQ73Djwn --- plan/tickets/6/plan.md | 29 +- src/dp_python_lib/client/__init__.py | 50 ++ src/dp_python_lib/client/annotation_client.py | 18 +- .../client/annotations_client.py | 838 ++++++++++++++++++ src/dp_python_lib/client/dataset_client.py | 728 +++++++++++++++ src/dp_python_lib/client/export_client.py | 248 ++++++ tests/unit/test_annotation_client.py | 57 ++ tests/unit/test_annotations_client.py | 715 +++++++++++++++ tests/unit/test_dataset_client.py | 682 ++++++++++++++ tests/unit/test_export_client.py | 305 +++++++ 10 files changed, 3664 insertions(+), 6 deletions(-) create mode 100644 src/dp_python_lib/client/annotations_client.py create mode 100644 src/dp_python_lib/client/dataset_client.py create mode 100644 src/dp_python_lib/client/export_client.py create mode 100644 tests/unit/test_annotation_client.py create mode 100644 tests/unit/test_annotations_client.py create mode 100644 tests/unit/test_dataset_client.py create mode 100644 tests/unit/test_export_client.py diff --git a/plan/tickets/6/plan.md b/plan/tickets/6/plan.md index 9abf77e..551a8e3 100644 --- a/plan/tickets/6/plan.md +++ b/plan/tickets/6/plan.md @@ -22,6 +22,15 @@ grpcio floor raised to 1.83.1), #13 closed, follow-ups #40 / #41 filed, #14 sequenced first, and the planning convention, the `valueStatus` rewording, and the `limit=0` test fix landed in #42. This is the first plan under the `plan/tickets//` convention in this repo (previous plans lived in the gitignored `.dev/plan/`). + **Phase 0 is complete**: #14 merged as `a0ce121`, so nothing blocks Phase 1. +- **Re-triaged 2026-09-09** before implementation, against the merged stubs and the upstream sources at the commits + above (unchanged since the plan was written). Every message shape in [section 2](#2-authoritative-message-shapes-from-the-regenerated-stubs) + was re-introspected from the committed stubs and matches; the `TextCriterion`, key-only-`attributes`, export + one-source, `EXPORT_FORMAT_UNSPECIFIED`, and calculations-validation behaviors in + [section 3](#3-server-behaviors-the-client-must-encode-dp-service-248-verified-against-the-merged-prs) were + re-verified in the dp-service Java source rather than taken from the ticket. Four things the first pass missed + were folded in: the delete-not-found tier, the absent server-side `begin < end` check, the epoch-0 `SamplingClock` + rejection (all three now rows in section 3), and the redundant-`count` question (resolved in D6). ## Overview @@ -154,6 +163,9 @@ attributes a reader of this table might reach for do not exist. | Criteria AND across the list, OR within a criterion; **at most one `TextCriterion` per request** (a second is a validation `REJECT`, since two `$text` clauses cannot be ANDed) | client-side `ValueError` naming the rule is cheap and matches the "fail with a message naming the problem" posture (Q7) | | Blank criterion values are rejected server-side; `IdCriterion` ids and get/delete ids must be valid ObjectIds (malformed → `REJECT`, not "not found") | helpers reject empty inputs as usual; no ObjectId validation client-side (format is a server implementation detail) | | `getDataSet` / `getAnnotation` / `getCalculations` not-found → `ExceptionalResult` | same "business error" tier as `get_pv_metadata` | +| `deleteDataSet` / `deleteAnnotation` **not-found is also a `REJECT`**, not a silent success (`DeleteAnnotationDispatcher:34-36`, "no Annotation record found for id: …") | same business-error tier, so no code change — but the cookbook teardown and the integration test's delete-twice leg must expect an error result on the second delete, not a success | +| `DataBlock` validation is **only** `beginTime.epochSeconds >= 1`, `endTime.epochSeconds >= 1`, and a non-empty `pvNames` (`AnnotationValidationUtility.validateDataBlock`) — the server never checks `begin < end` | D4's client-side `begin < end` check is the *only* one there is, not a duplicate of a server check. A reversed block is accepted today, which is also why the proto's silence on half-openness is a real ambiguity rather than a documentation gap | +| A `SamplingClock` axis must have `startTime.epochSeconds != 0` as well as non-zero `periodNanos` and `count` (`validateCalculationsDataFrame`) | an epoch-0 start time is **rejected**: fixtures must not build a calculations axis from `datetime(1970, 1, 1)`. `sampling_clock()` does not check this (it has no reason to — the sample-status axis has no such rule), so it is a fixture discipline, not a builder change | | `deleteDataSet` rejected while referenced; message names one referencing annotation id plus the total count | surface verbatim; no client-side cascade (decision D8) | | `deleteAnnotation` deletes the annotation first, then its calculations; incoming `annotationIds` / `derivedFrom` links dangle | readers must tolerate dangling ids — document, do not resolve | | `saveAnnotation` full-replace **includes calculations**: omitting them clears (and deletes) the stored object; a replaced object is deleted; the result returns the new `calculationsId` | params carry `calculations` explicitly; the cookbook's update recipe reads with `get_annotation()` and resends | @@ -246,7 +258,8 @@ annotation has none); `QueryAnnotationsApiResult.annotations` + `next_page_token **D6 — A shared `common.DataFrame` builder module, sized for what calculations need and shaped for #17.** New `client/data_frame.py` (no optional dependencies): -- axis: `sampling_clock()` / `timestamp_list()` move here (re-exported from `sample_status_client`); +- axis: `sampling_clock()` / `timestamp_list()` move here (re-exported from `sample_status_client`), with their + signatures unchanged — see the count note below; - typed scalar columns: `double_column(name, values, metadata=None)`, `float_column`, `int64_column`, `int32_column`, `bool_column`, `string_column`, `enum_column(name, values, enum_id, metadata=None)`; - `data_column(name, values, metadata=None)` — the legacy `DataColumn` escape hatch, where a `None` entry is a @@ -266,6 +279,17 @@ module is the substrate #17 extends rather than a parallel one. `calculations(f common_pb2.DataFrame]) -> annotation_pb2.Calculations` lives in `annotations_client.py`; taking a dict makes frame-name uniqueness true by construction. (Q4, resolved as proposed.) +*On the redundant `count` (triage, 2026-09-09).* `sampling_clock(start, period, count)` makes every calculations +call site restate a number it already has — `sampling_clock(t0, period, count=len(values))` — so a mismatch between +axis and column becomes a `data_frame()` `ValueError` rather than being unrepresentable. A `data_frame()` that took +a `(start_time, period_nanos)` pair and derived the count from the columns would remove that. **Rejected**, for two +reasons. It would fork the axis API by caller: `sample_status_client` genuinely knows its count independently (it +labels a subset of an archived clock the client did not produce), so the count-bearing form has to stay, and a second +derived form beside it means two ways to say the same thing. And `SampleStatusFrame` already set the house precedent +for exactly this shape — an explicit `data_timestamps` plus per-column validation against it — so `data_frame()` +matching it is what a reader of one will expect of the other. The redundancy is real but cheap, and it is caught +client-side with a message naming the column. Decided now because the signature is breaking to change later. + **D7 — Read side: typed columns → Python, then → pandas, sharing one converter with #16.** New `client/data_frame_conversions.py`: `data_frame_timestamps(frame) -> list[int]` (epoch nanos, via `expand_data_timestamps`), `data_frame_columns(frame) -> dict[str, list]` (all 14 typed kinds plus `DataColumn` @@ -352,7 +376,8 @@ connection and answers `getDataSet` with `UNIMPLEMENTED`, so reachability is not dataset → get → query by id / owner / pv name / tag (asserting the lowercase-normalized tag) → save annotation with a hand-built `Calculations` → `get_annotation` (calculations inline, `dataSetIds` ids-only) → `get_calculations` → `query_annotations([AQ.datasets([...])])` (calculations empty, id present) → - `delete_dataset` rejected while referenced → `delete_annotation` → `delete_dataset` → paging across a + `delete_dataset` rejected while referenced → `delete_annotation` → `delete_dataset` → deleting either a second + time is a business *error*, not a success (finding above) → paging across a run-unique tag with `limit=1` → malformed page token is a business error. Everything written is deleted, keyed by a run-unique owner id. The builder, pandas, and export legs are added in Phase 4. diff --git a/src/dp_python_lib/client/__init__.py b/src/dp_python_lib/client/__init__.py index e048982..8c139c8 100644 --- a/src/dp_python_lib/client/__init__.py +++ b/src/dp_python_lib/client/__init__.py @@ -1,4 +1,32 @@ from dp_python_lib.client.annotation_client import AnnotationClient +from dp_python_lib.client.annotations_client import ( + AnnotationQuery, + AnnotationsClient, + DeleteAnnotationApiResult, + GetAnnotationApiResult, + GetCalculationsApiResult, + QueryAnnotationsApiResult, + SaveAnnotationApiResult, + SaveAnnotationRequestParams, + calculations, +) +from dp_python_lib.client.dataset_client import ( + DataSetClient, + DataSetQuery, + DeleteDataSetApiResult, + GetDataSetApiResult, + QueryDataSetsApiResult, + SaveDataSetApiResult, + SaveDataSetRequestParams, + data_block, +) +from dp_python_lib.client.export_client import ( + ExportClient, + ExportDataApiResult, + ExportDataRequestParams, + ExportFormat, + calculations_spec, +) from dp_python_lib.client.ingestion_client import ( IngestionClient, RegisterProviderApiResult, @@ -55,16 +83,29 @@ __all__ = [ "AnnotationClient", + "AnnotationQuery", + "AnnotationsClient", "ConfigQuery", "ConfigurationActivationQuery", "ConfigurationQuery", + "DataSetClient", + "DataSetQuery", + "DeleteAnnotationApiResult", "DeleteConfigurationActivationApiResult", "DeleteConfigurationApiResult", + "DeleteDataSetApiResult", "DeletePvMetadataApiResult", "DeleteSampleStatusesApiResult", + "ExportClient", + "ExportDataApiResult", + "ExportDataRequestParams", + "ExportFormat", "GetActiveConfigurationsApiResult", + "GetAnnotationApiResult", + "GetCalculationsApiResult", "GetConfigurationActivationApiResult", "GetConfigurationApiResult", + "GetDataSetApiResult", "GetPvMetadataApiResult", "IngestionClient", "MachineConfigClient", @@ -72,9 +113,11 @@ "PvMetadataClient", "PvMetadataQuery", "PvQuery", + "QueryAnnotationsApiResult", "QueryClient", "QueryConfigurationActivationsApiResult", "QueryConfigurationsApiResult", + "QueryDataSetsApiResult", "QueryParams", "QueryPvMetadataApiResult", "QuerySampleStatusesApiResult", @@ -87,14 +130,21 @@ "SampleStatusFilter", "SampleStatusFrame", "SampleStatusRow", + "SaveAnnotationApiResult", + "SaveAnnotationRequestParams", "SaveConfigurationActivationApiResult", "SaveConfigurationActivationRequestParams", "SaveConfigurationApiResult", "SaveConfigurationRequestParams", + "SaveDataSetApiResult", + "SaveDataSetRequestParams", "SavePvMetadataApiResult", "SavePvMetadataRequestParams", "SaveSampleStatusesApiResult", "SaveSampleStatusesRequestParams", + "calculations", + "calculations_spec", + "data_block", "sampling_clock", "timestamp_list", "to_timestamp", diff --git a/src/dp_python_lib/client/annotation_client.py b/src/dp_python_lib/client/annotation_client.py index e857f26..a60bfb8 100644 --- a/src/dp_python_lib/client/annotation_client.py +++ b/src/dp_python_lib/client/annotation_client.py @@ -2,6 +2,9 @@ import grpc +from dp_python_lib.client.annotations_client import AnnotationsClient +from dp_python_lib.client.dataset_client import DataSetClient +from dp_python_lib.client.export_client import ExportClient from dp_python_lib.client.machine_config_client import MachineConfigClient from dp_python_lib.client.pv_metadata_client import PvMetadataClient from dp_python_lib.client.sample_status_client import SampleStatusClient @@ -10,15 +13,19 @@ class AnnotationClient: """ Facade for the MLDP Annotation Service. The upstream DpAnnotationService owns several distinct feature areas - (PV metadata, machine configuration, sample status, annotations, ...); this facade groups the corresponding - feature-scoped clients under one object, all sharing the single Annotation Service channel. + (PV metadata, machine configuration, sample status, datasets, annotations, export); this facade groups the + corresponding feature-scoped clients under one object, all sharing the single Annotation Service channel. - Currently exposes: + Exposes: - pv_metadata: PvMetadataClient for the PV metadata API methods. - machine_config: MachineConfigClient for the machine configuration API methods. - sample_status: SampleStatusClient for the sample status API methods. + - datasets: DataSetClient for the DataSet API methods. + - annotations: AnnotationsClient for the annotation and calculations API methods. + - export: ExportClient for the data export API method. - Future feature clients (annotations) will be added here as additional attributes. + Note the near-collision between this facade (AnnotationClient, singular) and the feature client it exposes as + .annotations (AnnotationsClient, plural). Users reach both through MldpClient and construct neither directly. """ def __init__(self, channel: grpc.Channel) -> None: @@ -30,4 +37,7 @@ def __init__(self, channel: grpc.Channel) -> None: self.pv_metadata = PvMetadataClient(channel) self.machine_config = MachineConfigClient(channel) self.sample_status = SampleStatusClient(channel) + self.datasets = DataSetClient(channel) + self.annotations = AnnotationsClient(channel) + self.export = ExportClient(channel) self.logger.debug("AnnotationClient initialized with channel: %s", channel) diff --git a/src/dp_python_lib/client/annotations_client.py b/src/dp_python_lib/client/annotations_client.py new file mode 100644 index 0000000..7130b62 --- /dev/null +++ b/src/dp_python_lib/client/annotations_client.py @@ -0,0 +1,838 @@ +import logging +from collections.abc import Iterator + +import grpc + +from dp_python_lib.client.dataset_client import _check_at_most_one_text_criterion +from dp_python_lib.client.result import ApiResultBase +from dp_python_lib.client.service_api_client_base import ServiceApiClientBase +from dp_python_lib.grpc import annotation_pb2, annotation_pb2_grpc, common_pb2 + + +def calculations(frames: dict[str, common_pb2.DataFrame]) -> annotation_pb2.Calculations: + """ + Builds a Calculations payload from named data frames -- the derived values an annotation attaches to its + DataSets, one frame per time axis. + + Taking a dict rather than a list makes frame-name uniqueness true by construction; the server rejects duplicate + frame names, and a list would let a caller build one. + + Build each frame with data_frame.data_frame(); it validates the frame's internal shape (column count against the + time axis, unique column names, non-empty names and values) so an error names the offending column. + + Note the proto's naming trap: the repeated field is 'calculationDataFrames' (singular "calculation") while the + message it holds is 'CalculationsDataFrame' (plural). + + :param frames: Mapping of frame name to the common.DataFrame carrying that frame's columns. + :return: An annotation.Calculations carrying one CalculationsDataFrame per entry. + :raises ValueError: if frames is empty or any frame name is empty. + """ + if not frames: + raise ValueError("calculations() requires at least one frame") + + result = annotation_pb2.Calculations() + for frame_name, frame in frames.items(): + if not frame_name: + raise ValueError("calculations() requires a non-empty name for every frame") + calculations_frame = result.calculationDataFrames.add() + calculations_frame.name = frame_name + calculations_frame.frame.CopyFrom(frame) + return result + + +class AnnotationQuery: + """ + Factory of lightweight helpers for building QueryAnnotationsRequest.QueryAnnotationsCriterion objects for use + with AnnotationsClient.query_annotations() and iter_annotations(). Each helper returns a single criterion; + callers pass a list of criteria to the query methods. + + Criteria AND across the list and OR within a single criterion. An empty or omitted criteria list matches all + annotations, so "browse everything, paged" is a legitimate call. + + Example: + from dp_python_lib.client import AnnotationQuery as AQ + criteria = [AQ.tags(["reviewed"]), AQ.datasets([dataset_id])] + result = client.annotation.annotations.query_annotations(criteria) + """ + + _Criterion = annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion + + @staticmethod + def ids(values: list[str]) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotations with any of the specified ids. + :param values: Annotation ids to match. + :return: A QueryAnnotationsCriterion with an idCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("ids() requires a non-empty values list") + criterion = AnnotationQuery._Criterion() + criterion.idCriterion.ids[:] = values + return criterion + + @staticmethod + def owners(values: list[str]) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotations owned by any of the specified owner ids. + :param values: Owner ids to match. + :return: A QueryAnnotationsCriterion with an ownerCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("owners() requires a non-empty values list") + criterion = AnnotationQuery._Criterion() + criterion.ownerCriterion.ownerIds[:] = values + return criterion + + @staticmethod + def datasets(values: list[str]) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotations that target any of the specified DataSet ids. + + This is the criterion to use before deleting a DataSet, to find the annotations blocking the delete. + + :param values: DataSet ids to match. + :return: A QueryAnnotationsCriterion with a dataSetsCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("datasets() requires a non-empty values list") + criterion = AnnotationQuery._Criterion() + criterion.dataSetsCriterion.dataSetIds[:] = values + return criterion + + @staticmethod + def annotations(values: list[str]) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotations that reference any of the specified annotation ids. + + These are soft references: deleting an annotation does not clean up incoming links, so a matched + annotationIds entry may name an annotation that no longer exists. + + :param values: Referenced annotation ids to match. + :return: A QueryAnnotationsCriterion with an annotationsCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("annotations() requires a non-empty values list") + criterion = AnnotationQuery._Criterion() + criterion.annotationsCriterion.annotationIds[:] = values + return criterion + + @staticmethod + def name( + exact: list[str] | None = None, + prefix: list[str] | None = None, + contains: list[str] | None = None, + ) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotation names by exact value, prefix, and/or substring. + :param exact: Names to match exactly. + :param prefix: Name prefixes to match. + :param contains: Substrings the name must contain. + :return: A QueryAnnotationsCriterion with a nameCriterion. + :raises ValueError: if none of exact/prefix/contains is provided and non-empty. + """ + if not (exact or prefix or contains): + raise ValueError("name() requires at least one non-empty of exact/prefix/contains") + criterion = AnnotationQuery._Criterion() + name_criterion = criterion.nameCriterion + if exact: + name_criterion.exact[:] = exact + if prefix: + name_criterion.prefix[:] = prefix + if contains: + name_criterion.contains[:] = contains + return criterion + + @staticmethod + def text(text: str) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion running a full-text search over the annotation's indexed text fields (name and + description). + + This is a collection-level text index search, not a per-field match; use name() when a match must be + restricted to the name. At most ONE text criterion is allowed per request -- two $text clauses cannot be + ANDed, so the server rejects a second one, and query_annotations() rejects it client-side first. + + :param text: The text to search for. + :return: A QueryAnnotationsCriterion with a textCriterion. + :raises ValueError: if text is empty. + """ + if not text: + raise ValueError("text() requires a non-empty text value") + criterion = AnnotationQuery._Criterion() + criterion.textCriterion.text = text + return criterion + + @staticmethod + def tags(values: list[str]) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotations having any of the specified tags. + + Tags are normalized (lowercased, deduplicated, sorted) when saved, so match against the lowercase form. + + :param values: Tag values to match. + :return: A QueryAnnotationsCriterion with a tagsCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("tags() requires a non-empty values list") + criterion = AnnotationQuery._Criterion() + criterion.tagsCriterion.values[:] = values + return criterion + + @staticmethod + def attributes( + key: str, values: list[str] | None = None + ) -> "annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion": + """ + Builds a criterion matching annotations by attribute key and optional value(s). + + Omitting values (or passing an empty list) performs a key-only existence search: any annotation possessing + the key matches, whatever its value. This differs from the older PvMetadataQuery/ConfigurationQuery + helpers, which require values; relaxing those is issue #40. + + :param key: Attribute key to match (maps to Attribute.name). + :param values: Attribute values to match for that key, or None for a key-only existence search. + :return: A QueryAnnotationsCriterion with an attributesCriterion. + :raises ValueError: if key is empty. + """ + if not key: + raise ValueError("attributes() requires a non-empty key") + criterion = AnnotationQuery._Criterion() + criterion.attributesCriterion.key = key + if values: + criterion.attributesCriterion.values[:] = values + return criterion + + +class SaveAnnotationRequestParams: + """ + Encapsulates client parameters for a call to the saveAnnotation() API method. + + Saving is an id-driven upsert that REPLACES IN FULL, and that includes the calculations: supplying annotation_id + without calculations CLEARS the stored calculations object and deletes it. To change an annotation without + losing its calculations, read the current state with get_annotation() (the only method that returns calculations + inline) and resend them. Omitting annotation_id creates a new annotation. + """ + + def __init__( + self, + name: str, + owner_id: str, + dataset_ids: list[str], + annotation_ids: list[str] | None = None, + description: str | None = None, + tags: list[str] | None = None, + attributes: dict[str, str] | None = None, + modified_by: str | None = None, + calculations: annotation_pb2.Calculations | None = None, + annotation_id: str | None = None, + ) -> None: + """ + :param name: Human-readable name of the annotation. Required by the server. + :param owner_id: Identifier of the annotation's owner. Required by the server. + :param dataset_ids: Ids of the DataSets this annotation targets. Must be non-empty. + :param annotation_ids: Ids of other annotations this one references (soft links; may dangle). + :param description: Human-readable description of the annotation. + :param tags: List of tags (keywords). Normalized lowercase/deduplicated/sorted on save. + :param attributes: Map of key/value attributes describing the annotation. + :param modified_by: Identifier of the user or process making the change. + :param calculations: Derived values to store with the annotation (see calculations()). saveAnnotation() is + the only write path for calculations, and omitting them on a replace clears the stored object. + :param annotation_id: Id of an existing annotation to replace in full. Omit to create a new one. + """ + self.name = name + self.owner_id = owner_id + self.dataset_ids = dataset_ids + self.annotation_ids = annotation_ids + self.description = description + self.tags = tags + self.attributes = attributes + self.modified_by = modified_by + self.calculations = calculations + self.annotation_id = annotation_id + + +class SaveAnnotationApiResult(ApiResultBase): + """ + Wraps the response from saveAnnotation(), with a status object including an error flag and message. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.SaveAnnotationResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The SaveAnnotationResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def annotation_id(self) -> str | None: + """Id of the annotation that was saved, or None on error.""" + if self.response is not None and self.response.HasField("saveAnnotationResult"): + return self.response.saveAnnotationResult.annotationId + return None + + @property + def calculations_id(self) -> str | None: + """ + Id of the calculations object stored with the annotation, or None on error. + + This is empty string -- not None -- when the request carried no calculations; None means the call failed. + A replace that changes the calculations returns a NEW id, and the previous object is deleted. + """ + if self.response is not None and self.response.HasField("saveAnnotationResult"): + return self.response.saveAnnotationResult.calculationsId + return None + + +class GetAnnotationApiResult(ApiResultBase): + """ + Wraps the response from getAnnotation(), with a status object including an error flag and message. + + This is the only method that returns an annotation's calculations inline; queryAnnotations() results carry the + calculationsId but leave the content empty. An annotation that does not exist is a business error. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.GetAnnotationResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The GetAnnotationResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def annotation(self) -> annotation_pb2.Annotation | None: + """The requested Annotation, or None on error.""" + if self.response is not None and self.response.HasField("getAnnotationResult"): + return self.response.getAnnotationResult.annotation + return None + + @property + def calculations(self) -> annotation_pb2.Calculations | None: + """ + The annotation's inline Calculations, or None on error or when the annotation has none. + + A calculationsId that resolves to nothing is a server-side error (data corruption), never an empty success, + so an annotation reported here as having no calculations genuinely has none. + """ + annotation = self.annotation + if annotation is not None and annotation.HasField("calculations"): + return annotation.calculations + return None + + +class QueryAnnotationsApiResult(ApiResultBase): + """ + Wraps a single page of the response from queryAnnotations(), with a status object including an error flag and + message. Use AnnotationsClient.iter_annotations() to transparently page through all results. + + Query results carry ids rather than content: an annotation's calculationsId is populated but its calculations + are not. Fetch them with get_annotation() or get_calculations(). + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.QueryAnnotationsResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The QueryAnnotationsResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def annotations(self) -> list[annotation_pb2.Annotation]: + """The Annotation records in this page, or an empty list on error.""" + if self.response is not None and self.response.HasField("annotationsResult"): + return list(self.response.annotationsResult.annotations) + return [] + + @property + def next_page_token(self) -> str: + """Token for retrieving the next page, or empty string if there are no more pages.""" + if self.response is not None and self.response.HasField("annotationsResult"): + return self.response.annotationsResult.nextPageToken + return "" + + +class DeleteAnnotationApiResult(ApiResultBase): + """ + Wraps the response from deleteAnnotation(), with a status object including an error flag and message. + + Deleting an annotation that does not exist is a business error, not a silent success. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.DeleteAnnotationResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The DeleteAnnotationResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def annotation_id(self) -> str | None: + """Id of the annotation that was deleted, or None on error.""" + if self.response is not None and self.response.HasField("deleteAnnotationResult"): + return self.response.deleteAnnotationResult.annotationId + return None + + +class GetCalculationsApiResult(ApiResultBase): + """ + Wraps the response from getCalculations(), with a status object including an error flag and message. + + This is the click-through path from a query result: queryAnnotations() gives you a calculationsId, and this + fetches the content it names without re-fetching the whole annotation. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.GetCalculationsResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The GetCalculationsResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def calculations(self) -> annotation_pb2.Calculations | None: + """The requested Calculations, or None on error.""" + if self.response is not None and self.response.HasField("getCalculationsResult"): + return self.response.getCalculationsResult.calculations + return None + + +class AnnotationsClient(ServiceApiClientBase): + """ + User-facing client for the annotation methods of the MLDP Annotation Service. An annotation describes one or + more DataSets -- a name, a description, tags, and optionally a Calculations payload of derived values with + column-level provenance. + + Provides low-level wrappers for saveAnnotation(), getAnnotation(), queryAnnotations(), deleteAnnotation(), and + getCalculations(), plus the iter_annotations() paging iterator. + + Note this is AnnotationsClient (plural), the feature client, as distinct from AnnotationClient (singular), the + facade that owns the Annotation Service channel and exposes this client as client.annotation.annotations. Users + reach this through the facade and never construct either directly. + + patchAnnotation() is not wrapped: it is a reserved placeholder that returns "not implemented". + """ + + def __init__(self, channel: grpc.Channel) -> None: + """ + :param channel: gRPC communication channel for the Annotation Service. + """ + super().__init__(channel, annotation_pb2_grpc.DpAnnotationServiceStub) + self.logger = logging.getLogger(__name__) + self.logger.debug("AnnotationsClient initialized with channel: %s", channel) + + # ------------------------------------------------------------------ + # saveAnnotation + # ------------------------------------------------------------------ + + def _build_save_annotation_request( + self, request_params: SaveAnnotationRequestParams + ) -> annotation_pb2.SaveAnnotationRequest: + """ + Builds a SaveAnnotationRequest from the supplied SaveAnnotationRequestParams. + :param request_params: User parameters for the call to saveAnnotation(). + :return: A SaveAnnotationRequest for the specified params. + """ + self.logger.debug("Building SaveAnnotationRequest for annotation: %s", request_params.name) + + request = annotation_pb2.SaveAnnotationRequest() + request.name = request_params.name + request.ownerId = request_params.owner_id + + if request_params.annotation_id: + request.id = request_params.annotation_id + + if request_params.dataset_ids: + request.dataSetIds[:] = request_params.dataset_ids + + if request_params.annotation_ids: + request.annotationIds[:] = request_params.annotation_ids + + if request_params.description: + request.description = request_params.description + + if request_params.tags: + request.tags[:] = request_params.tags + + if request_params.attributes: + for name, value in request_params.attributes.items(): + attribute = common_pb2.Attribute() + attribute.name = name + attribute.value = value + request.attributes.append(attribute) + + if request_params.modified_by: + request.modifiedBy = request_params.modified_by + + if request_params.calculations is not None: + request.calculations.CopyFrom(request_params.calculations) + + self.logger.debug("SaveAnnotationRequest built successfully") + return request + + def _send_save_annotation(self, request: annotation_pb2.SaveAnnotationRequest) -> SaveAnnotationApiResult: + """ + Invokes the saveAnnotation() API method with the supplied request. + :param request: SaveAnnotationRequest with parameters for the call. + :return: A SaveAnnotationApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.saveAnnotation, + request, + SaveAnnotationApiResult, + "saveAnnotationResult", + "saveAnnotation", + request_log=lambda: self.logger.info( + "Calling saveAnnotation API for annotation: %s targeting %d dataset(s)", + request.name, + len(request.dataSetIds), + ), + success_log=lambda response: self.logger.info( + "Successfully saved annotation: %s with id: %s", + request.name, + response.saveAnnotationResult.annotationId, + ), + ) + + def save_annotation(self, request_params: SaveAnnotationRequestParams) -> SaveAnnotationApiResult: + """ + User-facing method for invoking the saveAnnotation() API method. + + Saving replaces in full when request_params.annotation_id is set, and that INCLUDES the calculations: + omitting them clears and deletes the stored object. Read with get_annotation() and resend to preserve them. + + :param request_params: Contains user parameters for the call to saveAnnotation(). + :return: A SaveAnnotationApiResult with the method response and status information. + """ + self.logger.info("Starting saveAnnotation operation for annotation: %s", request_params.name) + + request = self._build_save_annotation_request(request_params) + result = self._send_save_annotation(request) + + if result.result_status.is_error: + self.logger.error("SaveAnnotation operation failed: %s", result.result_status.message) + else: + self.logger.info("SaveAnnotation operation completed successfully for: %s", request_params.name) + + return result + + # ------------------------------------------------------------------ + # getAnnotation + # ------------------------------------------------------------------ + + def _build_get_annotation_request(self, annotation_id: str) -> annotation_pb2.GetAnnotationRequest: + """ + Builds a GetAnnotationRequest for the supplied annotation id. + :param annotation_id: Id of the annotation to retrieve. + :return: A GetAnnotationRequest for the specified id. + """ + self.logger.debug("Building GetAnnotationRequest for id: %s", annotation_id) + request = annotation_pb2.GetAnnotationRequest() + request.annotationId = annotation_id + return request + + def _send_get_annotation(self, request: annotation_pb2.GetAnnotationRequest) -> GetAnnotationApiResult: + """ + Invokes the getAnnotation() API method with the supplied request. + :param request: GetAnnotationRequest with parameters for the call. + :return: A GetAnnotationApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.getAnnotation, + request, + GetAnnotationApiResult, + "getAnnotationResult", + "getAnnotation", + request_log=lambda: self.logger.info("Calling getAnnotation API for id: %s", request.annotationId), + success_log=lambda _response: self.logger.info( + "Successfully retrieved annotation for id: %s", request.annotationId + ), + ) + + def get_annotation(self, annotation_id: str) -> GetAnnotationApiResult: + """ + User-facing method for invoking the getAnnotation() API method. + + This is the only method that returns an annotation's calculations inline. An annotation that does not exist + comes back as a business error, not an empty success. + + :param annotation_id: Id of the annotation to retrieve. + :return: A GetAnnotationApiResult with the method response and status information. + """ + self.logger.info("Starting getAnnotation operation for id: %s", annotation_id) + + request = self._build_get_annotation_request(annotation_id) + result = self._send_get_annotation(request) + + if result.result_status.is_error: + self.logger.error("GetAnnotation operation failed: %s", result.result_status.message) + else: + self.logger.info("GetAnnotation operation completed successfully for id: %s", annotation_id) + + return result + + # ------------------------------------------------------------------ + # queryAnnotations + # ------------------------------------------------------------------ + + def _build_query_annotations_request( + self, + criteria: list[annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion] | None = None, + limit: int | None = None, + page_token: str | None = None, + ) -> annotation_pb2.QueryAnnotationsRequest: + """ + Builds a QueryAnnotationsRequest from the supplied criteria and paging parameters. + :param criteria: List of QueryAnnotationsCriterion objects (see AnnotationQuery helpers), or None to + match all. + :param limit: Maximum number of records to return per page (optional). + :param page_token: Token for retrieving a subsequent page (optional). + :return: A QueryAnnotationsRequest for the specified params. + """ + self.logger.debug("Building QueryAnnotationsRequest with %d criteria", len(criteria) if criteria else 0) + request = annotation_pb2.QueryAnnotationsRequest() + if criteria: + request.criteria.extend(criteria) + if limit is not None: + request.limit = limit + if page_token: + request.pageToken = page_token + return request + + def _send_query_annotations(self, request: annotation_pb2.QueryAnnotationsRequest) -> QueryAnnotationsApiResult: + """ + Invokes the queryAnnotations() API method with the supplied request. + :param request: QueryAnnotationsRequest with parameters for the call. + :return: A QueryAnnotationsApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.queryAnnotations, + request, + QueryAnnotationsApiResult, + "annotationsResult", + "queryAnnotations", + request_log=lambda: self.logger.info( + "Calling queryAnnotations API with %d criteria", len(request.criteria) + ), + success_log=lambda response: self.logger.info( + "QueryAnnotations returned %d records", len(response.annotationsResult.annotations) + ), + ) + + def query_annotations( + self, + criteria: list[annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion] | None = None, + limit: int | None = None, + page_token: str | None = None, + ) -> QueryAnnotationsApiResult: + """ + User-facing method for invoking the queryAnnotations() API method. Returns a single page of results; use + iter_annotations() to page through all results transparently. + + An omitted or empty criteria list matches all annotations. Results carry ids rather than content: the + calculationsId is populated but the calculations themselves are not. + + :param criteria: List of QueryAnnotationsCriterion objects (see AnnotationQuery helpers), or None to + match all. + :param limit: Maximum number of records to return PER PAGE -- not a cap on the total (optional). + :param page_token: Token for retrieving a subsequent page (optional). + :return: A QueryAnnotationsApiResult with a single page of results and status information. + :raises ValueError: if criteria contains more than one text criterion. + """ + criteria = criteria or [] + _check_at_most_one_text_criterion(criteria, "query_annotations()") + + self.logger.info("Starting queryAnnotations operation with %d criteria", len(criteria)) + + request = self._build_query_annotations_request(criteria, limit=limit, page_token=page_token) + result = self._send_query_annotations(request) + + if result.result_status.is_error: + self.logger.error("QueryAnnotations operation failed: %s", result.result_status.message) + else: + self.logger.info("QueryAnnotations operation completed successfully") + + return result + + def iter_annotations( + self, + criteria: list[annotation_pb2.QueryAnnotationsRequest.QueryAnnotationsCriterion] | None = None, + limit: int | None = None, + ) -> Iterator[annotation_pb2.Annotation]: + """ + Convenience generator that transparently pages through all queryAnnotations() results, following the + nextPageToken until the results are exhausted. Yields individual Annotation records. + + Raises RuntimeError if any page returns an error, so callers can distinguish failure from an empty result set. + + :param criteria: List of QueryAnnotationsCriterion objects (see AnnotationQuery helpers), or None to + match all. + :param limit: Maximum number of records to return per page (optional). + :return: An iterator over all matching Annotation records across all pages. + :raises ValueError: if criteria contains more than one text criterion. + :raises RuntimeError: if any page returns an error. + """ + page_token: str | None = None + while True: + result = self.query_annotations(criteria, limit=limit, page_token=page_token) + if result.result_status.is_error: + raise RuntimeError(f"queryAnnotations failed during paging: {result.result_status.message}") + + yield from result.annotations + + page_token = result.next_page_token + if not page_token: + break + + # ------------------------------------------------------------------ + # deleteAnnotation + # ------------------------------------------------------------------ + + def _build_delete_annotation_request(self, annotation_id: str) -> annotation_pb2.DeleteAnnotationRequest: + """ + Builds a DeleteAnnotationRequest for the supplied annotation id. + :param annotation_id: Id of the annotation to delete. + :return: A DeleteAnnotationRequest for the specified id. + """ + self.logger.debug("Building DeleteAnnotationRequest for id: %s", annotation_id) + request = annotation_pb2.DeleteAnnotationRequest() + request.annotationId = annotation_id + return request + + def _send_delete_annotation(self, request: annotation_pb2.DeleteAnnotationRequest) -> DeleteAnnotationApiResult: + """ + Invokes the deleteAnnotation() API method with the supplied request. + :param request: DeleteAnnotationRequest with parameters for the call. + :return: A DeleteAnnotationApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.deleteAnnotation, + request, + DeleteAnnotationApiResult, + "deleteAnnotationResult", + "deleteAnnotation", + request_log=lambda: self.logger.info("Calling deleteAnnotation API for id: %s", request.annotationId), + success_log=lambda _response: self.logger.info( + "Successfully deleted annotation id: %s", request.annotationId + ), + ) + + def delete_annotation(self, annotation_id: str) -> DeleteAnnotationApiResult: + """ + User-facing method for invoking the deleteAnnotation() API method. + + The delete cascades to the annotation's calculations. It does NOT clean up incoming soft references: other + annotations' annotationIds entries, and derivedFrom provenance links naming this annotation's calculations, + are left dangling for readers to tolerate. + + Deleting an annotation that does not exist is a business error, not a silent success. + + :param annotation_id: Id of the annotation to delete. + :return: A DeleteAnnotationApiResult with the method response and status information. + """ + self.logger.info("Starting deleteAnnotation operation for id: %s", annotation_id) + + request = self._build_delete_annotation_request(annotation_id) + result = self._send_delete_annotation(request) + + if result.result_status.is_error: + self.logger.error("DeleteAnnotation operation failed: %s", result.result_status.message) + else: + self.logger.info("DeleteAnnotation operation completed successfully for id: %s", annotation_id) + + return result + + # ------------------------------------------------------------------ + # getCalculations + # ------------------------------------------------------------------ + + def _build_get_calculations_request(self, calculations_id: str) -> annotation_pb2.GetCalculationsRequest: + """ + Builds a GetCalculationsRequest for the supplied calculations id. + :param calculations_id: Id of the calculations object to retrieve. + :return: A GetCalculationsRequest for the specified id. + """ + self.logger.debug("Building GetCalculationsRequest for id: %s", calculations_id) + request = annotation_pb2.GetCalculationsRequest() + request.calculationsId = calculations_id + return request + + def _send_get_calculations(self, request: annotation_pb2.GetCalculationsRequest) -> GetCalculationsApiResult: + """ + Invokes the getCalculations() API method with the supplied request. + :param request: GetCalculationsRequest with parameters for the call. + :return: A GetCalculationsApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.getCalculations, + request, + GetCalculationsApiResult, + "getCalculationsResult", + "getCalculations", + request_log=lambda: self.logger.info("Calling getCalculations API for id: %s", request.calculationsId), + success_log=lambda response: self.logger.info( + "Successfully retrieved calculations id: %s with %d frame(s)", + request.calculationsId, + len(response.getCalculationsResult.calculations.calculationDataFrames), + ), + ) + + def get_calculations(self, calculations_id: str) -> GetCalculationsApiResult: + """ + User-facing method for invoking the getCalculations() API method. + + This is the click-through path from a queryAnnotations() result, which carries the calculationsId but not the + content. Calculations that do not exist come back as a business error. + + :param calculations_id: Id of the calculations object to retrieve. + :return: A GetCalculationsApiResult with the method response and status information. + """ + self.logger.info("Starting getCalculations operation for id: %s", calculations_id) + + request = self._build_get_calculations_request(calculations_id) + result = self._send_get_calculations(request) + + if result.result_status.is_error: + self.logger.error("GetCalculations operation failed: %s", result.result_status.message) + else: + self.logger.info("GetCalculations operation completed successfully for id: %s", calculations_id) + + return result diff --git a/src/dp_python_lib/client/dataset_client.py b/src/dp_python_lib/client/dataset_client.py new file mode 100644 index 0000000..47f9a3b --- /dev/null +++ b/src/dp_python_lib/client/dataset_client.py @@ -0,0 +1,728 @@ +import logging +from collections.abc import Iterator + +import grpc + +from dp_python_lib.client.machine_config_client import TimestampInput, to_timestamp +from dp_python_lib.client.result import ApiResultBase +from dp_python_lib.client.service_api_client_base import ServiceApiClientBase +from dp_python_lib.grpc import annotation_pb2, annotation_pb2_grpc, common_pb2 + + +def data_block( + begin_time: TimestampInput, + end_time: TimestampInput, + pv_names: list[str], +) -> annotation_pb2.DataBlock: + """ + Builds a DataBlock: one time range plus the PV names covered over it. A DataSet is a list of these, and the + exportData() API accepts them inline for a one-off export that does not warrant saving a DataSet. + + The proto does not state whether the interval is half-open; the dp-grpc cookbook flags the same ambiguity. Treat + a block as covering [begin_time, end_time] inclusively unless the server documents otherwise. + + Note that begin < end is checked HERE and nowhere else: the server validates only that each bound is non-zero and + that pvNames is non-empty (AnnotationValidationUtility.validateDataBlock), so a reversed block would otherwise be + accepted and stored. + + :param begin_time: Start of the block's time range (tz-aware datetime, epoch seconds, or common.Timestamp). + :param end_time: End of the block's time range (same accepted forms). + :param pv_names: Names of the PVs the block covers. + :return: An annotation.DataBlock for the specified range and PVs. + :raises ValueError: if pv_names is empty or begin_time is not strictly before end_time. + """ + if not pv_names: + raise ValueError("data_block() requires a non-empty pv_names list") + + begin = to_timestamp(begin_time) + end = to_timestamp(end_time) + if (begin.epochSeconds, begin.nanoseconds) >= (end.epochSeconds, end.nanoseconds): + raise ValueError( + f"data_block() requires begin_time strictly before end_time; got begin " + f"{begin.epochSeconds}.{begin.nanoseconds:09d} and end {end.epochSeconds}.{end.nanoseconds:09d}" + ) + + block = annotation_pb2.DataBlock() + block.beginTime.CopyFrom(begin) + block.endTime.CopyFrom(end) + block.pvNames[:] = pv_names + return block + + +class DataSetQuery: + """ + Factory of lightweight helpers for building QueryDataSetsRequest.QueryDataSetsCriterion objects for use with + DataSetClient.query_datasets() and iter_datasets(). Each helper returns a single criterion; callers pass a list + of criteria to the query methods. + + Criteria AND across the list and OR within a single criterion. An empty or omitted criteria list matches all + DataSets, so "browse everything, paged" is a legitimate call. + + Example: + from dp_python_lib.client import DataSetQuery as DS + criteria = [DS.owners(["cmcchesney"]), DS.tags(["ramp-study"])] + result = client.annotation.datasets.query_datasets(criteria) + """ + + _Criterion = annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion + + @staticmethod + def ids(values: list[str]) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion matching DataSets with any of the specified ids. + :param values: DataSet ids to match. + :return: A QueryDataSetsCriterion with an idCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("ids() requires a non-empty values list") + criterion = DataSetQuery._Criterion() + criterion.idCriterion.ids[:] = values + return criterion + + @staticmethod + def owners(values: list[str]) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion matching DataSets owned by any of the specified owner ids. + :param values: Owner ids to match. + :return: A QueryDataSetsCriterion with an ownerCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("owners() requires a non-empty values list") + criterion = DataSetQuery._Criterion() + criterion.ownerCriterion.ownerIds[:] = values + return criterion + + @staticmethod + def name( + exact: list[str] | None = None, + prefix: list[str] | None = None, + contains: list[str] | None = None, + ) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion matching DataSet names by exact value, prefix, and/or substring. + :param exact: Names to match exactly. + :param prefix: Name prefixes to match. + :param contains: Substrings the name must contain. + :return: A QueryDataSetsCriterion with a nameCriterion. + :raises ValueError: if none of exact/prefix/contains is provided and non-empty. + """ + if not (exact or prefix or contains): + raise ValueError("name() requires at least one non-empty of exact/prefix/contains") + criterion = DataSetQuery._Criterion() + name_criterion = criterion.nameCriterion + if exact: + name_criterion.exact[:] = exact + if prefix: + name_criterion.prefix[:] = prefix + if contains: + name_criterion.contains[:] = contains + return criterion + + @staticmethod + def text(text: str) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion running a full-text search over the DataSet's indexed text fields (name and description). + + This is a collection-level text index search, not a per-field match; use name() when a match must be + restricted to the name. At most ONE text criterion is allowed per request -- two $text clauses cannot be + ANDed, so the server rejects a second one, and query_datasets() rejects it client-side first. + + :param text: The text to search for. + :return: A QueryDataSetsCriterion with a textCriterion. + :raises ValueError: if text is empty. + """ + if not text: + raise ValueError("text() requires a non-empty text value") + criterion = DataSetQuery._Criterion() + criterion.textCriterion.text = text + return criterion + + @staticmethod + def pv_names(values: list[str]) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion matching DataSets whose data blocks cover any of the specified PV names. + :param values: PV names to match. + :return: A QueryDataSetsCriterion with a pvNameCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("pv_names() requires a non-empty values list") + criterion = DataSetQuery._Criterion() + criterion.pvNameCriterion.names[:] = values + return criterion + + @staticmethod + def tags(values: list[str]) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion matching DataSets having any of the specified tags. + + Tags are normalized (lowercased, deduplicated, sorted) when saved, so match against the lowercase form. + + :param values: Tag values to match. + :return: A QueryDataSetsCriterion with a tagsCriterion. + :raises ValueError: if values is empty. + """ + if not values: + raise ValueError("tags() requires a non-empty values list") + criterion = DataSetQuery._Criterion() + criterion.tagsCriterion.values[:] = values + return criterion + + @staticmethod + def attributes( + key: str, values: list[str] | None = None + ) -> "annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion": + """ + Builds a criterion matching DataSets by attribute key and optional value(s). + + Omitting values (or passing an empty list) performs a key-only existence search: any DataSet possessing the + key matches, whatever its value. This differs from the older PvMetadataQuery/ConfigurationQuery helpers, + which require values; relaxing those is issue #40. + + :param key: Attribute key to match (maps to Attribute.name). + :param values: Attribute values to match for that key, or None for a key-only existence search. + :return: A QueryDataSetsCriterion with an attributesCriterion. + :raises ValueError: if key is empty. + """ + if not key: + raise ValueError("attributes() requires a non-empty key") + criterion = DataSetQuery._Criterion() + criterion.attributesCriterion.key = key + if values: + criterion.attributesCriterion.values[:] = values + return criterion + + +def _check_at_most_one_text_criterion(criteria: list, op_name: str) -> None: + """ + Rejects a criteria list carrying more than one text criterion. + + Two $text clauses cannot be ANDed -- Mongo rejects the query with "Too many text expressions" -- so the server + rejects the second one during validation. Catching it here fails with a message naming the rule instead of + surfacing a server rejection. + + :param criteria: The criteria list to check. + :param op_name: The calling operation, used in the error message. + :raises ValueError: if more than one criterion carries a textCriterion. + """ + text_count = sum(1 for criterion in criteria if criterion.WhichOneof("criterion") == "textCriterion") + if text_count > 1: + raise ValueError( + f"{op_name} accepts at most one text criterion, got {text_count}; two full-text clauses cannot be " + f"combined with AND. Combine the terms into a single text() criterion instead." + ) + + +class SaveDataSetRequestParams: + """ + Encapsulates client parameters for a call to the saveDataSet() API method. + + Saving is an id-driven upsert that REPLACES IN FULL: supplying dataset_id replaces that DataSet with exactly the + content given here, so omitted fields are cleared rather than left alone. Read the current state with + get_dataset() and resend it with your changes. Omitting dataset_id creates a new DataSet. + """ + + def __init__( + self, + name: str, + owner_id: str, + data_blocks: list[annotation_pb2.DataBlock], + description: str | None = None, + tags: list[str] | None = None, + attributes: dict[str, str] | None = None, + modified_by: str | None = None, + dataset_id: str | None = None, + ) -> None: + """ + :param name: Human-readable name of the DataSet. Required by the server. + :param owner_id: Identifier of the DataSet's owner. Required by the server. + :param data_blocks: The time ranges and PV names the DataSet covers (see data_block()). Must be non-empty. + :param description: Human-readable description of the DataSet. + :param tags: List of tags (keywords) describing the DataSet. Normalized lowercase/deduplicated/sorted on save. + :param attributes: Map of key/value attributes describing the DataSet. + :param modified_by: Identifier of the user or process making the change. + :param dataset_id: Id of an existing DataSet to replace in full. Omit to create a new one. + """ + self.name = name + self.owner_id = owner_id + self.data_blocks = data_blocks + self.description = description + self.tags = tags + self.attributes = attributes + self.modified_by = modified_by + self.dataset_id = dataset_id + + +class SaveDataSetApiResult(ApiResultBase): + """ + Wraps the response from saveDataSet(), with a status object including an error flag and message. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.SaveDataSetResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The SaveDataSetResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def dataset_id(self) -> str | None: + """Id of the DataSet that was saved, or None on error.""" + if self.response is not None and self.response.HasField("saveDataSetResult"): + return self.response.saveDataSetResult.dataSetId + return None + + +class GetDataSetApiResult(ApiResultBase): + """ + Wraps the response from getDataSet(), with a status object including an error flag and message. + + A DataSet that does not exist is a business error, not an empty success. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.GetDataSetResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The GetDataSetResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def dataset(self) -> annotation_pb2.DataSet | None: + """The requested DataSet, or None on error.""" + if self.response is not None and self.response.HasField("getDataSetResult"): + return self.response.getDataSetResult.dataSet + return None + + +class QueryDataSetsApiResult(ApiResultBase): + """ + Wraps a single page of the response from queryDataSets(), with a status object including an error flag and + message. Use DataSetClient.iter_datasets() to transparently page through all results. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.QueryDataSetsResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The QueryDataSetsResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def datasets(self) -> list[annotation_pb2.DataSet]: + """The DataSet records in this page, or an empty list on error.""" + if self.response is not None and self.response.HasField("dataSetsResult"): + return list(self.response.dataSetsResult.dataSets) + return [] + + @property + def next_page_token(self) -> str: + """Token for retrieving the next page, or empty string if there are no more pages.""" + if self.response is not None and self.response.HasField("dataSetsResult"): + return self.response.dataSetsResult.nextPageToken + return "" + + +class DeleteDataSetApiResult(ApiResultBase): + """ + Wraps the response from deleteDataSet(), with a status object including an error flag and message. + + Two cases surface as business errors rather than successes: deleting a DataSet still referenced by an annotation + (the message names one referencing annotation id and the total count), and deleting one that does not exist. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.DeleteDataSetResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The DeleteDataSetResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def dataset_id(self) -> str | None: + """Id of the DataSet that was deleted, or None on error.""" + if self.response is not None and self.response.HasField("deleteDataSetResult"): + return self.response.deleteDataSetResult.dataSetId + return None + + +class DataSetClient(ServiceApiClientBase): + """ + User-facing client for the DataSet methods of the MLDP Annotation Service. A DataSet names a region of the + archive -- a list of DataBlocks, each a time range plus the PVs covered over it -- so that region can be found, + annotated, and exported later. + + Provides low-level wrappers for saveDataSet(), getDataSet(), queryDataSets(), and deleteDataSet(), plus the + iter_datasets() paging iterator and the get_datasets() batch fetch. + + patchDataSet() is not wrapped: it is a reserved placeholder that returns "not implemented". + """ + + def __init__(self, channel: grpc.Channel) -> None: + """ + :param channel: gRPC communication channel for the Annotation Service. + """ + super().__init__(channel, annotation_pb2_grpc.DpAnnotationServiceStub) + self.logger = logging.getLogger(__name__) + self.logger.debug("DataSetClient initialized with channel: %s", channel) + + # ------------------------------------------------------------------ + # saveDataSet + # ------------------------------------------------------------------ + + def _build_save_dataset_request( + self, request_params: SaveDataSetRequestParams + ) -> annotation_pb2.SaveDataSetRequest: + """ + Builds a SaveDataSetRequest from the supplied SaveDataSetRequestParams. + :param request_params: User parameters for the call to saveDataSet(). + :return: A SaveDataSetRequest for the specified params. + """ + self.logger.debug("Building SaveDataSetRequest for DataSet: %s", request_params.name) + + request = annotation_pb2.SaveDataSetRequest() + request.name = request_params.name + request.ownerId = request_params.owner_id + + if request_params.dataset_id: + request.id = request_params.dataset_id + + if request_params.data_blocks: + request.dataBlocks.extend(request_params.data_blocks) + + if request_params.description: + request.description = request_params.description + + if request_params.tags: + request.tags[:] = request_params.tags + + if request_params.attributes: + for name, value in request_params.attributes.items(): + attribute = common_pb2.Attribute() + attribute.name = name + attribute.value = value + request.attributes.append(attribute) + + if request_params.modified_by: + request.modifiedBy = request_params.modified_by + + self.logger.debug("SaveDataSetRequest built successfully") + return request + + def _send_save_dataset(self, request: annotation_pb2.SaveDataSetRequest) -> SaveDataSetApiResult: + """ + Invokes the saveDataSet() API method with the supplied request. + :param request: SaveDataSetRequest with parameters for the call. + :return: A SaveDataSetApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.saveDataSet, + request, + SaveDataSetApiResult, + "saveDataSetResult", + "saveDataSet", + request_log=lambda: self.logger.info( + "Calling saveDataSet API for DataSet: %s with %d data blocks", request.name, len(request.dataBlocks) + ), + success_log=lambda response: self.logger.info( + "Successfully saved DataSet: %s with id: %s", request.name, response.saveDataSetResult.dataSetId + ), + ) + + def save_dataset(self, request_params: SaveDataSetRequestParams) -> SaveDataSetApiResult: + """ + User-facing method for invoking the saveDataSet() API method. + + Saving replaces in full when request_params.dataset_id is set: omitted fields are cleared, not preserved. + + :param request_params: Contains user parameters for the call to saveDataSet(). + :return: A SaveDataSetApiResult with the method response and status information. + """ + self.logger.info("Starting saveDataSet operation for DataSet: %s", request_params.name) + + request = self._build_save_dataset_request(request_params) + result = self._send_save_dataset(request) + + if result.result_status.is_error: + self.logger.error("SaveDataSet operation failed: %s", result.result_status.message) + else: + self.logger.info("SaveDataSet operation completed successfully for DataSet: %s", request_params.name) + + return result + + # ------------------------------------------------------------------ + # getDataSet + # ------------------------------------------------------------------ + + def _build_get_dataset_request(self, dataset_id: str) -> annotation_pb2.GetDataSetRequest: + """ + Builds a GetDataSetRequest for the supplied DataSet id. + :param dataset_id: Id of the DataSet to retrieve. + :return: A GetDataSetRequest for the specified id. + """ + self.logger.debug("Building GetDataSetRequest for id: %s", dataset_id) + request = annotation_pb2.GetDataSetRequest() + request.dataSetId = dataset_id + return request + + def _send_get_dataset(self, request: annotation_pb2.GetDataSetRequest) -> GetDataSetApiResult: + """ + Invokes the getDataSet() API method with the supplied request. + :param request: GetDataSetRequest with parameters for the call. + :return: A GetDataSetApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.getDataSet, + request, + GetDataSetApiResult, + "getDataSetResult", + "getDataSet", + request_log=lambda: self.logger.info("Calling getDataSet API for id: %s", request.dataSetId), + success_log=lambda _response: self.logger.info( + "Successfully retrieved DataSet for id: %s", request.dataSetId + ), + ) + + def get_dataset(self, dataset_id: str) -> GetDataSetApiResult: + """ + User-facing method for invoking the getDataSet() API method. + + A DataSet that does not exist comes back as a business error, not an empty success. + + :param dataset_id: Id of the DataSet to retrieve. + :return: A GetDataSetApiResult with the method response and status information. + """ + self.logger.info("Starting getDataSet operation for id: %s", dataset_id) + + request = self._build_get_dataset_request(dataset_id) + result = self._send_get_dataset(request) + + if result.result_status.is_error: + self.logger.error("GetDataSet operation failed: %s", result.result_status.message) + else: + self.logger.info("GetDataSet operation completed successfully for id: %s", dataset_id) + + return result + + # ------------------------------------------------------------------ + # queryDataSets + # ------------------------------------------------------------------ + + def _build_query_datasets_request( + self, + criteria: list[annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion] | None = None, + limit: int | None = None, + page_token: str | None = None, + ) -> annotation_pb2.QueryDataSetsRequest: + """ + Builds a QueryDataSetsRequest from the supplied criteria and paging parameters. + :param criteria: List of QueryDataSetsCriterion objects (see DataSetQuery helpers), or None to match all. + :param limit: Maximum number of records to return per page (optional). + :param page_token: Token for retrieving a subsequent page (optional). + :return: A QueryDataSetsRequest for the specified params. + """ + self.logger.debug("Building QueryDataSetsRequest with %d criteria", len(criteria) if criteria else 0) + request = annotation_pb2.QueryDataSetsRequest() + if criteria: + request.criteria.extend(criteria) + if limit is not None: + request.limit = limit + if page_token: + request.pageToken = page_token + return request + + def _send_query_datasets(self, request: annotation_pb2.QueryDataSetsRequest) -> QueryDataSetsApiResult: + """ + Invokes the queryDataSets() API method with the supplied request. + :param request: QueryDataSetsRequest with parameters for the call. + :return: A QueryDataSetsApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.queryDataSets, + request, + QueryDataSetsApiResult, + "dataSetsResult", + "queryDataSets", + request_log=lambda: self.logger.info("Calling queryDataSets API with %d criteria", len(request.criteria)), + success_log=lambda response: self.logger.info( + "QueryDataSets returned %d records", len(response.dataSetsResult.dataSets) + ), + ) + + def query_datasets( + self, + criteria: list[annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion] | None = None, + limit: int | None = None, + page_token: str | None = None, + ) -> QueryDataSetsApiResult: + """ + User-facing method for invoking the queryDataSets() API method. Returns a single page of results; use + iter_datasets() to page through all results transparently. + + An omitted or empty criteria list matches all DataSets. + + :param criteria: List of QueryDataSetsCriterion objects (see DataSetQuery helpers), or None to match all. + :param limit: Maximum number of records to return PER PAGE -- not a cap on the total (optional). + :param page_token: Token for retrieving a subsequent page (optional). + :return: A QueryDataSetsApiResult with a single page of results and status information. + :raises ValueError: if criteria contains more than one text criterion. + """ + criteria = criteria or [] + _check_at_most_one_text_criterion(criteria, "query_datasets()") + + self.logger.info("Starting queryDataSets operation with %d criteria", len(criteria)) + + request = self._build_query_datasets_request(criteria, limit=limit, page_token=page_token) + result = self._send_query_datasets(request) + + if result.result_status.is_error: + self.logger.error("QueryDataSets operation failed: %s", result.result_status.message) + else: + self.logger.info("QueryDataSets operation completed successfully") + + return result + + def iter_datasets( + self, + criteria: list[annotation_pb2.QueryDataSetsRequest.QueryDataSetsCriterion] | None = None, + limit: int | None = None, + ) -> Iterator[annotation_pb2.DataSet]: + """ + Convenience generator that transparently pages through all queryDataSets() results, following the + nextPageToken until the results are exhausted. Yields individual DataSet records. + + Raises RuntimeError if any page returns an error, so callers can distinguish failure from an empty result set. + + :param criteria: List of QueryDataSetsCriterion objects (see DataSetQuery helpers), or None to match all. + :param limit: Maximum number of records to return per page (optional). + :return: An iterator over all matching DataSet records across all pages. + :raises ValueError: if criteria contains more than one text criterion. + :raises RuntimeError: if any page returns an error. + """ + page_token: str | None = None + while True: + result = self.query_datasets(criteria, limit=limit, page_token=page_token) + if result.result_status.is_error: + raise RuntimeError(f"queryDataSets failed during paging: {result.result_status.message}") + + yield from result.datasets + + page_token = result.next_page_token + if not page_token: + break + + def get_datasets(self, ids: list[str]) -> dict[str, annotation_pb2.DataSet]: + """ + Batch-fetches DataSets by id in a single paged query, rather than one getDataSet() call apiece. + + This exists for the common "a page of annotations needs its datasets" case: annotation records carry + dataSetIds rather than the DataSets themselves, so resolving them one at a time is an N+1. + + Ids are deduplicated with their order preserved. An empty ids list returns {} without issuing an RPC, + because callers feed this from annotation dataSetIds lists that may legitimately be empty. Ids that resolve + to nothing are simply absent from the returned dict: a dangling dataSetIds entry is not an error. + + :param ids: DataSet ids to fetch. + :return: A dict mapping DataSet id to DataSet, for those ids that resolved. + :raises RuntimeError: if any page returns an error. + """ + if not ids: + self.logger.debug("get_datasets() called with no ids; returning empty result without an RPC") + return {} + + unique_ids = list(dict.fromkeys(ids)) + self.logger.info("Starting get_datasets batch fetch for %d unique id(s)", len(unique_ids)) + + found = {dataset.id: dataset for dataset in self.iter_datasets([DataSetQuery.ids(unique_ids)])} + + self.logger.info("get_datasets resolved %d of %d requested id(s)", len(found), len(unique_ids)) + return found + + # ------------------------------------------------------------------ + # deleteDataSet + # ------------------------------------------------------------------ + + def _build_delete_dataset_request(self, dataset_id: str) -> annotation_pb2.DeleteDataSetRequest: + """ + Builds a DeleteDataSetRequest for the supplied DataSet id. + :param dataset_id: Id of the DataSet to delete. + :return: A DeleteDataSetRequest for the specified id. + """ + self.logger.debug("Building DeleteDataSetRequest for id: %s", dataset_id) + request = annotation_pb2.DeleteDataSetRequest() + request.dataSetId = dataset_id + return request + + def _send_delete_dataset(self, request: annotation_pb2.DeleteDataSetRequest) -> DeleteDataSetApiResult: + """ + Invokes the deleteDataSet() API method with the supplied request. + :param request: DeleteDataSetRequest with parameters for the call. + :return: A DeleteDataSetApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.deleteDataSet, + request, + DeleteDataSetApiResult, + "deleteDataSetResult", + "deleteDataSet", + request_log=lambda: self.logger.info("Calling deleteDataSet API for id: %s", request.dataSetId), + success_log=lambda _response: self.logger.info("Successfully deleted DataSet id: %s", request.dataSetId), + ) + + def delete_dataset(self, dataset_id: str) -> DeleteDataSetApiResult: + """ + User-facing method for invoking the deleteDataSet() API method. + + The server refuses to delete a DataSet that any annotation still references, and reports which one; there is + deliberately no cascade here. To remove a referenced DataSet, delete its annotations first: + + for annotation in client.annotation.annotations.iter_annotations([AnnotationQuery.datasets([id])]): + client.annotation.annotations.delete_annotation(annotation.id) + client.annotation.datasets.delete_dataset(id) + + Deleting a DataSet that does not exist is likewise a business error, not a silent success. + + :param dataset_id: Id of the DataSet to delete. + :return: A DeleteDataSetApiResult with the method response and status information. + """ + self.logger.info("Starting deleteDataSet operation for id: %s", dataset_id) + + request = self._build_delete_dataset_request(dataset_id) + result = self._send_delete_dataset(request) + + if result.result_status.is_error: + self.logger.error("DeleteDataSet operation failed: %s", result.result_status.message) + else: + self.logger.info("DeleteDataSet operation completed successfully for id: %s", dataset_id) + + return result diff --git a/src/dp_python_lib/client/export_client.py b/src/dp_python_lib/client/export_client.py new file mode 100644 index 0000000..76e5efb --- /dev/null +++ b/src/dp_python_lib/client/export_client.py @@ -0,0 +1,248 @@ +import logging +from enum import Enum + +import grpc + +from dp_python_lib.client.result import ApiResultBase +from dp_python_lib.client.service_api_client_base import ServiceApiClientBase +from dp_python_lib.grpc import annotation_pb2, annotation_pb2_grpc, common_pb2 + + +class ExportFormat(str, Enum): + """ + Output file formats accepted by the exportData() API method. + + A str enum, so ExportFormat.CSV and ExportFormat("csv") are both valid and the member compares equal to its + string value. The proto's EXPORT_FORMAT_UNSPECIFIED zero value is deliberately absent: the server rejects it, + so making it unreachable through this enum turns a would-be server rejection into a client-side ValueError + naming the valid formats. + + Note the tabular formats can only represent scalar columns: exporting data containing array, image, or struct + columns as CSV or XLSX is rejected server-side. Use HDF5 for those. + """ + + HDF5 = "hdf5" + CSV = "csv" + XLSX = "xlsx" + + def to_proto(self) -> "annotation_pb2.ExportDataRequest.ExportOutputFormat": + """ + Converts this format into its protobuf enum value. + :return: The corresponding ExportDataRequest.ExportOutputFormat value. + """ + return _EXPORT_FORMAT_TO_PROTO[self] + + +_EXPORT_FORMAT_TO_PROTO = { + ExportFormat.HDF5: annotation_pb2.ExportDataRequest.ExportOutputFormat.EXPORT_FORMAT_HDF5, + ExportFormat.CSV: annotation_pb2.ExportDataRequest.ExportOutputFormat.EXPORT_FORMAT_CSV, + ExportFormat.XLSX: annotation_pb2.ExportDataRequest.ExportOutputFormat.EXPORT_FORMAT_XLSX, +} + + +def calculations_spec( + calculations_id: str, + frame_columns: dict[str, list[str]] | None = None, +) -> common_pb2.CalculationsSpec: + """ + Builds a CalculationsSpec naming the calculations to include in an export, and optionally which of their columns. + + Omitting frame_columns includes every frame and every column of the named calculations. Supplying it restricts + the export to the named columns of the named frames; frames not mentioned are excluded entirely. + + :param calculations_id: Id of the calculations object to export. + :param frame_columns: Mapping of frame name to the column names to include from that frame. Omit for all. + :return: A common.CalculationsSpec for the specified calculations and column filter. + :raises ValueError: if calculations_id is empty, or if any frame name or column name list is empty. + """ + if not calculations_id: + raise ValueError("calculations_spec() requires a non-empty calculations_id") + + spec = common_pb2.CalculationsSpec() + spec.calculationsId = calculations_id + + if frame_columns: + for frame_name, column_names in frame_columns.items(): + if not frame_name: + raise ValueError("calculations_spec() requires a non-empty name for every frame") + if not column_names: + raise ValueError( + f"calculations_spec() requires a non-empty column name list for frame '{frame_name}'; " + f"omit the frame entirely to exclude it, or omit frame_columns to include all columns" + ) + spec.dataFrameColumns[frame_name].columnNames[:] = column_names + + return spec + + +class ExportDataRequestParams: + """ + Encapsulates client parameters for a call to the exportData() API method. + + At least one data source is required, and the sources merge: a saved DataSet by id, ad-hoc DataBlocks specified + inline (treated by the server as a transient DataSet), and/or a CalculationsSpec. A calculations-only export is + legal and needs no ingested time-series data. + """ + + def __init__( + self, + output_format: "ExportFormat | str", + dataset_id: str | None = None, + data_blocks: list[annotation_pb2.DataBlock] | None = None, + calculations_spec: common_pb2.CalculationsSpec | None = None, + ) -> None: + """ + :param output_format: The export file format, as an ExportFormat or its string value ("hdf5"/"csv"/"xlsx"). + A bare string is coerced through ExportFormat(), so a misspelling raises here rather than server-side. + :param dataset_id: Id of a saved DataSet to export. + :param data_blocks: Ad-hoc time ranges and PV names to export without saving a DataSet (see data_block()). + :param calculations_spec: Calculations to include in the export (see calculations_spec()). + :raises ValueError: if output_format is not a valid format, or if no data source is specified. + """ + try: + self.output_format = ExportFormat(output_format) + except ValueError: + valid = ", ".join(repr(member.value) for member in ExportFormat) + raise ValueError( + f"ExportDataRequestParams received an invalid output_format {output_format!r}; " + f"valid formats are {valid}" + ) from None + + if not dataset_id and not data_blocks and calculations_spec is None: + raise ValueError( + "ExportDataRequestParams requires at least one data source: dataset_id, data_blocks, " + "or calculations_spec" + ) + + self.dataset_id = dataset_id + self.data_blocks = data_blocks + self.calculations_spec = calculations_spec + + +class ExportDataApiResult(ApiResultBase): + """ + Wraps the response from exportData(), with a status object including an error flag and message. + + The exported file lives on the SERVER's filesystem: file_path is a server-side path, and file_url is populated + only when the deployment publishes exports over HTTP. There is no RPC for retrieving the file, so this library + offers no download convenience. + """ + + def __init__( + self, + is_error: bool, + message: str, + response: annotation_pb2.ExportDataResponse | None = None, + ) -> None: + """ + :param is_error: Boolean flag indicating if an error occurred in the API call. + :param message: Error message describing the error condition. + :param response: The ExportDataResponse returned by the API call, or None. + """ + super().__init__(is_error, message) + self.response = response + + @property + def file_path(self) -> str | None: + """Server-side path of the exported file, or None on error.""" + if self.response is not None and self.response.HasField("exportDataResult"): + return self.response.exportDataResult.filePath + return None + + @property + def file_url(self) -> str | None: + """ + URL of the exported file, or None on error. + + Empty string is normal, not a failure: it means the deployment does not publish exports over HTTP. + """ + if self.response is not None and self.response.HasField("exportDataResult"): + return self.response.exportDataResult.fileUrl + return None + + +class ExportClient(ServiceApiClientBase): + """ + User-facing client for the data export method of the MLDP Annotation Service. Exports a saved DataSet, ad-hoc + data blocks, and/or calculations to a file on the server, in HDF5, CSV, or XLSX format. + + Provides a low-level wrapper for exportData(). + """ + + def __init__(self, channel: grpc.Channel) -> None: + """ + :param channel: gRPC communication channel for the Annotation Service. + """ + super().__init__(channel, annotation_pb2_grpc.DpAnnotationServiceStub) + self.logger = logging.getLogger(__name__) + self.logger.debug("ExportClient initialized with channel: %s", channel) + + # ------------------------------------------------------------------ + # exportData + # ------------------------------------------------------------------ + + def _build_export_data_request(self, request_params: ExportDataRequestParams) -> annotation_pb2.ExportDataRequest: + """ + Builds an ExportDataRequest from the supplied ExportDataRequestParams. + :param request_params: User parameters for the call to exportData(). + :return: An ExportDataRequest for the specified params. + """ + self.logger.debug("Building ExportDataRequest with format: %s", request_params.output_format.value) + + request = annotation_pb2.ExportDataRequest() + request.outputFormat = request_params.output_format.to_proto() + + if request_params.dataset_id: + request.dataSetId = request_params.dataset_id + + if request_params.data_blocks: + request.dataBlocks.extend(request_params.data_blocks) + + if request_params.calculations_spec is not None: + request.calculationsSpec.CopyFrom(request_params.calculations_spec) + + self.logger.debug("ExportDataRequest built successfully") + return request + + def _send_export_data(self, request: annotation_pb2.ExportDataRequest) -> ExportDataApiResult: + """ + Invokes the exportData() API method with the supplied request. + :param request: ExportDataRequest with parameters for the call. + :return: An ExportDataApiResult with the method response and status information. + """ + return self._dispatch( + self._stub.exportData, + request, + ExportDataApiResult, + "exportDataResult", + "exportData", + request_log=lambda: self.logger.info( + "Calling exportData API with format: %s", + annotation_pb2.ExportDataRequest.ExportOutputFormat.Name(request.outputFormat), + ), + success_log=lambda response: self.logger.info( + "Successfully exported data to: %s", response.exportDataResult.filePath + ), + ) + + def export_data(self, request_params: ExportDataRequestParams) -> ExportDataApiResult: + """ + User-facing method for invoking the exportData() API method. + + The exported file is written on the server; the result carries a server-side path and, when the deployment + publishes over HTTP, a URL. + + :param request_params: Contains user parameters for the call to exportData(). + :return: An ExportDataApiResult with the method response and status information. + """ + self.logger.info("Starting exportData operation with format: %s", request_params.output_format.value) + + request = self._build_export_data_request(request_params) + result = self._send_export_data(request) + + if result.result_status.is_error: + self.logger.error("ExportData operation failed: %s", result.result_status.message) + else: + self.logger.info("ExportData operation completed successfully") + + return result diff --git a/tests/unit/test_annotation_client.py b/tests/unit/test_annotation_client.py new file mode 100644 index 0000000..e77f934 --- /dev/null +++ b/tests/unit/test_annotation_client.py @@ -0,0 +1,57 @@ +import os +import sys +import unittest +from unittest.mock import Mock + +# Add src directory to path for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../src")) + +from dp_python_lib.client.annotation_client import AnnotationClient +from dp_python_lib.client.annotations_client import AnnotationsClient +from dp_python_lib.client.dataset_client import DataSetClient +from dp_python_lib.client.export_client import ExportClient +from dp_python_lib.client.machine_config_client import MachineConfigClient +from dp_python_lib.client.pv_metadata_client import PvMetadataClient +from dp_python_lib.client.sample_status_client import SampleStatusClient + + +class TestAnnotationClientFacade(unittest.TestCase): + """ + The facade groups every DpAnnotationService feature client under one object sharing one channel. These tests + pin the wiring: a missing or misnamed attribute is otherwise only caught at a call site. + """ + + def setUp(self): + self.channel = Mock() + self.client = AnnotationClient(self.channel) + + def test_exposes_every_feature_client(self): + expected = { + "pv_metadata": PvMetadataClient, + "machine_config": MachineConfigClient, + "sample_status": SampleStatusClient, + "datasets": DataSetClient, + "annotations": AnnotationsClient, + "export": ExportClient, + } + for attribute, client_class in expected.items(): + with self.subTest(attribute=attribute): + self.assertIsInstance(getattr(self.client, attribute), client_class) + + def test_all_feature_clients_share_the_one_channel(self): + for attribute in ("pv_metadata", "machine_config", "sample_status", "datasets", "annotations", "export"): + with self.subTest(attribute=attribute): + self.assertIs(getattr(self.client, attribute)._channel, self.channel) + + def test_each_feature_client_has_its_own_stub(self): + # ServiceApiClientBase creates the stub once at init; the facade's clients must not share one instance. + stubs = [ + self.client.datasets._stub, + self.client.annotations._stub, + self.client.export._stub, + ] + self.assertEqual(len({id(stub) for stub in stubs}), len(stubs)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_annotations_client.py b/tests/unit/test_annotations_client.py new file mode 100644 index 0000000..0253f67 --- /dev/null +++ b/tests/unit/test_annotations_client.py @@ -0,0 +1,715 @@ +import os +import sys +import unittest +from unittest.mock import Mock + +import grpc + +# Add src directory to path for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../src")) + +from assignment_spy import watch_assignments + +from dp_python_lib.client.annotations_client import ( + AnnotationQuery, + AnnotationsClient, + DeleteAnnotationApiResult, + GetAnnotationApiResult, + GetCalculationsApiResult, + QueryAnnotationsApiResult, + SaveAnnotationApiResult, + SaveAnnotationRequestParams, + calculations, +) +from dp_python_lib.grpc import annotation_pb2, common_pb2 + + +def _response_with_field(field_name): + """ + Build a Mock response whose HasField(field) returns True only for field_name. This keeps both the _send_* + oneof check and the *ApiResult property accessors (which also call HasField) consistent. + """ + response = Mock() + response.HasField = Mock(side_effect=lambda field: field == field_name) + return response + + +def _frame(column_name="x_rms"): + """A minimal common.DataFrame carrying one named double column, for calculations construction tests.""" + frame = common_pb2.DataFrame() + column = frame.doubleColumns.add() + column.name = column_name + column.values[:] = [1.0, 2.0] + return frame + + +class TestCalculationsBuilder(unittest.TestCase): + """Unit tests for the calculations() builder.""" + + def test_builds_one_frame_per_entry(self): + result = calculations({"bpm-statistics": _frame(), "rf-statistics": _frame("y_rms")}) + + self.assertEqual(len(result.calculationDataFrames), 2) + self.assertEqual( + {f.name for f in result.calculationDataFrames}, + {"bpm-statistics", "rf-statistics"}, + ) + + def test_carries_frame_content(self): + result = calculations({"f1": _frame("x_rms")}) + + frame = result.calculationDataFrames[0].frame + self.assertEqual([c.name for c in frame.doubleColumns], ["x_rms"]) + self.assertEqual(list(frame.doubleColumns[0].values), [1.0, 2.0]) + + def test_rejects_empty_frames(self): + with self.assertRaises(ValueError) as ctx: + calculations({}) + self.assertIn("at least one frame", str(ctx.exception)) + + def test_rejects_empty_frame_name(self): + with self.assertRaises(ValueError) as ctx: + calculations({"": _frame()}) + self.assertIn("non-empty name", str(ctx.exception)) + + +class TestAnnotationsClientBuildRequests(unittest.TestCase): + """Unit tests for the request-building helpers (no gRPC calls).""" + + def setUp(self): + self.client = AnnotationsClient(Mock()) + + def test_build_save_request_all_fields(self): + calcs = calculations({"f1": _frame()}) + params = SaveAnnotationRequestParams( + name="BPM X RMS", + owner_id="cmcchesney", + dataset_ids=["ds-1", "ds-2"], + annotation_ids=["an-9"], + description="a test annotation", + tags=["reviewed"], + attributes={"runNumber": "4471"}, + modified_by="tester", + calculations=calcs, + annotation_id="an-1", + ) + request = self.client._build_save_annotation_request(params) + + self.assertEqual(request.id, "an-1") + self.assertEqual(request.name, "BPM X RMS") + self.assertEqual(request.ownerId, "cmcchesney") + self.assertEqual(list(request.dataSetIds), ["ds-1", "ds-2"]) + self.assertEqual(list(request.annotationIds), ["an-9"]) + self.assertEqual(request.description, "a test annotation") + self.assertEqual(list(request.tags), ["reviewed"]) + self.assertEqual(request.modifiedBy, "tester") + self.assertEqual({(a.name, a.value) for a in request.attributes}, {("runNumber", "4471")}) + self.assertTrue(request.HasField("calculations")) + self.assertEqual([f.name for f in request.calculations.calculationDataFrames], ["f1"]) + + def test_build_save_request_required_fields_only(self): + params = SaveAnnotationRequestParams(name="n", owner_id="o", dataset_ids=["ds-1"]) + request = self.client._build_save_annotation_request(params) + + self.assertEqual(request.name, "n") + self.assertEqual(request.ownerId, "o") + self.assertEqual(list(request.dataSetIds), ["ds-1"]) + self.assertEqual(request.id, "") + self.assertEqual(list(request.annotationIds), []) + self.assertEqual(request.description, "") + self.assertEqual(list(request.tags), []) + self.assertEqual(len(request.attributes), 0) + self.assertEqual(request.modifiedBy, "") + + def test_build_save_request_omitted_calculations_is_absent(self): + # Absent, not an empty message: omitting calculations on a replace is what CLEARS the stored object, so the + # field must genuinely not be set rather than carry an empty Calculations. + params = SaveAnnotationRequestParams(name="n", owner_id="o", dataset_ids=["ds-1"]) + request = self.client._build_save_annotation_request(params) + self.assertFalse(request.HasField("calculations")) + + def test_build_save_request_omitted_id_is_not_assigned(self): + params = SaveAnnotationRequestParams(name="n", owner_id="o", dataset_ids=["ds-1"]) + with watch_assignments(annotation_pb2, "SaveAnnotationRequest", "id") as assigned: + self.client._build_save_annotation_request(params) + self.assertEqual(assigned, [], "an omitted annotation_id must not be assigned") + + def test_build_get_request(self): + self.assertEqual(self.client._build_get_annotation_request("an-1").annotationId, "an-1") + + def test_build_delete_request(self): + self.assertEqual(self.client._build_delete_annotation_request("an-1").annotationId, "an-1") + + def test_build_get_calculations_request(self): + self.assertEqual(self.client._build_get_calculations_request("calc-1").calculationsId, "calc-1") + + def test_build_query_request(self): + criteria = [AnnotationQuery.tags(["reviewed"]), AnnotationQuery.datasets(["ds-1"])] + request = self.client._build_query_annotations_request(criteria, limit=25, page_token="tok") + + self.assertEqual(len(request.criteria), 2) + self.assertTrue(request.criteria[0].HasField("tagsCriterion")) + self.assertTrue(request.criteria[1].HasField("dataSetsCriterion")) + self.assertEqual(list(request.criteria[1].dataSetsCriterion.dataSetIds), ["ds-1"]) + self.assertEqual(request.limit, 25) + self.assertEqual(request.pageToken, "tok") + + def test_build_query_request_no_criteria(self): + self.assertEqual(len(self.client._build_query_annotations_request().criteria), 0) + + def test_build_query_request_limit_zero_is_set(self): + # limit=0 must be forwarded (distinct from "not provided"); guard against a truthiness regression (#13). + with watch_assignments(annotation_pb2, "QueryAnnotationsRequest", "limit") as assigned: + self.client._build_query_annotations_request([AnnotationQuery.tags(["x"])], limit=0) + self.assertEqual(assigned, [0], "limit=0 must be assigned, not dropped by a truthiness guard") + + def test_build_query_request_limit_omitted_is_unset(self): + with watch_assignments(annotation_pb2, "QueryAnnotationsRequest", "limit") as assigned: + self.client._build_query_annotations_request([AnnotationQuery.tags(["x"])]) + self.assertEqual(assigned, [], "an omitted limit must not be assigned") + + def test_build_query_request_empty_page_token_is_unset(self): + with watch_assignments(annotation_pb2, "QueryAnnotationsRequest", "pageToken") as assigned: + self.client._build_query_annotations_request([AnnotationQuery.tags(["x"])], page_token="") + self.assertEqual(assigned, [], "an empty page token must not be assigned") + + +class TestAnnotationQueryHelpers(unittest.TestCase): + """Unit tests for the AnnotationQuery criterion builders.""" + + def test_ids(self): + c = AnnotationQuery.ids(["a"]) + self.assertTrue(c.HasField("idCriterion")) + self.assertEqual(list(c.idCriterion.ids), ["a"]) + + def test_owners(self): + c = AnnotationQuery.owners(["o1"]) + self.assertTrue(c.HasField("ownerCriterion")) + self.assertEqual(list(c.ownerCriterion.ownerIds), ["o1"]) + + def test_datasets(self): + c = AnnotationQuery.datasets(["ds-1"]) + self.assertTrue(c.HasField("dataSetsCriterion")) + self.assertEqual(list(c.dataSetsCriterion.dataSetIds), ["ds-1"]) + + def test_annotations(self): + c = AnnotationQuery.annotations(["an-1"]) + self.assertTrue(c.HasField("annotationsCriterion")) + self.assertEqual(list(c.annotationsCriterion.annotationIds), ["an-1"]) + + def test_name_all_options(self): + c = AnnotationQuery.name(exact=["A"], prefix=["B"], contains=["C"]) + self.assertTrue(c.HasField("nameCriterion")) + self.assertEqual(list(c.nameCriterion.exact), ["A"]) + self.assertEqual(list(c.nameCriterion.prefix), ["B"]) + self.assertEqual(list(c.nameCriterion.contains), ["C"]) + + def test_text(self): + c = AnnotationQuery.text("rms") + self.assertTrue(c.HasField("textCriterion")) + self.assertEqual(c.textCriterion.text, "rms") + + def test_tags(self): + c = AnnotationQuery.tags(["t1"]) + self.assertTrue(c.HasField("tagsCriterion")) + self.assertEqual(list(c.tagsCriterion.values), ["t1"]) + + def test_attributes_with_values(self): + c = AnnotationQuery.attributes("runNumber", ["4471"]) + self.assertTrue(c.HasField("attributesCriterion")) + self.assertEqual(c.attributesCriterion.key, "runNumber") + self.assertEqual(list(c.attributesCriterion.values), ["4471"]) + + def test_attributes_key_only(self): + for criterion in (AnnotationQuery.attributes("runNumber"), AnnotationQuery.attributes("runNumber", [])): + self.assertTrue(criterion.HasField("attributesCriterion")) + self.assertEqual(criterion.attributesCriterion.key, "runNumber") + self.assertEqual(list(criterion.attributesCriterion.values), []) + + def test_empty_inputs_raise(self): + with self.assertRaises(ValueError): + AnnotationQuery.ids([]) + with self.assertRaises(ValueError): + AnnotationQuery.owners([]) + with self.assertRaises(ValueError): + AnnotationQuery.datasets([]) + with self.assertRaises(ValueError): + AnnotationQuery.annotations([]) + with self.assertRaises(ValueError): + AnnotationQuery.name() + with self.assertRaises(ValueError): + AnnotationQuery.text("") + with self.assertRaises(ValueError): + AnnotationQuery.tags([]) + with self.assertRaises(ValueError): + AnnotationQuery.attributes("") + + +class TestQueryAnnotationsTextCriterionRule(unittest.TestCase): + """Two text criteria cannot be ANDed, so the client rejects them before the RPC.""" + + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.client._stub = Mock() + + def test_two_text_criteria_raise(self): + with self.assertRaises(ValueError) as ctx: + self.client.query_annotations([AnnotationQuery.text("a"), AnnotationQuery.text("b")]) + self.assertIn("at most one text criterion", str(ctx.exception)) + self.client._stub.queryAnnotations.assert_not_called() + + def test_one_text_criterion_is_allowed(self): + response = _response_with_field("annotationsResult") + response.annotationsResult = annotation_pb2.QueryAnnotationsResponse.AnnotationsResult() + self.client._stub.queryAnnotations.return_value = response + + result = self.client.query_annotations([AnnotationQuery.text("a"), AnnotationQuery.tags(["t"])]) + + self.assertFalse(result.result_status.is_error) + + def test_iter_annotations_also_rejects_two_text_criteria(self): + with self.assertRaises(ValueError): + list(self.client.iter_annotations([AnnotationQuery.text("a"), AnnotationQuery.text("b")])) + + +class TestSendSaveAnnotation(unittest.TestCase): + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.request = annotation_pb2.SaveAnnotationRequest(name="n", ownerId="o") + + def test_success(self): + response = _response_with_field("saveAnnotationResult") + response.saveAnnotationResult.annotationId = "an-1" + response.saveAnnotationResult.calculationsId = "calc-1" + mock_stub = Mock() + mock_stub.saveAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_annotation(self.request) + + self.assertIsInstance(result, SaveAnnotationApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.annotation_id, "an-1") + self.assertEqual(result.calculations_id, "calc-1") + mock_stub.saveAnnotation.assert_called_once_with(self.request) + + def test_no_calculations_yields_empty_string_not_none(self): + # "" means the request carried no calculations; None is reserved for a failed call. + response = _response_with_field("saveAnnotationResult") + response.saveAnnotationResult.annotationId = "an-1" + response.saveAnnotationResult.calculationsId = "" + mock_stub = Mock() + mock_stub.saveAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_annotation(self.request) + + self.assertEqual(result.calculations_id, "") + self.assertIsNotNone(result.calculations_id) + + def test_exceptional_result(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "dataSetIds must not be empty" + mock_stub = Mock() + mock_stub.saveAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertEqual(result.result_status.message, "dataSetIds must not be empty") + self.assertIsNone(result.annotation_id) + self.assertIsNone(result.calculations_id) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.saveAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="Connection timeout") + mock_stub.saveAnnotation.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_save_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: Connection timeout", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.saveAnnotation.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_save_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestSendGetAnnotation(unittest.TestCase): + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.request = annotation_pb2.GetAnnotationRequest(annotationId="an-1") + + def _response_with_annotation(self, annotation): + response = _response_with_field("getAnnotationResult") + response.getAnnotationResult.annotation = annotation + return response + + def test_success(self): + annotation = annotation_pb2.Annotation(id="an-1", name="BPM X RMS") + mock_stub = Mock() + mock_stub.getAnnotation.return_value = self._response_with_annotation(annotation) + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertIsInstance(result, GetAnnotationApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.annotation.name, "BPM X RMS") + mock_stub.getAnnotation.assert_called_once_with(self.request) + + def test_calculations_are_returned_inline(self): + annotation = annotation_pb2.Annotation(id="an-1", calculationsId="calc-1") + annotation.calculations.CopyFrom(calculations({"f1": _frame()})) + mock_stub = Mock() + mock_stub.getAnnotation.return_value = self._response_with_annotation(annotation) + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertIsNotNone(result.calculations) + self.assertEqual([f.name for f in result.calculations.calculationDataFrames], ["f1"]) + + def test_calculations_none_when_annotation_has_none(self): + annotation = annotation_pb2.Annotation(id="an-1") + mock_stub = Mock() + mock_stub.getAnnotation.return_value = self._response_with_annotation(annotation) + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertIsNone(result.calculations) + self.assertIsNotNone(result.annotation) + + def test_not_found_is_a_business_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "no Annotation record found for id: an-1" + mock_stub = Mock() + mock_stub.getAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIsNone(result.annotation) + self.assertIsNone(result.calculations) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.getAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.getAnnotation.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.getAnnotation.side_effect = KeyError("bad") + self.client._stub = mock_stub + + result = self.client._send_get_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error", result.result_status.message) + + +class TestSendQueryAnnotations(unittest.TestCase): + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.request = annotation_pb2.QueryAnnotationsRequest() + + def _result(self, names, next_token=""): + result = annotation_pb2.QueryAnnotationsResponse.AnnotationsResult() + for name in names: + result.annotations.append(annotation_pb2.Annotation(name=name)) + result.nextPageToken = next_token + return result + + def test_success(self): + response = _response_with_field("annotationsResult") + response.annotationsResult = self._result(["a", "b"], "tok") + mock_stub = Mock() + mock_stub.queryAnnotations.return_value = response + self.client._stub = mock_stub + + result = self.client._send_query_annotations(self.request) + + self.assertIsInstance(result, QueryAnnotationsApiResult) + self.assertEqual([a.name for a in result.annotations], ["a", "b"]) + self.assertEqual(result.next_page_token, "tok") + + def test_exceptional_result(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "invalid page token" + mock_stub = Mock() + mock_stub.queryAnnotations.return_value = response + self.client._stub = mock_stub + + result = self.client._send_query_annotations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertEqual(result.annotations, []) + self.assertEqual(result.next_page_token, "") + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.queryAnnotations.return_value = response + self.client._stub = mock_stub + + result = self.client._send_query_annotations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.queryAnnotations.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_query_annotations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.queryAnnotations.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_query_annotations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestSendDeleteAnnotation(unittest.TestCase): + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.request = annotation_pb2.DeleteAnnotationRequest(annotationId="an-1") + + def test_success(self): + response = _response_with_field("deleteAnnotationResult") + response.deleteAnnotationResult.annotationId = "an-1" + mock_stub = Mock() + mock_stub.deleteAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_annotation(self.request) + + self.assertIsInstance(result, DeleteAnnotationApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.annotation_id, "an-1") + mock_stub.deleteAnnotation.assert_called_once_with(self.request) + + def test_not_found_is_a_business_error(self): + # Deleting an annotation that does not exist is a REJECT, not a silent success. + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "no Annotation record found for id: an-1" + mock_stub = Mock() + mock_stub.deleteAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("no Annotation record found", result.result_status.message) + self.assertIsNone(result.annotation_id) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.deleteAnnotation.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.deleteAnnotation.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_delete_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.deleteAnnotation.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_delete_annotation(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestSendGetCalculations(unittest.TestCase): + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.request = annotation_pb2.GetCalculationsRequest(calculationsId="calc-1") + + def test_success(self): + response = _response_with_field("getCalculationsResult") + response.getCalculationsResult.calculations = calculations({"f1": _frame()}) + mock_stub = Mock() + mock_stub.getCalculations.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_calculations(self.request) + + self.assertIsInstance(result, GetCalculationsApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual([f.name for f in result.calculations.calculationDataFrames], ["f1"]) + mock_stub.getCalculations.assert_called_once_with(self.request) + + def test_not_found_is_a_business_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "no Calculations record found for id: calc-1" + mock_stub = Mock() + mock_stub.getCalculations.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_calculations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIsNone(result.calculations) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.getCalculations.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_calculations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.getCalculations.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_get_calculations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.getCalculations.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_get_calculations(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestIterAnnotations(unittest.TestCase): + def setUp(self): + self.client = AnnotationsClient(Mock()) + self.criteria = [AnnotationQuery.tags(["reviewed"])] + + def _page(self, names, next_token): + response = _response_with_field("annotationsResult") + result = annotation_pb2.QueryAnnotationsResponse.AnnotationsResult() + for name in names: + result.annotations.append(annotation_pb2.Annotation(name=name)) + result.nextPageToken = next_token + response.annotationsResult = result + return response + + def test_pages_through_all_results(self): + mock_stub = Mock() + mock_stub.queryAnnotations.side_effect = [self._page(["a", "b"], "tok1"), self._page(["c"], "")] + self.client._stub = mock_stub + + names = [a.name for a in self.client.iter_annotations(self.criteria, limit=2)] + + self.assertEqual(names, ["a", "b", "c"]) + self.assertEqual(mock_stub.queryAnnotations.call_count, 2) + second_request = mock_stub.queryAnnotations.call_args_list[1].args[0] + self.assertEqual(second_request.pageToken, "tok1") + + def test_single_page(self): + mock_stub = Mock() + mock_stub.queryAnnotations.return_value = self._page(["a"], "") + self.client._stub = mock_stub + + self.assertEqual([a.name for a in self.client.iter_annotations(self.criteria)], ["a"]) + self.assertEqual(mock_stub.queryAnnotations.call_count, 1) + + def test_no_criteria_matches_all(self): + mock_stub = Mock() + mock_stub.queryAnnotations.return_value = self._page(["a"], "") + self.client._stub = mock_stub + + self.assertEqual([a.name for a in self.client.iter_annotations()], ["a"]) + self.assertEqual(len(mock_stub.queryAnnotations.call_args_list[0].args[0].criteria), 0) + + def test_page_error_raises_runtime_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "page token rejected" + mock_stub = Mock() + mock_stub.queryAnnotations.return_value = response + self.client._stub = mock_stub + + with self.assertRaises(RuntimeError) as ctx: + list(self.client.iter_annotations(self.criteria)) + self.assertIn("page token rejected", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_dataset_client.py b/tests/unit/test_dataset_client.py new file mode 100644 index 0000000..e893d11 --- /dev/null +++ b/tests/unit/test_dataset_client.py @@ -0,0 +1,682 @@ +import os +import sys +import unittest +from datetime import datetime, timezone +from unittest.mock import Mock + +import grpc + +# Add src directory to path for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../src")) + +from assignment_spy import watch_assignments + +from dp_python_lib.client.dataset_client import ( + DataSetClient, + DataSetQuery, + DeleteDataSetApiResult, + GetDataSetApiResult, + QueryDataSetsApiResult, + SaveDataSetApiResult, + SaveDataSetRequestParams, + data_block, +) +from dp_python_lib.grpc import annotation_pb2 + +BEGIN = datetime(2026, 7, 14, 18, tzinfo=timezone.utc) +END = datetime(2026, 7, 14, 19, tzinfo=timezone.utc) + + +def _response_with_field(field_name): + """ + Build a Mock response whose HasField(field) returns True only for field_name. This keeps both the _send_* + oneof check and the *ApiResult property accessors (which also call HasField) consistent. + """ + response = Mock() + response.HasField = Mock(side_effect=lambda field: field == field_name) + return response + + +class TestDataBlockBuilder(unittest.TestCase): + """Unit tests for the data_block() builder.""" + + def test_builds_block(self): + block = data_block(BEGIN, END, ["A:1", "A:2"]) + self.assertEqual(block.beginTime.epochSeconds, int(BEGIN.timestamp())) + self.assertEqual(block.endTime.epochSeconds, int(END.timestamp())) + self.assertEqual(list(block.pvNames), ["A:1", "A:2"]) + + def test_accepts_epoch_seconds_and_timestamps(self): + block = data_block(100, 200, ["A:1"]) + self.assertEqual(block.beginTime.epochSeconds, 100) + self.assertEqual(block.endTime.epochSeconds, 200) + + def test_rejects_empty_pv_names(self): + with self.assertRaises(ValueError) as ctx: + data_block(BEGIN, END, []) + self.assertIn("pv_names", str(ctx.exception)) + + def test_rejects_reversed_range(self): + # The server validates only that each bound is non-zero, so this check is the only one there is. + with self.assertRaises(ValueError) as ctx: + data_block(END, BEGIN, ["A:1"]) + self.assertIn("strictly before", str(ctx.exception)) + + def test_rejects_equal_bounds(self): + with self.assertRaises(ValueError): + data_block(BEGIN, BEGIN, ["A:1"]) + + def test_rejects_reversed_range_at_nanosecond_precision(self): + # Equal seconds, decreasing nanoseconds -- caught only if the comparison includes the nanos field. + with self.assertRaises(ValueError): + data_block(100.000_002, 100.000_001, ["A:1"]) + + def test_rejects_naive_datetime(self): + with self.assertRaises(ValueError): + data_block(datetime(2026, 7, 14), END, ["A:1"]) # noqa: DTZ001 -- naive input is the condition under test + + +class TestDataSetClientBuildRequests(unittest.TestCase): + """Unit tests for the request-building helpers (no gRPC calls).""" + + def setUp(self): + self.mock_channel = Mock() + self.client = DataSetClient(self.mock_channel) + self.block = data_block(BEGIN, END, ["A:1"]) + + def test_build_save_request_all_fields(self): + params = SaveDataSetRequestParams( + name="ramp study", + owner_id="cmcchesney", + data_blocks=[self.block], + description="a test dataset", + tags=["ramp-study", "reviewed"], + attributes={"runNumber": "4471", "station": "gunb"}, + modified_by="tester", + dataset_id="ds-1", + ) + request = self.client._build_save_dataset_request(params) + + self.assertEqual(request.id, "ds-1") + self.assertEqual(request.name, "ramp study") + self.assertEqual(request.ownerId, "cmcchesney") + self.assertEqual(request.description, "a test dataset") + self.assertEqual(list(request.tags), ["ramp-study", "reviewed"]) + self.assertEqual(request.modifiedBy, "tester") + self.assertEqual(len(request.dataBlocks), 1) + self.assertEqual(list(request.dataBlocks[0].pvNames), ["A:1"]) + self.assertEqual( + {(a.name, a.value) for a in request.attributes}, + {("runNumber", "4471"), ("station", "gunb")}, + ) + + def test_build_save_request_required_fields_only(self): + params = SaveDataSetRequestParams(name="n", owner_id="o", data_blocks=[self.block]) + request = self.client._build_save_dataset_request(params) + + self.assertEqual(request.name, "n") + self.assertEqual(request.ownerId, "o") + self.assertEqual(request.id, "") + self.assertEqual(request.description, "") + self.assertEqual(list(request.tags), []) + self.assertEqual(len(request.attributes), 0) + self.assertEqual(request.modifiedBy, "") + + def test_build_save_request_omitted_id_is_not_assigned(self): + # An omitted dataset_id must leave 'id' unassigned, since a present id means "replace in full". + params = SaveDataSetRequestParams(name="n", owner_id="o", data_blocks=[self.block]) + with watch_assignments(annotation_pb2, "SaveDataSetRequest", "id") as assigned: + self.client._build_save_dataset_request(params) + self.assertEqual(assigned, [], "an omitted dataset_id must not be assigned") + + def test_build_get_request(self): + request = self.client._build_get_dataset_request("ds-1") + self.assertEqual(request.dataSetId, "ds-1") + + def test_build_delete_request(self): + request = self.client._build_delete_dataset_request("ds-1") + self.assertEqual(request.dataSetId, "ds-1") + + def test_build_query_request(self): + criteria = [DataSetQuery.owners(["cmcchesney"]), DataSetQuery.tags(["ramp-study"])] + request = self.client._build_query_datasets_request(criteria, limit=25, page_token="tok") + + self.assertEqual(len(request.criteria), 2) + self.assertTrue(request.criteria[0].HasField("ownerCriterion")) + self.assertEqual(list(request.criteria[0].ownerCriterion.ownerIds), ["cmcchesney"]) + self.assertTrue(request.criteria[1].HasField("tagsCriterion")) + self.assertEqual(request.limit, 25) + self.assertEqual(request.pageToken, "tok") + + def test_build_query_request_no_criteria(self): + # An empty criteria list is a legal match-all, not an error. + request = self.client._build_query_datasets_request() + self.assertEqual(len(request.criteria), 0) + + def test_build_query_request_limit_zero_is_set(self): + # limit=0 must be forwarded (distinct from "not provided"); guard against a truthiness regression (#13). + # A proto3 scalar reads 0 whether or not it was assigned, so watch the assignment itself. + with watch_assignments(annotation_pb2, "QueryDataSetsRequest", "limit") as assigned: + self.client._build_query_datasets_request([DataSetQuery.tags(["x"])], limit=0) + self.assertEqual(assigned, [0], "limit=0 must be assigned, not dropped by a truthiness guard") + + def test_build_query_request_limit_omitted_is_unset(self): + with watch_assignments(annotation_pb2, "QueryDataSetsRequest", "limit") as assigned: + self.client._build_query_datasets_request([DataSetQuery.tags(["x"])]) + self.assertEqual(assigned, [], "an omitted limit must not be assigned") + + def test_build_query_request_empty_page_token_is_unset(self): + with watch_assignments(annotation_pb2, "QueryDataSetsRequest", "pageToken") as assigned: + self.client._build_query_datasets_request([DataSetQuery.tags(["x"])], page_token="") + self.assertEqual(assigned, [], "an empty page token must not be assigned") + + +class TestDataSetQueryHelpers(unittest.TestCase): + """Unit tests for the DataSetQuery criterion builders.""" + + def test_ids(self): + c = DataSetQuery.ids(["a", "b"]) + self.assertTrue(c.HasField("idCriterion")) + self.assertEqual(list(c.idCriterion.ids), ["a", "b"]) + + def test_owners(self): + c = DataSetQuery.owners(["o1"]) + self.assertTrue(c.HasField("ownerCriterion")) + self.assertEqual(list(c.ownerCriterion.ownerIds), ["o1"]) + + def test_name_all_options(self): + c = DataSetQuery.name(exact=["A"], prefix=["B"], contains=["C"]) + self.assertTrue(c.HasField("nameCriterion")) + self.assertEqual(list(c.nameCriterion.exact), ["A"]) + self.assertEqual(list(c.nameCriterion.prefix), ["B"]) + self.assertEqual(list(c.nameCriterion.contains), ["C"]) + + def test_text(self): + c = DataSetQuery.text("ramp") + self.assertTrue(c.HasField("textCriterion")) + self.assertEqual(c.textCriterion.text, "ramp") + + def test_pv_names(self): + c = DataSetQuery.pv_names(["A:1"]) + self.assertTrue(c.HasField("pvNameCriterion")) + self.assertEqual(list(c.pvNameCriterion.names), ["A:1"]) + + def test_tags(self): + c = DataSetQuery.tags(["t1", "t2"]) + self.assertTrue(c.HasField("tagsCriterion")) + self.assertEqual(list(c.tagsCriterion.values), ["t1", "t2"]) + + def test_attributes_with_values(self): + c = DataSetQuery.attributes("runNumber", ["4471"]) + self.assertTrue(c.HasField("attributesCriterion")) + self.assertEqual(c.attributesCriterion.key, "runNumber") + self.assertEqual(list(c.attributesCriterion.values), ["4471"]) + + def test_attributes_key_only(self): + # Unlike the older helpers, an absent/empty values list is legal: it is a key-only existence search. + for criterion in (DataSetQuery.attributes("runNumber"), DataSetQuery.attributes("runNumber", [])): + self.assertTrue(criterion.HasField("attributesCriterion")) + self.assertEqual(criterion.attributesCriterion.key, "runNumber") + self.assertEqual(list(criterion.attributesCriterion.values), []) + + def test_empty_inputs_raise(self): + with self.assertRaises(ValueError): + DataSetQuery.ids([]) + with self.assertRaises(ValueError): + DataSetQuery.owners([]) + with self.assertRaises(ValueError): + DataSetQuery.name() + with self.assertRaises(ValueError): + DataSetQuery.text("") + with self.assertRaises(ValueError): + DataSetQuery.pv_names([]) + with self.assertRaises(ValueError): + DataSetQuery.tags([]) + with self.assertRaises(ValueError): + DataSetQuery.attributes("") + + +class TestQueryDataSetsTextCriterionRule(unittest.TestCase): + """Two text criteria cannot be ANDed, so the client rejects them before the RPC.""" + + def setUp(self): + self.client = DataSetClient(Mock()) + self.client._stub = Mock() + + def test_two_text_criteria_raise(self): + with self.assertRaises(ValueError) as ctx: + self.client.query_datasets([DataSetQuery.text("a"), DataSetQuery.text("b")]) + self.assertIn("at most one text criterion", str(ctx.exception)) + self.client._stub.queryDataSets.assert_not_called() + + def test_one_text_criterion_is_allowed(self): + response = _response_with_field("dataSetsResult") + response.dataSetsResult = annotation_pb2.QueryDataSetsResponse.DataSetsResult() + self.client._stub.queryDataSets.return_value = response + + result = self.client.query_datasets([DataSetQuery.text("a"), DataSetQuery.tags(["t"])]) + + self.assertFalse(result.result_status.is_error) + + def test_iter_datasets_also_rejects_two_text_criteria(self): + with self.assertRaises(ValueError): + list(self.client.iter_datasets([DataSetQuery.text("a"), DataSetQuery.text("b")])) + + +class TestSendSaveDataSet(unittest.TestCase): + def setUp(self): + self.client = DataSetClient(Mock()) + self.request = annotation_pb2.SaveDataSetRequest(name="n", ownerId="o") + + def test_success(self): + response = _response_with_field("saveDataSetResult") + response.saveDataSetResult.dataSetId = "ds-1" + mock_stub = Mock() + mock_stub.saveDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_dataset(self.request) + + self.assertIsInstance(result, SaveDataSetApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.dataset_id, "ds-1") + mock_stub.saveDataSet.assert_called_once_with(self.request) + + def test_exceptional_result(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "name must be specified" + mock_stub = Mock() + mock_stub.saveDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertEqual(result.result_status.message, "name must be specified") + self.assertIsNone(result.dataset_id) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.saveDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_save_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="Connection timeout") + mock_stub.saveDataSet.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_save_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: Connection timeout", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.saveDataSet.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_save_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestSendGetDataSet(unittest.TestCase): + def setUp(self): + self.client = DataSetClient(Mock()) + self.request = annotation_pb2.GetDataSetRequest(dataSetId="ds-1") + + def test_success(self): + response = _response_with_field("getDataSetResult") + response.getDataSetResult.dataSet = annotation_pb2.DataSet(id="ds-1", name="ramp study") + mock_stub = Mock() + mock_stub.getDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_dataset(self.request) + + self.assertIsInstance(result, GetDataSetApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.dataset.name, "ramp study") + mock_stub.getDataSet.assert_called_once_with(self.request) + + def test_not_found_is_a_business_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "no DataSet record found for id: ds-1" + mock_stub = Mock() + mock_stub.getDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("no DataSet record found", result.result_status.message) + self.assertIsNone(result.dataset) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.getDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_get_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.getDataSet.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_get_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.getDataSet.side_effect = KeyError("bad") + self.client._stub = mock_stub + + result = self.client._send_get_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error", result.result_status.message) + + +class TestSendQueryDataSets(unittest.TestCase): + def setUp(self): + self.client = DataSetClient(Mock()) + self.request = annotation_pb2.QueryDataSetsRequest() + + def _result(self, names, next_token=""): + result = annotation_pb2.QueryDataSetsResponse.DataSetsResult() + for name in names: + result.dataSets.append(annotation_pb2.DataSet(name=name)) + result.nextPageToken = next_token + return result + + def test_success(self): + response = _response_with_field("dataSetsResult") + response.dataSetsResult = self._result(["a", "b"], "tok") + mock_stub = Mock() + mock_stub.queryDataSets.return_value = response + self.client._stub = mock_stub + + result = self.client._send_query_datasets(self.request) + + self.assertIsInstance(result, QueryDataSetsApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual([d.name for d in result.datasets], ["a", "b"]) + self.assertEqual(result.next_page_token, "tok") + + def test_exceptional_result(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "invalid page token" + mock_stub = Mock() + mock_stub.queryDataSets.return_value = response + self.client._stub = mock_stub + + result = self.client._send_query_datasets(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertEqual(result.datasets, []) + self.assertEqual(result.next_page_token, "") + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.queryDataSets.return_value = response + self.client._stub = mock_stub + + result = self.client._send_query_datasets(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.queryDataSets.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_query_datasets(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.queryDataSets.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_query_datasets(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestSendDeleteDataSet(unittest.TestCase): + def setUp(self): + self.client = DataSetClient(Mock()) + self.request = annotation_pb2.DeleteDataSetRequest(dataSetId="ds-1") + + def test_success(self): + response = _response_with_field("deleteDataSetResult") + response.deleteDataSetResult.dataSetId = "ds-1" + mock_stub = Mock() + mock_stub.deleteDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_dataset(self.request) + + self.assertIsInstance(result, DeleteDataSetApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.dataset_id, "ds-1") + mock_stub.deleteDataSet.assert_called_once_with(self.request) + + def test_referenced_dataset_is_a_business_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "DataSet ds-1 is referenced by annotation an-1 (3 total)" + mock_stub = Mock() + mock_stub.deleteDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("referenced by annotation", result.result_status.message) + self.assertIsNone(result.dataset_id) + + def test_not_found_is_a_business_error(self): + # Deleting a DataSet that does not exist is a REJECT, not a silent success. + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "no DataSet record found for id: ds-1" + mock_stub = Mock() + mock_stub.deleteDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIsNone(result.dataset_id) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.deleteDataSet.return_value = response + self.client._stub = mock_stub + + result = self.client._send_delete_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="unavailable") + mock_stub.deleteDataSet.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_delete_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: unavailable", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.deleteDataSet.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_delete_dataset(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestIterDataSets(unittest.TestCase): + def setUp(self): + self.client = DataSetClient(Mock()) + self.criteria = [DataSetQuery.tags(["ramp-study"])] + + def _page(self, names, next_token): + response = _response_with_field("dataSetsResult") + result = annotation_pb2.QueryDataSetsResponse.DataSetsResult() + for name in names: + result.dataSets.append(annotation_pb2.DataSet(name=name)) + result.nextPageToken = next_token + response.dataSetsResult = result + return response + + def test_pages_through_all_results(self): + mock_stub = Mock() + mock_stub.queryDataSets.side_effect = [self._page(["a", "b"], "tok1"), self._page(["c"], "")] + self.client._stub = mock_stub + + names = [d.name for d in self.client.iter_datasets(self.criteria, limit=2)] + + self.assertEqual(names, ["a", "b", "c"]) + self.assertEqual(mock_stub.queryDataSets.call_count, 2) + second_request = mock_stub.queryDataSets.call_args_list[1].args[0] + self.assertEqual(second_request.pageToken, "tok1") + + def test_single_page(self): + mock_stub = Mock() + mock_stub.queryDataSets.return_value = self._page(["a"], "") + self.client._stub = mock_stub + + self.assertEqual([d.name for d in self.client.iter_datasets(self.criteria)], ["a"]) + self.assertEqual(mock_stub.queryDataSets.call_count, 1) + + def test_no_criteria_matches_all(self): + mock_stub = Mock() + mock_stub.queryDataSets.return_value = self._page(["a"], "") + self.client._stub = mock_stub + + self.assertEqual([d.name for d in self.client.iter_datasets()], ["a"]) + self.assertEqual(len(mock_stub.queryDataSets.call_args_list[0].args[0].criteria), 0) + + def test_page_error_raises_runtime_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "page token rejected" + mock_stub = Mock() + mock_stub.queryDataSets.return_value = response + self.client._stub = mock_stub + + with self.assertRaises(RuntimeError) as ctx: + list(self.client.iter_datasets(self.criteria)) + self.assertIn("page token rejected", str(ctx.exception)) + + +class TestGetDataSetsBatch(unittest.TestCase): + """Unit tests for the get_datasets() batch fetch (D9).""" + + def setUp(self): + self.client = DataSetClient(Mock()) + + def _page(self, ids, next_token=""): + response = _response_with_field("dataSetsResult") + result = annotation_pb2.QueryDataSetsResponse.DataSetsResult() + for dataset_id in ids: + result.dataSets.append(annotation_pb2.DataSet(id=dataset_id, name=f"name-{dataset_id}")) + result.nextPageToken = next_token + response.dataSetsResult = result + return response + + def test_returns_dict_keyed_by_id(self): + mock_stub = Mock() + mock_stub.queryDataSets.return_value = self._page(["a", "b"]) + self.client._stub = mock_stub + + found = self.client.get_datasets(["a", "b"]) + + self.assertEqual(set(found), {"a", "b"}) + self.assertEqual(found["a"].name, "name-a") + + def test_empty_ids_makes_no_rpc(self): + mock_stub = Mock() + self.client._stub = mock_stub + + self.assertEqual(self.client.get_datasets([]), {}) + mock_stub.queryDataSets.assert_not_called() + + def test_deduplicates_ids_preserving_order(self): + mock_stub = Mock() + mock_stub.queryDataSets.return_value = self._page(["a", "b"]) + self.client._stub = mock_stub + + self.client.get_datasets(["a", "b", "a", "b", "a"]) + + request = mock_stub.queryDataSets.call_args_list[0].args[0] + self.assertEqual(list(request.criteria[0].idCriterion.ids), ["a", "b"]) + + def test_unresolved_ids_are_simply_absent(self): + # A dangling dataSetIds entry is not an error; it just does not appear in the result. + mock_stub = Mock() + mock_stub.queryDataSets.return_value = self._page(["a"]) + self.client._stub = mock_stub + + found = self.client.get_datasets(["a", "missing"]) + + self.assertEqual(set(found), {"a"}) + + def test_pages_through_results(self): + mock_stub = Mock() + mock_stub.queryDataSets.side_effect = [self._page(["a"], "tok"), self._page(["b"], "")] + self.client._stub = mock_stub + + found = self.client.get_datasets(["a", "b"]) + + self.assertEqual(set(found), {"a", "b"}) + self.assertEqual(mock_stub.queryDataSets.call_count, 2) + + def test_page_error_raises_runtime_error(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "boom" + mock_stub = Mock() + mock_stub.queryDataSets.return_value = response + self.client._stub = mock_stub + + with self.assertRaises(RuntimeError): + self.client.get_datasets(["a"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_export_client.py b/tests/unit/test_export_client.py new file mode 100644 index 0000000..ffc2fcd --- /dev/null +++ b/tests/unit/test_export_client.py @@ -0,0 +1,305 @@ +import os +import sys +import unittest +from datetime import datetime, timezone +from unittest.mock import Mock + +import grpc + +# Add src directory to path for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../src")) + +from dp_python_lib.client.dataset_client import data_block +from dp_python_lib.client.export_client import ( + ExportClient, + ExportDataApiResult, + ExportDataRequestParams, + ExportFormat, + calculations_spec, +) +from dp_python_lib.grpc import annotation_pb2 + +BEGIN = datetime(2026, 7, 14, 18, tzinfo=timezone.utc) +END = datetime(2026, 7, 14, 19, tzinfo=timezone.utc) + +_FORMAT = annotation_pb2.ExportDataRequest.ExportOutputFormat + + +def _response_with_field(field_name): + """ + Build a Mock response whose HasField(field) returns True only for field_name. This keeps both the _send_* + oneof check and the *ApiResult property accessors (which also call HasField) consistent. + """ + response = Mock() + response.HasField = Mock(side_effect=lambda field: field == field_name) + return response + + +class TestExportFormat(unittest.TestCase): + """Unit tests for the ExportFormat enum.""" + + def test_maps_to_proto_values(self): + self.assertEqual(ExportFormat.HDF5.to_proto(), _FORMAT.EXPORT_FORMAT_HDF5) + self.assertEqual(ExportFormat.CSV.to_proto(), _FORMAT.EXPORT_FORMAT_CSV) + self.assertEqual(ExportFormat.XLSX.to_proto(), _FORMAT.EXPORT_FORMAT_XLSX) + + def test_constructible_from_string_value(self): + self.assertIs(ExportFormat("csv"), ExportFormat.CSV) + self.assertIs(ExportFormat("hdf5"), ExportFormat.HDF5) + self.assertIs(ExportFormat("xlsx"), ExportFormat.XLSX) + + def test_compares_equal_to_its_string(self): + self.assertEqual(ExportFormat.CSV, "csv") + + def test_unspecified_is_unreachable(self): + # EXPORT_FORMAT_UNSPECIFIED is rejected server-side; it must have no enum member here. + self.assertNotIn("EXPORT_FORMAT_UNSPECIFIED", {member.name for member in ExportFormat}) + self.assertNotIn(_FORMAT.EXPORT_FORMAT_UNSPECIFIED, {member.to_proto() for member in ExportFormat}) + + +class TestCalculationsSpecBuilder(unittest.TestCase): + """Unit tests for the calculations_spec() builder.""" + + def test_id_only_means_all_frames_and_columns(self): + spec = calculations_spec("calc-1") + self.assertEqual(spec.calculationsId, "calc-1") + self.assertEqual(len(spec.dataFrameColumns), 0) + + def test_with_frame_columns(self): + spec = calculations_spec("calc-1", {"f1": ["x_rms", "y_rms"], "f2": ["z"]}) + self.assertEqual(spec.calculationsId, "calc-1") + self.assertEqual(set(spec.dataFrameColumns), {"f1", "f2"}) + self.assertEqual(list(spec.dataFrameColumns["f1"].columnNames), ["x_rms", "y_rms"]) + self.assertEqual(list(spec.dataFrameColumns["f2"].columnNames), ["z"]) + + def test_rejects_empty_calculations_id(self): + with self.assertRaises(ValueError) as ctx: + calculations_spec("") + self.assertIn("calculations_id", str(ctx.exception)) + + def test_rejects_empty_frame_name(self): + with self.assertRaises(ValueError) as ctx: + calculations_spec("calc-1", {"": ["x"]}) + self.assertIn("non-empty name", str(ctx.exception)) + + def test_rejects_empty_column_name_list(self): + # The server rejects an empty list, so name the alternatives here rather than let it bounce. + with self.assertRaises(ValueError) as ctx: + calculations_spec("calc-1", {"f1": []}) + self.assertIn("f1", str(ctx.exception)) + + +class TestExportDataRequestParams(unittest.TestCase): + """Unit tests for the params class's validation and format coercion.""" + + def setUp(self): + self.block = data_block(BEGIN, END, ["A:1"]) + + def test_accepts_enum_member(self): + params = ExportDataRequestParams(ExportFormat.HDF5, dataset_id="ds-1") + self.assertIs(params.output_format, ExportFormat.HDF5) + + def test_coerces_bare_string(self): + params = ExportDataRequestParams("csv", dataset_id="ds-1") + self.assertIs(params.output_format, ExportFormat.CSV) + + def test_rejects_unknown_format(self): + with self.assertRaises(ValueError) as ctx: + ExportDataRequestParams("parquet", dataset_id="ds-1") + message = str(ctx.exception) + self.assertIn("parquet", message) + self.assertIn("'csv'", message) + self.assertIn("'hdf5'", message) + self.assertIn("'xlsx'", message) + + def test_rejects_wrong_case_format(self): + # The enum values are lowercase; "CSV" is not a member. + with self.assertRaises(ValueError): + ExportDataRequestParams("CSV", dataset_id="ds-1") + + def test_rejects_no_data_source(self): + with self.assertRaises(ValueError) as ctx: + ExportDataRequestParams(ExportFormat.CSV) + self.assertIn("at least one data source", str(ctx.exception)) + + def test_accepts_each_source_alone(self): + ExportDataRequestParams(ExportFormat.CSV, dataset_id="ds-1") + ExportDataRequestParams(ExportFormat.CSV, data_blocks=[self.block]) + ExportDataRequestParams(ExportFormat.CSV, calculations_spec=calculations_spec("calc-1")) + + def test_accepts_merged_sources(self): + params = ExportDataRequestParams( + ExportFormat.HDF5, + dataset_id="ds-1", + data_blocks=[self.block], + calculations_spec=calculations_spec("calc-1"), + ) + self.assertEqual(params.dataset_id, "ds-1") + self.assertEqual(len(params.data_blocks), 1) + self.assertEqual(params.calculations_spec.calculationsId, "calc-1") + + def test_empty_data_blocks_list_is_not_a_source(self): + with self.assertRaises(ValueError): + ExportDataRequestParams(ExportFormat.CSV, data_blocks=[]) + + +class TestExportClientBuildRequest(unittest.TestCase): + """Unit tests for the request-building helper (no gRPC calls).""" + + def setUp(self): + self.client = ExportClient(Mock()) + self.block = data_block(BEGIN, END, ["A:1"]) + + def test_build_request_all_sources(self): + params = ExportDataRequestParams( + ExportFormat.HDF5, + dataset_id="ds-1", + data_blocks=[self.block], + calculations_spec=calculations_spec("calc-1", {"f1": ["x_rms"]}), + ) + request = self.client._build_export_data_request(params) + + self.assertEqual(request.outputFormat, _FORMAT.EXPORT_FORMAT_HDF5) + self.assertEqual(request.dataSetId, "ds-1") + self.assertEqual(len(request.dataBlocks), 1) + self.assertEqual(list(request.dataBlocks[0].pvNames), ["A:1"]) + self.assertTrue(request.HasField("calculationsSpec")) + self.assertEqual(request.calculationsSpec.calculationsId, "calc-1") + self.assertEqual(list(request.calculationsSpec.dataFrameColumns["f1"].columnNames), ["x_rms"]) + + def test_build_request_dataset_only(self): + params = ExportDataRequestParams(ExportFormat.CSV, dataset_id="ds-1") + request = self.client._build_export_data_request(params) + + self.assertEqual(request.outputFormat, _FORMAT.EXPORT_FORMAT_CSV) + self.assertEqual(request.dataSetId, "ds-1") + self.assertEqual(len(request.dataBlocks), 0) + self.assertFalse(request.HasField("calculationsSpec")) + + def test_build_request_calculations_only(self): + # A calculations-only export is legal and needs no ingested time-series data. + params = ExportDataRequestParams(ExportFormat.CSV, calculations_spec=calculations_spec("calc-1")) + request = self.client._build_export_data_request(params) + + self.assertEqual(request.dataSetId, "") + self.assertEqual(len(request.dataBlocks), 0) + self.assertTrue(request.HasField("calculationsSpec")) + + def test_build_request_blocks_only(self): + params = ExportDataRequestParams(ExportFormat.XLSX, data_blocks=[self.block]) + request = self.client._build_export_data_request(params) + + self.assertEqual(request.outputFormat, _FORMAT.EXPORT_FORMAT_XLSX) + self.assertEqual(request.dataSetId, "") + self.assertEqual(len(request.dataBlocks), 1) + self.assertFalse(request.HasField("calculationsSpec")) + + +class TestSendExportData(unittest.TestCase): + def setUp(self): + self.client = ExportClient(Mock()) + self.request = annotation_pb2.ExportDataRequest(dataSetId="ds-1", outputFormat=_FORMAT.EXPORT_FORMAT_CSV) + + def test_success(self): + response = _response_with_field("exportDataResult") + response.exportDataResult.filePath = "/srv/exports/out.csv" + response.exportDataResult.fileUrl = "https://example.test/exports/out.csv" + mock_stub = Mock() + mock_stub.exportData.return_value = response + self.client._stub = mock_stub + + result = self.client._send_export_data(self.request) + + self.assertIsInstance(result, ExportDataApiResult) + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.file_path, "/srv/exports/out.csv") + self.assertEqual(result.file_url, "https://example.test/exports/out.csv") + mock_stub.exportData.assert_called_once_with(self.request) + + def test_empty_file_url_is_normal(self): + # An empty fileUrl means the deployment does not publish over HTTP -- not a failure. + response = _response_with_field("exportDataResult") + response.exportDataResult.filePath = "/srv/exports/out.csv" + response.exportDataResult.fileUrl = "" + mock_stub = Mock() + mock_stub.exportData.return_value = response + self.client._stub = mock_stub + + result = self.client._send_export_data(self.request) + + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.file_url, "") + self.assertIsNotNone(result.file_url) + + def test_exceptional_result(self): + response = _response_with_field("exceptionalResult") + response.exceptionalResult.message = "non-scalar column cannot be exported as CSV" + mock_stub = Mock() + mock_stub.exportData.return_value = response + self.client._stub = mock_stub + + result = self.client._send_export_data(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("non-scalar column", result.result_status.message) + self.assertIsNone(result.file_path) + self.assertIsNone(result.file_url) + + def test_unexpected_response(self): + response = Mock() + response.HasField = Mock(return_value=False) + mock_stub = Mock() + mock_stub.exportData.return_value = response + self.client._stub = mock_stub + + result = self.client._send_export_data(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected response format", result.result_status.message) + + def test_grpc_error(self): + mock_stub = Mock() + err = grpc.RpcError() + err.details = Mock(return_value="Connection timeout") + mock_stub.exportData.side_effect = err + self.client._stub = mock_stub + + result = self.client._send_export_data(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("gRPC error: Connection timeout", result.result_status.message) + + def test_general_exception(self): + mock_stub = Mock() + mock_stub.exportData.side_effect = ValueError("boom") + self.client._stub = mock_stub + + result = self.client._send_export_data(self.request) + + self.assertTrue(result.result_status.is_error) + self.assertIn("Unexpected error: boom", result.result_status.message) + + +class TestExportDataUserFacing(unittest.TestCase): + """The user-facing method threads params through the builder and sender.""" + + def test_export_data_end_to_end(self): + client = ExportClient(Mock()) + response = _response_with_field("exportDataResult") + response.exportDataResult.filePath = "/srv/exports/out.h5" + response.exportDataResult.fileUrl = "" + mock_stub = Mock() + mock_stub.exportData.return_value = response + client._stub = mock_stub + + result = client.export_data(ExportDataRequestParams(ExportFormat.HDF5, dataset_id="ds-1")) + + self.assertFalse(result.result_status.is_error) + self.assertEqual(result.file_path, "/srv/exports/out.h5") + sent = mock_stub.exportData.call_args_list[0].args[0] + self.assertEqual(sent.dataSetId, "ds-1") + self.assertEqual(sent.outputFormat, _FORMAT.EXPORT_FORMAT_HDF5) + + +if __name__ == "__main__": + unittest.main() From 005ad87acae8979dd3f38766bd2cd5d44bd0875d Mon Sep 17 00:00:00 2001 From: Craig McChesney Date: Wed, 9 Sep 2026 14:04:14 -0600 Subject: [PATCH 2/3] test: live-server integration test for datasets and annotations (issue #6) Closes out Phase 1 with the wrapper-level round trip the plan puts in PR 1 (Q9), verified against an ecosystem built from dp-service fddf692: 18 tests, 12 subtests. Covers save/get round trips, every DataSetQuery and AnnotationQuery criterion (including the key-only attributes search), lowercase tag normalization, full-replace clearing omitted fields, paging with limit=1, a malformed page token, the get_datasets() batch fetch, calculations inline on getAnnotation versus id-only on queryAnnotations, get_calculations() click-through, delete-refused-while-referenced, and the delete cascade. Writing it surfaced a server behavior neither the proto nor the triage found: saveDataSet requires every PV named in a data block to already exist IN THE ARCHIVE. The error text says "no PV metadata found for names: [...]", but the check is a distinct on pvName over the buckets collection (MongoAnnotationHandler.validateSaveDataSetRequest -> MongoSyncQueryClient.executeQueryPvExistence), so saved PV metadata does not satisfy it. The test therefore ingests its own samples first, through the generated ingestion stub because IngestionClient wraps only registerProvider() until #17, and probes until the bucket is queryable -- ingestData() acks before the data is visible to the check. Recorded in the plan and CLAUDE.md, since it constrains every future example and test in this area. Also adds the CLAUDE.md Key Files entries and an invariants section for the new clients; the full usage section and cookbook recipe land in PR 2. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019he3UCsAnqTDE2VQ73Djwn --- CLAUDE.md | 43 +- plan/tickets/6/plan.md | 6 + .../test_datasets_annotations_integration.py | 534 ++++++++++++++++++ 3 files changed, 582 insertions(+), 1 deletion(-) create mode 100644 tests/integration/test_datasets_annotations_integration.py diff --git a/CLAUDE.md b/CLAUDE.md index bcd7696..a9dbf9c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -148,11 +148,14 @@ plan documents one change, `CLAUDE.md` documents the invariant it established. - `src/dp_python_lib/client/mldp_client.py` - Main client wrapper for the gRPC services - `src/dp_python_lib/client/ingestion_client.py` - Ingestion service client with methods like `register_provider()` -- `src/dp_python_lib/client/annotation_client.py` - Annotation service facade; groups feature-scoped clients sharing the one `DpAnnotationService` channel (exposes `.pv_metadata` and `.machine_config`, with room to grow `.annotations`) +- `src/dp_python_lib/client/annotation_client.py` - Annotation service facade; groups feature-scoped clients sharing the one `DpAnnotationService` channel (`.pv_metadata`, `.machine_config`, `.sample_status`, `.datasets`, `.annotations`, `.export` — every implemented `DpAnnotationService` feature area) - `src/dp_python_lib/client/pv_metadata_client.py` - PV metadata client (`save_pv_metadata()`, `get_pv_metadata()`, `query_pv_metadata()`, `iter_pv_metadata()`, `delete_pv_metadata()`) plus the `PvMetadataQuery` (`Q`) criterion helpers - `src/dp_python_lib/client/machine_config_client.py` - Machine configuration client covering both configurations (`save_configuration()`, `get_configuration()`, `query_configurations()`, `iter_configurations()`, `delete_configuration()`) and their temporal activations (`save_configuration_activation()`, `get_configuration_activation()`, `query_configuration_activations()`, `iter_configuration_activations()`, `delete_configuration_activation()`, `get_active_configurations()`). Includes the `ConfigurationQuery` (`C`) and `ConfigurationActivationQuery` (`CA`) criterion helpers and the `to_timestamp()` helper (tz-aware datetime / epoch seconds / `common.Timestamp`). Get/delete activation take a composite key (`client_activation_id` XOR `configuration_name`+`start_time`). Activation `end_time` is optional — omit it for an open-ended activation ("still in effect"); the field is then genuinely absent on the wire - `src/dp_python_lib/client/sample_status_client.py` - Sample status client (`save_sample_statuses()`, `query_sample_statuses()`, `iter_sample_statuses()`, `iter_sample_statuses_stream()`, `delete_sample_statuses()`) plus the `sampling_clock()` / `timestamp_list()` axis builders and the `SampleStatusColumn` / `SampleStatusFrame` construction classes. A status's identity key is `(pvName, timestamp, domain, layer)`; `delete_sample_statuses()` requires either `pv_names` or an explicit `all_pvs=True` opt-in for the destructive wildcard - `src/dp_python_lib/client/sample_status_conversions.py` - Per-sample expansion of query results (no optional extras required): `expand_data_timestamps()` (SamplingClock positions computed in **integer nanoseconds**, never float seconds — the exact-match contract depends on it), `bucket_to_rows()` / `buckets_to_rows()` / `iter_rows()` yielding `SampleStatusRow` objects with absent confidence/reason surfaced as `None` rather than fabricated `0.0`/`""` +- `src/dp_python_lib/client/dataset_client.py` - DataSet client (`save_dataset()`, `get_dataset()`, `query_datasets()`, `iter_datasets()`, `delete_dataset()`, plus the `get_datasets(ids)` batch fetch that avoids the annotation-listing N+1) with the `DataSetQuery` (`DS`) criterion helpers and the `data_block()` builder. `data_block()` is the only place `begin < end` is checked — the server does not +- `src/dp_python_lib/client/annotations_client.py` - Annotations client (`save_annotation()`, `get_annotation()`, `query_annotations()`, `iter_annotations()`, `delete_annotation()`, `get_calculations()`) with the `AnnotationQuery` (`AQ`) criterion helpers and the `calculations()` builder, which takes a `dict[str, DataFrame]` so frame-name uniqueness is true by construction. Note `AnnotationsClient` (feature client) vs `AnnotationClient` (facade) +- `src/dp_python_lib/client/export_client.py` - Export client (`export_data()`) with the `ExportFormat` str enum and the `calculations_spec()` builder - `src/dp_python_lib/client/query_client.py` - v2 time-series query client (sample-oriented) exposed as `client.query`. Low-level wrappers `query_samples()` (unary, one resumable page) and `iter_query_samples()` (transparent paging), plus `iter_query_samples_stream()` (server-streaming, fire-and-consume, lazy). Queries are described by a kind-neutral `QueryParams` built from the `PvQuery` (`PV`) and `ConfigQuery` (`CFG`) criterion helpers; shares a `_build_query_spec()` seam so a future bucket request builder reuses it. Results wrap the raw `ColumnTable` (`.column_table`, `.next_page_token`); `.to_dataframe()`/`.to_numpy()` delegate to `query_conversions` (Phase 2, optional `[analysis]` extra) - `src/dp_python_lib/client/query_conversions.py` - Pythonic conversions for query results (optional `[analysis]` extra: pandas/numpy/openpyxl, imported lazily). `data_value_to_python()` (oneof extractor: scalars→native, timestamp→epoch-nanos, array→list, structure→dict, image→`Image` wrapper, fail-loud on unhandled arm), `column_table_to_dataframe()` (UTC datetime index + one column per DataColumn; dense-alignment and duplicate-column-name fail-loud; ColumnMetadata in `df.attrs`), `column_table_to_numpy()` (dict of 1-D arrays; complex arms stay 1-D object arrays rather than collapsing to 2-D), `dataframe_to_excel()` (thin `to_excel()` wrapper: row-limit guard, tz-drop, complex-cell stringification), and `query_samples_to_dataframe()`/`stream_query_samples_to_dataframes()` whole-query conveniences (unary concats by column name; streaming yields per-page frames lazily) - `src/dp_python_lib/client/service_api_client_base.py` - Base class for the service clients: owns the channel and the one-per-client gRPC stub, and provides `_dispatch()`, the shared three-tier sender that all 18 unary `_send_*` methods delegate to @@ -163,6 +166,11 @@ plan documents one change, `CLAUDE.md` documents the invariant it established. - `tests/unit/test_machine_config_activation_client.py` - Unit tests for the ConfigurationActivation side of MachineConfigClient (incl. composite-key validation, timestamp handling, getActiveConfigurations) - `tests/unit/test_sample_status_client.py` - Unit tests for SampleStatusClient (frame/column validation, axis builders, three-tier error handling, paging, streaming, delete opt-in, `limit=0` regression) - `tests/unit/test_sample_status_conversions.py` - Unit tests for sample_status_conversions (nanosecond-exact axis expansion, absent-vs-zero confidence, alignment fail-loud, laziness) +- `tests/unit/test_dataset_client.py` - Unit tests for DataSetClient (request building, `DataSetQuery` incl. key-only attributes, `data_block()` validation, three-tier error handling, paging, `get_datasets()` dedup/empty/absent-id) +- `tests/unit/test_annotations_client.py` - Unit tests for AnnotationsClient (incl. `calculations()`, absent-vs-empty calculations on save, the `calculations_id` empty-string-vs-None rule, three-tier error handling, paging) +- `tests/unit/test_export_client.py` - Unit tests for ExportClient (`ExportFormat` mapping and unreachable `UNSPECIFIED`, `calculations_spec()`, the zero-source rejection, three-tier error handling) +- `tests/unit/test_annotation_client.py` - Unit tests pinning the `AnnotationClient` facade wiring (every feature client present, one shared channel, one stub apiece) +- `tests/integration/test_datasets_annotations_integration.py` - Live-server round trip for datasets/annotations/calculations; ingests its own samples first, because `saveDataSet` requires archived PVs - `tests/unit/test_query_client.py` - Unit tests for QueryClient (request building, three-tier error handling, unary paging, streaming, `PvQuery`/`ConfigQuery` helpers, `QueryParams` validation) - `tests/unit/test_query_conversions.py` - Unit tests for query_conversions (each DataValue arm, dense-alignment and duplicate-column-name fail-loud, int-gap float-upcast, timestamp columns, 1-D object arrays for complex arms, metadata in attrs, concat-by-name, Excel row-limit/stringification/native-bytes; DataFrame/NumPy/Excel tests skip cleanly when the `[analysis]` extra is absent) - `pyproject.toml` - Project metadata and dependencies @@ -514,6 +522,39 @@ Notes: `column_table_to_torch()` behind a separate optional `[torch]` extra — no change to `QueryClient` or the NumPy path. Not built yet. +### DataSets, Annotations, and Export API (Annotation Service) + +Phase 1 of issue #6 (`plan/tickets/6/plan.md`) added three feature clients on the `annotation` facade: +`client.annotation.datasets` (`DataSetClient`), `client.annotation.annotations` (`AnnotationsClient`), and +`client.annotation.export` (`ExportClient`). The full usage section lands with the cookbook recipe in PR 2; the +invariants worth knowing before touching this code: + +- **`saveDataSet` requires every PV named in a data block to already exist in the archive** — that is, to have + *ingested data*. The server's error text says `no PV metadata found for names: [...]`, but the check is a + `distinct` on `pvName` over the **buckets** collection (`MongoAnnotationHandler.validateSaveDataSetRequest` → + `MongoSyncQueryClient.executeQueryPvExistence`), so saving PV metadata does **not** satisfy it. Nothing is + validated client-side (the client cannot know what is archived), but any test or example must use an archived PV. + `ingestData()` acks *before* the bucket becomes queryable, so a `saveDataSet` issued immediately after ingesting + still fails; `tests/integration/test_datasets_annotations_integration.py` probes until it succeeds rather than + sleeping a fixed interval, and ingests through the generated stub because `IngestionClient` wraps only + `registerProvider()` until #17. +- **The server does not check `begin < end` on a `DataBlock`** (it checks only that each bound is non-zero and that + `pvNames` is non-empty), so `data_block()`'s check is the only one there is. +- **Delete-not-found is a business error**, not a silent success, on both `delete_dataset()` and + `delete_annotation()`. `delete_dataset()` is also refused while any annotation references the dataset — delete the + annotations first; there is deliberately no cascade. +- **`save_annotation()` replaces in full, including calculations**: omitting them clears *and deletes* the stored + object, and a replacement returns a new `calculationsId`. `get_annotation()` is the only method returning + calculations inline; `query_annotations()` results carry the id with empty content. +- Two criterion-helper differences from the older `PvMetadataQuery` / `ConfigurationQuery` helpers, both following + the proto: `attributes(key)` accepts an absent `values` list as a key-only existence search, and `criteria` is + optional because the server treats an empty list as match-all. Back-porting these to the five existing helpers is + [#40](https://github.com/osprey-dcs/dp-python-lib/issues/40) / [#41](https://github.com/osprey-dcs/dp-python-lib/issues/41). +- `ExportFormat` makes the server-rejected `EXPORT_FORMAT_UNSPECIFIED` unreachable, and `ExportDataRequestParams` + requires at least one of `dataset_id` / `data_blocks` / `calculations_spec`. The exported file lives on the + **server's** filesystem and there is no retrieval RPC, so there is no download convenience. +- `patchDataSet` / `patchAnnotation` are reserved "not implemented" placeholders and are not wrapped. + ### Sample Status API (Annotation Service) Sample status methods are exposed under the `annotation` facade at `client.annotation.sample_status`. A sample status diff --git a/plan/tickets/6/plan.md b/plan/tickets/6/plan.md index 551a8e3..7e48e79 100644 --- a/plan/tickets/6/plan.md +++ b/plan/tickets/6/plan.md @@ -31,6 +31,11 @@ re-verified in the dp-service Java source rather than taken from the ticket. Four things the first pass missed were folded in: the delete-not-found tier, the absent server-side `begin < end` check, the epoch-0 `SamplingClock` rejection (all three now rows in section 3), and the redundant-`count` question (resolved in D6). +- **Phase 1 implemented and merged-ready 2026-09-09.** Three clients, criterion helpers, params/results, facade + wiring, 149 unit tests (570 total). The wrapper-level integration test passes against a live ecosystem built from + dp-service `fddf692` (annotation on `localhost:50053`, ingestion on `:50051`): 18 tests, 12 subtests. Writing it + surfaced one further server behavior the triage had not found — `saveDataSet` requires its PVs to exist in the + archive — now recorded in section 3 and in `CLAUDE.md`. ## Overview @@ -162,6 +167,7 @@ attributes a reader of this table might reach for do not exist. | Page tokens are **opaque keyset tokens** carrying a query discriminator; a malformed, whitespace, legacy skip-offset, or wrong-query token is **rejected** (`REJECT`) | `iter_*` surfaces that as `RuntimeError`; nothing to parse. Note the three metadata queries still use skip tokens with silent restart (dp-service #193) — `conventions.md` must describe both behaviors | | Criteria AND across the list, OR within a criterion; **at most one `TextCriterion` per request** (a second is a validation `REJECT`, since two `$text` clauses cannot be ANDed) | client-side `ValueError` naming the rule is cheap and matches the "fail with a message naming the problem" posture (Q7) | | Blank criterion values are rejected server-side; `IdCriterion` ids and get/delete ids must be valid ObjectIds (malformed → `REJECT`, not "not found") | helpers reject empty inputs as usual; no ObjectId validation client-side (format is a server implementation detail) | +| **`saveDataSet` requires every PV named in a data block to already exist IN THE ARCHIVE.** Despite the error text (`no PV metadata found for names: [...]`), the check is a `distinct` on `pvName` over the *buckets* collection (`MongoAnnotationHandler.validateSaveDataSetRequest` → `MongoSyncQueryClient.executeQueryPvExistence`) — saved PV metadata does **not** satisfy it. Found 2026-09-09 while writing the Phase 1 integration test; neither the proto nor the ticket mentions it | nothing to validate client-side (the client cannot know what is archived), but it shapes the tests and the cookbook: a dataset can only name PVs with ingested data, so the integration test ingests its own samples first, and the cookbook's worked example must use an archived PV. `ingestData()` acks *before* the bucket is queryable, so a save issued immediately after ingesting still fails — the test probes until it succeeds rather than sleeping a fixed interval | | `getDataSet` / `getAnnotation` / `getCalculations` not-found → `ExceptionalResult` | same "business error" tier as `get_pv_metadata` | | `deleteDataSet` / `deleteAnnotation` **not-found is also a `REJECT`**, not a silent success (`DeleteAnnotationDispatcher:34-36`, "no Annotation record found for id: …") | same business-error tier, so no code change — but the cookbook teardown and the integration test's delete-twice leg must expect an error result on the second delete, not a success | | `DataBlock` validation is **only** `beginTime.epochSeconds >= 1`, `endTime.epochSeconds >= 1`, and a non-empty `pvNames` (`AnnotationValidationUtility.validateDataBlock`) — the server never checks `begin < end` | D4's client-side `begin < end` check is the *only* one there is, not a duplicate of a server check. A reversed block is accepted today, which is also why the proto's silence on half-openness is a real ambiguity rather than a documentation gap | diff --git a/tests/integration/test_datasets_annotations_integration.py b/tests/integration/test_datasets_annotations_integration.py new file mode 100644 index 0000000..888edfe --- /dev/null +++ b/tests/integration/test_datasets_annotations_integration.py @@ -0,0 +1,534 @@ +import logging +import os +import sys +import time +import unittest +from datetime import datetime, timedelta, timezone + +import grpc + +# Add src directory to path for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../src")) + +from dp_python_lib.client.annotations_client import ( + AnnotationQuery, + SaveAnnotationRequestParams, + calculations, +) +from dp_python_lib.client.dataset_client import ( + DataSetQuery, + SaveDataSetRequestParams, + data_block, +) +from dp_python_lib.client.mldp_client import MldpClient +from dp_python_lib.grpc import common_pb2, ingestion_pb2, ingestion_pb2_grpc + + +class TestDataSetsAnnotationsIntegration(unittest.TestCase): + """ + Integration tests for DataSetClient and AnnotationsClient that require a running MLDP ecosystem. + + Prerequisites: + - MLDP services running (annotation service at localhost:50053 AND ingestion service at localhost:50051), + however started + - The Annotation Service must carry the modernized DataSet/Annotation API: dp-grpc 1.16.0 or later, built from + dp-service main at or after PR #264 + + To run these tests: + 1. Start the MLDP ecosystem + 2. Run: python -m unittest tests.integration.test_datasets_annotations_integration -v + + Why ingestion is a prerequisite here: saveDataSet validates that every PV named in a data block already exists + IN THE ARCHIVE. Despite its error text ("no PV metadata found for names: ..."), the server's check is a + distinct on pvName over the buckets collection (MongoAnnotationHandler.validateSaveDataSetRequest -> + MongoSyncQueryClient.executeQueryPvExistence), so saving PV metadata is not enough -- the PV must have ingested + data. This class therefore ingests a few samples for its run-unique PV in setUpClass, using the generated + ingestion stub directly: the library's IngestionClient wraps only registerProvider() today, and ingestData() is + issue #17. + + Each test writes under a run-unique owner id and tag and deletes what it wrote, so runs neither collide with + each other nor with real data. + """ + + ANNOTATION_ADDRESS = "localhost:50053" + INGESTION_ADDRESS = "localhost:50051" + + # Samples ingested for the run's PV, so a data block over them has something to reference. + SAMPLE_COUNT = 5 + SAMPLE_PERIOD_NANOS = 1_000_000_000 + + @classmethod + def setUpClass(cls): + logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") + cls.logger = logging.getLogger(__name__) + cls.logger.info("Setting up datasets/annotations integration test environment") + + cls._verify_services_available() + + cls.client = MldpClient() + cls.datasets = cls.client.annotation.datasets + cls.annotations = cls.client.annotation.annotations + + cls._verify_modernized_api_available() + + # Run-unique namespace, so concurrent or repeated runs cannot see each other's records. + cls.run_id = str(int(time.time() * 1000)) + cls.owner_id = f"itest_owner_{cls.run_id}" + cls.tag = f"itest-tag-{cls.run_id}" + cls.pv_name = f"ITEST:DATASET:{cls.run_id}" + + # A fixed, whole-second base time well clear of "now", so the range is stable across the run. + cls.begin_time = datetime(2024, 2, 2, 18, 0, 0, tzinfo=timezone.utc) + cls.end_time = cls.begin_time + timedelta(hours=1) + + cls.logger.info("Using owner=%s tag=%s pv=%s", cls.owner_id, cls.tag, cls.pv_name) + + cls._ingest_samples_for_run_pv() + + @classmethod + def _verify_services_available(cls): + cls.logger.info("Checking if MLDP annotation and ingestion services are available") + for label, address in (("annotation", cls.ANNOTATION_ADDRESS), ("ingestion", cls.INGESTION_ADDRESS)): + try: + channel = grpc.insecure_channel(address) + grpc.channel_ready_future(channel).result(timeout=5) + cls.logger.info("%s service is reachable at %s", label.capitalize(), address) + channel.close() + except grpc.FutureTimeoutError: + raise unittest.SkipTest( + f"MLDP {label} service not available at {address}. " + "Please start the MLDP ecosystem before running integration tests." + ) from None + except Exception as e: + raise unittest.SkipTest( + f"Cannot connect to MLDP {label} service: {e}. Please ensure the MLDP ecosystem is running." + ) from None + + @classmethod + def _ingest_samples_for_run_pv(cls): + """ + Ingests a handful of samples for this run's PV, so data blocks naming it pass saveDataSet's + archive-existence check (see the class docstring). + + Uses the generated ingestion stub directly rather than the library, whose IngestionClient covers only + registerProvider() today; wrapping ingestData() is issue #17. + """ + channel = grpc.insecure_channel(cls.INGESTION_ADDRESS) + cls._ingestion_channel = channel + stub = ingestion_pb2_grpc.DpIngestionServiceStub(channel) + + registration = stub.registerProvider( + ingestion_pb2.RegisterProviderRequest(providerName=f"itest_provider_{cls.run_id}"), timeout=10 + ) + if registration.HasField("exceptionalResult"): + raise unittest.SkipTest( + f"could not register an ingestion provider: {registration.exceptionalResult.message}" + ) + provider_id = registration.registrationResult.providerId + + request = ingestion_pb2.IngestDataRequest(providerId=provider_id, clientRequestId=f"itest-{cls.run_id}") + clock = request.ingestionDataFrame.dataTimestamps.samplingClock + clock.startTime.epochSeconds = int(cls.begin_time.timestamp()) + clock.periodNanos = cls.SAMPLE_PERIOD_NANOS + clock.count = cls.SAMPLE_COUNT + + column = request.ingestionDataFrame.dataColumns.add() + column.name = cls.pv_name + for i in range(cls.SAMPLE_COUNT): + column.dataValues.add().doubleValue = float(i) + + response = stub.ingestData(request, timeout=15) + if response.HasField("exceptionalResult"): + raise unittest.SkipTest(f"could not ingest test data: {response.exceptionalResult.message}") + + cls._await_pv_in_archive() + cls.logger.info("Ingested %d samples for %s", cls.SAMPLE_COUNT, cls.pv_name) + + @classmethod + def _await_pv_in_archive(cls, attempts=20, delay_seconds=0.5): + """ + Waits until the ingested PV is visible to saveDataSet's existence check. + + ingestData() acks once the request is accepted, which is before the bucket is committed and queryable, so a + saveDataSet issued immediately afterwards still fails the check. Probe with a throwaway save rather than + sleeping a fixed interval. + """ + probe_params = SaveDataSetRequestParams( + name=f"itest archive probe {cls.run_id}", + owner_id=cls.owner_id, + data_blocks=[data_block(cls.begin_time, cls.end_time, [cls.pv_name])], + ) + for _ in range(attempts): + result = cls.datasets.save_dataset(probe_params) + if not result.result_status.is_error: + cls.datasets.delete_dataset(result.dataset_id) + return + time.sleep(delay_seconds) + + raise unittest.SkipTest( + f"ingested data for {cls.pv_name} did not become visible to saveDataSet within " + f"{attempts * delay_seconds:.0f}s: {result.result_status.message}" + ) + + @classmethod + def tearDownClass(cls): + channel = getattr(cls, "_ingestion_channel", None) + if channel is not None: + channel.close() + + @classmethod + def _verify_modernized_api_available(cls): + """ + Skips if the reachable Annotation Service predates the modernized DataSet/Annotation API. + + Reachability alone is not enough: a pre-1.16.0 server accepts the connection and then answers getDataSet + with UNIMPLEMENTED, which would surface as a wall of assertion failures rather than as the "backend not + available" skip these tests intend. Probe with a harmless get of a well-formed but absent ObjectId; a + server that has the API answers with a "not found" business error rather than a missing-method gRPC error. + """ + result = cls.datasets.get_dataset("000000000000000000000000") + message = result.result_status.message or "" + if result.result_status.is_error and "Method not found" in message: + raise unittest.SkipTest( + f"Annotation service at {cls.ANNOTATION_ADDRESS} does not implement the modernized DataSet API " + f"({message}). It is new in dp-grpc 1.16.0; upgrade the server to run these tests." + ) + cls.logger.info("Modernized DataSet/Annotation API is available") + + def setUp(self): + # Ids written by the running test, torn down in reverse order (annotations before datasets, since + # delete_dataset is refused while a dataset is referenced). + self.created_dataset_ids = [] + self.created_annotation_ids = [] + + def tearDown(self): + for annotation_id in reversed(self.created_annotation_ids): + self.annotations.delete_annotation(annotation_id) + for dataset_id in reversed(self.created_dataset_ids): + self.datasets.delete_dataset(dataset_id) + + def _save_dataset(self, name_suffix="", **overrides): + """Saves a dataset in this run's namespace, registers it for teardown, and returns its id.""" + kwargs = { + "name": f"itest dataset {self.run_id}{name_suffix}", + "owner_id": self.owner_id, + "data_blocks": [data_block(self.begin_time, self.end_time, [self.pv_name])], + "tags": [self.tag], + "attributes": {"runId": self.run_id}, + "modified_by": self.owner_id, + } + kwargs.update(overrides) + + result = self.datasets.save_dataset(SaveDataSetRequestParams(**kwargs)) + self.assertFalse(result.result_status.is_error, f"save_dataset failed: {result.result_status.message}") + self.assertTrue(result.dataset_id) + self.created_dataset_ids.append(result.dataset_id) + return result.dataset_id + + def _save_annotation(self, dataset_ids, name_suffix="", **overrides): + """Saves an annotation in this run's namespace, registers it for teardown, and returns the result.""" + kwargs = { + "name": f"itest annotation {self.run_id}{name_suffix}", + "owner_id": self.owner_id, + "dataset_ids": dataset_ids, + "tags": [self.tag], + "modified_by": self.owner_id, + } + kwargs.update(overrides) + + result = self.annotations.save_annotation(SaveAnnotationRequestParams(**kwargs)) + self.assertFalse(result.result_status.is_error, f"save_annotation failed: {result.result_status.message}") + self.assertTrue(result.annotation_id) + self.created_annotation_ids.append(result.annotation_id) + return result + + @staticmethod + def _calculations_frame(values=(12.7, 12.8, 12.9)): + """ + A hand-built single-column Calculations payload. + + Phase 1 has no data_frame builder yet (that is Phase 2), so this constructs the protos directly -- which is + also the escape hatch the builder is meant to complement rather than replace. + + Note the SamplingClock start time must be a real time: the server rejects an axis whose startTime is epoch + zero, so a fixture built from datetime(1970, 1, 1) would fail validation rather than the assertion. + """ + frame = common_pb2.DataFrame() + clock = frame.dataTimestamps.samplingClock + clock.startTime.epochSeconds = int(datetime(2024, 2, 2, 18, 0, 0, tzinfo=timezone.utc).timestamp()) + clock.startTime.nanoseconds = 0 + clock.periodNanos = 1_000_000_000 + clock.count = len(values) + + column = frame.doubleColumns.add() + column.name = "x_rms" + column.values[:] = list(values) + column.metadata.provenance.process = "1 Hz RMS" + source = column.metadata.provenance.derivedFrom.add() + source.pvName = "BPMS:GUNB:314:X" + + return calculations({"bpm-statistics": frame}) + + # ------------------------------------------------------------------ + # DataSet round trip + # ------------------------------------------------------------------ + + def test_dataset_save_get_round_trip(self): + """save -> get returns the same content, with server-set audit fields populated.""" + dataset_id = self._save_dataset() + + result = self.datasets.get_dataset(dataset_id) + + self.assertFalse(result.result_status.is_error, result.result_status.message) + dataset = result.dataset + self.assertEqual(dataset.id, dataset_id) + self.assertEqual(dataset.ownerId, self.owner_id) + self.assertEqual(dataset.name, f"itest dataset {self.run_id}") + self.assertEqual(len(dataset.dataBlocks), 1) + self.assertEqual(list(dataset.dataBlocks[0].pvNames), [self.pv_name]) + self.assertEqual(dataset.dataBlocks[0].beginTime.epochSeconds, int(self.begin_time.timestamp())) + self.assertEqual(dataset.dataBlocks[0].endTime.epochSeconds, int(self.end_time.timestamp())) + self.assertEqual({(a.name, a.value) for a in dataset.attributes}, {("runId", self.run_id)}) + # createdTime is server-set; updatedTime is unset on create. + self.assertTrue(dataset.HasField("createdTime")) + + def test_dataset_tags_are_normalized_lowercase(self): + """Tags are lowercased, deduplicated, and sorted on save, so they read back normalized.""" + mixed_case_tag = f"ITest-Mixed-{self.run_id}" + dataset_id = self._save_dataset(tags=[mixed_case_tag, mixed_case_tag.upper(), self.tag]) + + dataset = self.datasets.get_dataset(dataset_id).dataset + + self.assertIn(mixed_case_tag.lower(), list(dataset.tags)) + self.assertNotIn(mixed_case_tag, list(dataset.tags)) + # Deduplicated: the two case variants collapse to one entry. + self.assertEqual(list(dataset.tags).count(mixed_case_tag.lower()), 1) + + def test_dataset_get_not_found_is_a_business_error(self): + result = self.datasets.get_dataset("000000000000000000000000") + + self.assertTrue(result.result_status.is_error) + self.assertIsNone(result.dataset) + + def test_dataset_query_by_each_criterion(self): + """Each DataSetQuery criterion finds the dataset this run saved.""" + dataset_id = self._save_dataset() + + criteria_by_name = { + "ids": [DataSetQuery.ids([dataset_id])], + "owners": [DataSetQuery.owners([self.owner_id])], + "name": [DataSetQuery.name(prefix=[f"itest dataset {self.run_id}"])], + "pv_names": [DataSetQuery.pv_names([self.pv_name])], + "tags": [DataSetQuery.tags([self.tag])], + "attributes": [DataSetQuery.attributes("runId", [self.run_id])], + "attributes_key_only": [ + DataSetQuery.attributes("runId"), + DataSetQuery.owners([self.owner_id]), + ], + } + + for label, criteria in criteria_by_name.items(): + with self.subTest(criterion=label): + found = [d.id for d in self.datasets.iter_datasets(criteria)] + self.assertIn(dataset_id, found, f"{label} did not find the saved dataset") + + def test_dataset_replace_in_full_clears_omitted_fields(self): + """Saving with an id replaces in full: fields omitted from the second save are cleared.""" + dataset_id = self._save_dataset(description="original description") + self.assertEqual(self.datasets.get_dataset(dataset_id).dataset.description, "original description") + + # Re-save the same id without a description. + self._save_dataset(dataset_id=dataset_id) + + self.assertEqual(self.datasets.get_dataset(dataset_id).dataset.description, "") + + def test_dataset_paging(self): + """iter_datasets() follows page tokens across a run-unique tag with limit=1.""" + ids = {self._save_dataset(name_suffix=f" #{i}") for i in range(3)} + + found = {d.id for d in self.datasets.iter_datasets([DataSetQuery.tags([self.tag])], limit=1)} + + self.assertTrue(ids.issubset(found), f"paging missed some datasets: expected {ids}, found {found}") + + def test_malformed_page_token_is_a_business_error(self): + """Page tokens are opaque keyset tokens; a malformed one is rejected rather than silently restarted.""" + result = self.datasets.query_datasets([DataSetQuery.tags([self.tag])], page_token="not-a-real-token") + + self.assertTrue(result.result_status.is_error) + self.assertEqual(result.datasets, []) + + def test_get_datasets_batch_fetch(self): + """get_datasets() resolves several ids in one query, and tolerates one that resolves to nothing.""" + ids = [self._save_dataset(name_suffix=f" #{i}") for i in range(2)] + + found = self.datasets.get_datasets([*ids, "000000000000000000000000"]) + + self.assertEqual(set(found), set(ids)) + for dataset_id in ids: + self.assertEqual(found[dataset_id].id, dataset_id) + + # ------------------------------------------------------------------ + # Annotation round trip + # ------------------------------------------------------------------ + + def test_annotation_save_get_round_trip_with_calculations(self): + """save with calculations -> get returns them inline, keyed by the returned calculationsId.""" + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id], calculations=self._calculations_frame()) + + self.assertTrue(saved.calculations_id, "saving with calculations must return a calculationsId") + + result = self.annotations.get_annotation(saved.annotation_id) + + self.assertFalse(result.result_status.is_error, result.result_status.message) + annotation = result.annotation + self.assertEqual(annotation.id, saved.annotation_id) + self.assertEqual(list(annotation.dataSetIds), [dataset_id]) + self.assertEqual(annotation.calculationsId, saved.calculations_id) + + # getAnnotation is the only method that returns calculations inline. + inline = result.calculations + self.assertIsNotNone(inline) + self.assertEqual([f.name for f in inline.calculationDataFrames], ["bpm-statistics"]) + frame = inline.calculationDataFrames[0].frame + self.assertEqual([c.name for c in frame.doubleColumns], ["x_rms"]) + self.assertEqual(list(frame.doubleColumns[0].values), [12.7, 12.8, 12.9]) + self.assertEqual(frame.doubleColumns[0].metadata.provenance.process, "1 Hz RMS") + self.assertEqual(frame.doubleColumns[0].metadata.provenance.derivedFrom[0].pvName, "BPMS:GUNB:314:X") + + def test_annotation_without_calculations_returns_empty_calculations_id(self): + """No calculations in the request means an empty calculationsId -- empty string, not an error.""" + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id]) + + self.assertEqual(saved.calculations_id, "") + self.assertIsNone(self.annotations.get_annotation(saved.annotation_id).calculations) + + def test_get_calculations_click_through(self): + """queryAnnotations gives an id but no content; get_calculations() fetches the content it names.""" + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id], calculations=self._calculations_frame()) + + # Query results carry the id but leave the calculations empty. + from_query = [a for a in self.annotations.iter_annotations([AnnotationQuery.tags([self.tag])])] + matching = [a for a in from_query if a.id == saved.annotation_id] + self.assertEqual(len(matching), 1) + self.assertEqual(matching[0].calculationsId, saved.calculations_id) + self.assertEqual(len(matching[0].calculations.calculationDataFrames), 0) + + # The click-through fetches it. + result = self.annotations.get_calculations(saved.calculations_id) + + self.assertFalse(result.result_status.is_error, result.result_status.message) + self.assertEqual([f.name for f in result.calculations.calculationDataFrames], ["bpm-statistics"]) + + def test_annotation_replace_without_calculations_clears_them(self): + """A full replace omitting calculations clears AND deletes the stored object.""" + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id], calculations=self._calculations_frame()) + original_calculations_id = saved.calculations_id + + # Re-save the same annotation id with no calculations. + replaced = self._save_annotation([dataset_id], annotation_id=saved.annotation_id) + + self.assertEqual(replaced.annotation_id, saved.annotation_id) + self.assertEqual(replaced.calculations_id, "") + self.assertIsNone(self.annotations.get_annotation(saved.annotation_id).calculations) + + # The replaced calculations object is deleted, not orphaned. + orphan = self.annotations.get_calculations(original_calculations_id) + self.assertTrue(orphan.result_status.is_error) + + def test_annotation_query_by_each_criterion(self): + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id]) + + criteria_by_name = { + "ids": [AnnotationQuery.ids([saved.annotation_id])], + "owners": [AnnotationQuery.owners([self.owner_id])], + "datasets": [AnnotationQuery.datasets([dataset_id])], + "name": [AnnotationQuery.name(prefix=[f"itest annotation {self.run_id}"])], + "tags": [AnnotationQuery.tags([self.tag])], + } + + for label, criteria in criteria_by_name.items(): + with self.subTest(criterion=label): + found = [a.id for a in self.annotations.iter_annotations(criteria)] + self.assertIn(saved.annotation_id, found, f"{label} did not find the saved annotation") + + def test_annotation_get_not_found_is_a_business_error(self): + result = self.annotations.get_annotation("000000000000000000000000") + + self.assertTrue(result.result_status.is_error) + self.assertIsNone(result.annotation) + + # ------------------------------------------------------------------ + # Delete semantics + # ------------------------------------------------------------------ + + def test_delete_dataset_is_refused_while_referenced(self): + """A dataset an annotation targets cannot be deleted; the two-step teardown is the honest shape.""" + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id]) + + refused = self.datasets.delete_dataset(dataset_id) + + self.assertTrue(refused.result_status.is_error, "deleting a referenced dataset must be refused") + self.assertIsNone(refused.dataset_id) + + # Delete the annotation first, then the dataset succeeds. + deleted_annotation = self.annotations.delete_annotation(saved.annotation_id) + self.assertFalse(deleted_annotation.result_status.is_error, deleted_annotation.result_status.message) + self.created_annotation_ids.remove(saved.annotation_id) + + deleted_dataset = self.datasets.delete_dataset(dataset_id) + self.assertFalse(deleted_dataset.result_status.is_error, deleted_dataset.result_status.message) + self.assertEqual(deleted_dataset.dataset_id, dataset_id) + self.created_dataset_ids.remove(dataset_id) + + def test_deleting_twice_is_a_business_error(self): + """ + Delete-not-found is a REJECT, not a silent success. + + This is the behavior the pre-implementation triage corrected in the plan: the first delete succeeds, and the + second reports "no ... record found" rather than reporting success for a no-op. + """ + dataset_id = self._save_dataset() + + first = self.datasets.delete_dataset(dataset_id) + self.assertFalse(first.result_status.is_error, first.result_status.message) + self.created_dataset_ids.remove(dataset_id) + + second = self.datasets.delete_dataset(dataset_id) + self.assertTrue(second.result_status.is_error, "deleting an absent dataset must be an error") + self.assertIsNone(second.dataset_id) + + def test_delete_annotation_cascades_to_its_calculations(self): + """Deleting an annotation deletes the calculations it owns.""" + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id], calculations=self._calculations_frame()) + + self.assertFalse(self.annotations.get_calculations(saved.calculations_id).result_status.is_error) + + deleted = self.annotations.delete_annotation(saved.annotation_id) + self.assertFalse(deleted.result_status.is_error, deleted.result_status.message) + self.created_annotation_ids.remove(saved.annotation_id) + + gone = self.annotations.get_calculations(saved.calculations_id) + self.assertTrue(gone.result_status.is_error, "the annotation's calculations must be deleted with it") + + def test_delete_annotation_twice_is_a_business_error(self): + dataset_id = self._save_dataset() + saved = self._save_annotation([dataset_id]) + + first = self.annotations.delete_annotation(saved.annotation_id) + self.assertFalse(first.result_status.is_error, first.result_status.message) + self.created_annotation_ids.remove(saved.annotation_id) + + second = self.annotations.delete_annotation(saved.annotation_id) + self.assertTrue(second.result_status.is_error) + self.assertIsNone(second.annotation_id) + + +if __name__ == "__main__": + unittest.main() From 5f301030c83248584a8f35656d462ddbf2eddb02 Mon Sep 17 00:00:00 2001 From: Craig McChesney Date: Thu, 10 Sep 2026 11:09:37 -0600 Subject: [PATCH 3/3] fix: address PR #44 review findings (issue #6, PR 1) Resolves the review findings on PR #44, plus the DataBlock interval question it raised. - calculations(): the docstring pointed at data_frame.data_frame(), which ships in PR 2. Reworded to describe building the frame directly and to name the builders as forthcoming (Copilot review comment). - get_datasets(): chunk the id list (ID_QUERY_CHUNK_SIZE, 100) instead of sending every id as one $in. The list comes from the dataSetIds of a whole page of annotations and is effectively unbounded, so a single criterion could exceed the gRPC max message size. A short result is now logged at WARNING: an id withheld for any other reason is indistinguishable from a dangling one, so the shortfall should not pass silently. - Moved _check_at_most_one_text_criterion out of dataset_client into a new query_support module as check_at_most_one_text_criterion. It was being imported private and cross-module by annotations_client, which made that module's dependencies misleading and the name easy to break. This also removes annotations_client's only dependency on dataset_client. - data_block(): document the range as half-open [begin, end). Established by reading dp-service, since annotation.proto does not say: export reaches the same bucket-overlap filter and the same per-sample TabularDataUtility.isRetained() that querySamples() does, and that function excludes a sample landing exactly on end_time. So back-to-back blocks cover the boundary sample exactly once, matching QueryParams. HDF5 export is the documented exception -- it writes overlapping buckets whole and untrimmed, so it can include out-of-range samples. - data_block(): reject a bare string for pv_names, which otherwise assigns one PV name per character and stores a silently wrong DataSet. - SaveDataSetRequestParams and SaveAnnotationRequestParams now validate the fields the server requires (name, owner_id, and data_blocks/dataset_ids), matching ExportDataRequestParams, which already validated its own. - CLAUDE.md: record the interval finding, the empty-string-vs-None rule on calculations_id and file_url, the params validation, the get_datasets chunking, and the new query_support module. 583 unit tests pass (13 new); ruff lint and format clean; cookbook snippet checker passes. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019he3UCsAnqTDE2VQ73Djwn --- CLAUDE.md | 23 +++- .../client/annotations_client.py | 18 ++- src/dp_python_lib/client/dataset_client.py | 107 ++++++++++++------ src/dp_python_lib/client/query_support.py | 32 ++++++ tests/unit/test_annotations_client.py | 26 +++++ tests/unit/test_dataset_client.py | 75 ++++++++++++ 6 files changed, 239 insertions(+), 42 deletions(-) create mode 100644 src/dp_python_lib/client/query_support.py diff --git a/CLAUDE.md b/CLAUDE.md index a9dbf9c..68d6d64 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -159,6 +159,7 @@ plan documents one change, `CLAUDE.md` documents the invariant it established. - `src/dp_python_lib/client/query_client.py` - v2 time-series query client (sample-oriented) exposed as `client.query`. Low-level wrappers `query_samples()` (unary, one resumable page) and `iter_query_samples()` (transparent paging), plus `iter_query_samples_stream()` (server-streaming, fire-and-consume, lazy). Queries are described by a kind-neutral `QueryParams` built from the `PvQuery` (`PV`) and `ConfigQuery` (`CFG`) criterion helpers; shares a `_build_query_spec()` seam so a future bucket request builder reuses it. Results wrap the raw `ColumnTable` (`.column_table`, `.next_page_token`); `.to_dataframe()`/`.to_numpy()` delegate to `query_conversions` (Phase 2, optional `[analysis]` extra) - `src/dp_python_lib/client/query_conversions.py` - Pythonic conversions for query results (optional `[analysis]` extra: pandas/numpy/openpyxl, imported lazily). `data_value_to_python()` (oneof extractor: scalars→native, timestamp→epoch-nanos, array→list, structure→dict, image→`Image` wrapper, fail-loud on unhandled arm), `column_table_to_dataframe()` (UTC datetime index + one column per DataColumn; dense-alignment and duplicate-column-name fail-loud; ColumnMetadata in `df.attrs`), `column_table_to_numpy()` (dict of 1-D arrays; complex arms stay 1-D object arrays rather than collapsing to 2-D), `dataframe_to_excel()` (thin `to_excel()` wrapper: row-limit guard, tz-drop, complex-cell stringification), and `query_samples_to_dataframe()`/`stream_query_samples_to_dataframes()` whole-query conveniences (unary concats by column name; streaming yields per-page frames lazily) - `src/dp_python_lib/client/service_api_client_base.py` - Base class for the service clients: owns the channel and the one-per-client gRPC stub, and provides `_dispatch()`, the shared three-tier sender that all 18 unary `_send_*` methods delegate to +- `src/dp_python_lib/client/query_support.py` - Helpers shared by the criteria-based query clients, currently `check_at_most_one_text_criterion()` (Mongo cannot AND two `$text` clauses). It is generic over the different criterion types because each names its oneof `criterion` and its full-text arm `textCriterion`. New shared query helpers belong here rather than in whichever feature client happened to need one first — importing a private name across feature modules makes the importing module's dependencies misleading - `tests/unit/test_service_api_client_base.py` - Unit tests for `_dispatch` itself (success, business error, unrecognized response, `RpcError` with and without a resolvable `code()`, unexpected exception, and the `request_log`/`success_log` hooks) - `tests/unit/test_ingestion_client.py` - Unit tests for IngestionClient functionality - `tests/unit/test_pv_metadata_client.py` - Unit tests for PvMetadataClient functionality @@ -539,7 +540,16 @@ invariants worth knowing before touching this code: sleeping a fixed interval, and ingests through the generated stub because `IngestionClient` wraps only `registerProvider()` until #17. - **The server does not check `begin < end` on a `DataBlock`** (it checks only that each bound is non-zero and that - `pvNames` is non-empty), so `data_block()`'s check is the only one there is. + `pvNames` is non-empty, and never compares the two bounds), so `data_block()`'s check is the only one there is. +- **A `DataBlock`'s range is half-open, `[begin, end)`** — the same convention as the v2 query API's `QueryParams`, + established by reading dp-service rather than the proto, which says nothing. `saveDataSet` never compares the + bounds at all; the interval acquires meaning only at export. There, bucket selection and per-sample trimming are + literally the *same* functions `querySamples()` uses (`MongoQueryFilterBuilder.bucketOverlapsRangeFilter` and + `TabularDataUtility.isRetained`, whose contract is "a sample exactly at an interval's end belongs to the next + interval, not this one"). So back-to-back blocks cover the boundary sample exactly once. **The exception is + HDF5**: `ExportDataJobAbstractBucketed` writes every *overlapping bucket* whole and untrimmed, with no time range + passed to the writer, so an HDF5 export can contain samples outside the requested range and can write a straddling + bucket twice. Nothing client-side can change that — it is a property of the format, and worth telling users. - **Delete-not-found is a business error**, not a silent success, on both `delete_dataset()` and `delete_annotation()`. `delete_dataset()` is also refused while any annotation references the dataset — delete the annotations first; there is deliberately no cascade. @@ -553,6 +563,17 @@ invariants worth knowing before touching this code: - `ExportFormat` makes the server-rejected `EXPORT_FORMAT_UNSPECIFIED` unreachable, and `ExportDataRequestParams` requires at least one of `dataset_id` / `data_blocks` / `calculations_spec`. The exported file lives on the **server's** filesystem and there is no retrieval RPC, so there is no download convenience. +- **Empty string and `None` mean different things on two result properties**, and the distinction is load-bearing in + both: `SaveAnnotationApiResult.calculations_id` is `""` when the request carried no calculations but `None` when + the call *failed*, and `ExportDataApiResult.file_url` is `""` when the deployment simply does not publish exports + over HTTP but `None` on error. In both cases `""` is a successful outcome — test with `is None`, not truthiness. +- **The params classes validate the server's required fields client-side.** `SaveDataSetRequestParams` requires + `name` / `owner_id` / `data_blocks`, and `SaveAnnotationRequestParams` requires `name` / `owner_id` / + `dataset_ids`, each raising `ValueError` naming the field rather than spending a round trip to learn it. +- **`get_datasets()` chunks its id list** (`ID_QUERY_CHUNK_SIZE`, 100) because the ids come from the `dataSetIds` of + a whole page of annotations and are effectively unbounded; one criterion carrying all of them becomes a single + large `$in` and an oversized request message. A short result is logged at WARNING, since an id withheld for any + other reason is indistinguishable from a dangling one. - `patchDataSet` / `patchAnnotation` are reserved "not implemented" placeholders and are not wrapped. ### Sample Status API (Annotation Service) diff --git a/src/dp_python_lib/client/annotations_client.py b/src/dp_python_lib/client/annotations_client.py index 7130b62..9922609 100644 --- a/src/dp_python_lib/client/annotations_client.py +++ b/src/dp_python_lib/client/annotations_client.py @@ -3,7 +3,7 @@ import grpc -from dp_python_lib.client.dataset_client import _check_at_most_one_text_criterion +from dp_python_lib.client.query_support import check_at_most_one_text_criterion from dp_python_lib.client.result import ApiResultBase from dp_python_lib.client.service_api_client_base import ServiceApiClientBase from dp_python_lib.grpc import annotation_pb2, annotation_pb2_grpc, common_pb2 @@ -17,8 +17,10 @@ def calculations(frames: dict[str, common_pb2.DataFrame]) -> annotation_pb2.Calc Taking a dict rather than a list makes frame-name uniqueness true by construction; the server rejects duplicate frame names, and a list would let a caller build one. - Build each frame with data_frame.data_frame(); it validates the frame's internal shape (column count against the - time axis, unique column names, non-empty names and values) so an error names the offending column. + Each frame is a common.DataFrame: a time axis plus the columns sampled on it. Assemble that message directly + for now -- the data_frame builders (dp_python_lib.client.data_frame), which validate a frame's internal shape + (column count against the time axis, unique column names, non-empty names and values) so an error names the + offending column, arrive in the follow-up PR for issue #6. Note the proto's naming trap: the repeated field is 'calculationDataFrames' (singular "calculation") while the message it holds is 'CalculationsDataFrame' (plural). @@ -243,7 +245,15 @@ def __init__( :param calculations: Derived values to store with the annotation (see calculations()). saveAnnotation() is the only write path for calculations, and omitting them on a replace clears the stored object. :param annotation_id: Id of an existing annotation to replace in full. Omit to create a new one. + :raises ValueError: if name, owner_id, or dataset_ids is empty. """ + if not name: + raise ValueError("SaveAnnotationRequestParams requires a non-empty name") + if not owner_id: + raise ValueError("SaveAnnotationRequestParams requires a non-empty owner_id") + if not dataset_ids: + raise ValueError("SaveAnnotationRequestParams requires a non-empty dataset_ids list") + self.name = name self.owner_id = owner_id self.dataset_ids = dataset_ids @@ -678,7 +688,7 @@ def query_annotations( :raises ValueError: if criteria contains more than one text criterion. """ criteria = criteria or [] - _check_at_most_one_text_criterion(criteria, "query_annotations()") + check_at_most_one_text_criterion(criteria, "query_annotations()") self.logger.info("Starting queryAnnotations operation with %d criteria", len(criteria)) diff --git a/src/dp_python_lib/client/dataset_client.py b/src/dp_python_lib/client/dataset_client.py index 47f9a3b..372bb23 100644 --- a/src/dp_python_lib/client/dataset_client.py +++ b/src/dp_python_lib/client/dataset_client.py @@ -4,10 +4,17 @@ import grpc from dp_python_lib.client.machine_config_client import TimestampInput, to_timestamp +from dp_python_lib.client.query_support import check_at_most_one_text_criterion from dp_python_lib.client.result import ApiResultBase from dp_python_lib.client.service_api_client_base import ServiceApiClientBase from dp_python_lib.grpc import annotation_pb2, annotation_pb2_grpc, common_pb2 +ID_QUERY_CHUNK_SIZE = 100 +""" +Default number of ids per query in get_datasets(). Keeps a single request's $in clause and message size bounded +when the caller passes an arbitrarily long id list; the value is a conservative round number, not a server limit. +""" + def data_block( begin_time: TimestampInput, @@ -18,19 +25,31 @@ def data_block( Builds a DataBlock: one time range plus the PV names covered over it. A DataSet is a list of these, and the exportData() API accepts them inline for a one-off export that does not warrant saving a DataSet. - The proto does not state whether the interval is half-open; the dp-grpc cookbook flags the same ambiguity. Treat - a block as covering [begin_time, end_time] inclusively unless the server documents otherwise. + The range is HALF-OPEN, [begin_time, end_time), matching the v2 query API's QueryParams. annotation.proto does + not say so, but the server's sample-level retention test does: TabularDataUtility.isRetained() excludes a sample + landing exactly on end_time ("a sample exactly at an interval's end belongs to the next interval, not this one"), + and the export job reaches that same function that querySamples() does. So back-to-back blocks -- one ending at + T, the next beginning at T -- cover the sample at T exactly once. - Note that begin < end is checked HERE and nowhere else: the server validates only that each bound is non-zero and - that pvNames is non-empty (AnnotationValidationUtility.validateDataBlock), so a reversed block would otherwise be - accepted and stored. + That holds for CSV and XLSX exports. HDF5 export is bucket-granular: ExportDataJobAbstractBucketed writes every + bucket that OVERLAPS the block, whole and untrimmed, so an HDF5 file can contain samples outside the requested + range, and back-to-back blocks sharing a straddling bucket write it twice. Nothing client-side can change that; + it is a property of the export format. - :param begin_time: Start of the block's time range (tz-aware datetime, epoch seconds, or common.Timestamp). - :param end_time: End of the block's time range (same accepted forms). - :param pv_names: Names of the PVs the block covers. + Note that begin < end is checked HERE and nowhere else: the server validates only that each bound is non-zero and + that pvNames is non-empty (AnnotationValidationUtility.validateDataBlock), and never compares the two bounds, so + a reversed block would otherwise be accepted and stored. + + :param begin_time: Start of the block's time range, inclusive (tz-aware datetime, epoch seconds, + or common.Timestamp). + :param end_time: End of the block's time range, exclusive (same accepted forms). + :param pv_names: Names of the PVs the block covers. A list -- passing a bare string is rejected rather than + silently iterated into one PV name per character. :return: An annotation.DataBlock for the specified range and PVs. - :raises ValueError: if pv_names is empty or begin_time is not strictly before end_time. + :raises ValueError: if pv_names is empty or a bare string, or begin_time is not strictly before end_time. """ + if isinstance(pv_names, str): + raise ValueError(f"data_block() requires a list of PV names, not a bare string; got {pv_names!r}") if not pv_names: raise ValueError("data_block() requires a non-empty pv_names list") @@ -195,26 +214,6 @@ def attributes( return criterion -def _check_at_most_one_text_criterion(criteria: list, op_name: str) -> None: - """ - Rejects a criteria list carrying more than one text criterion. - - Two $text clauses cannot be ANDed -- Mongo rejects the query with "Too many text expressions" -- so the server - rejects the second one during validation. Catching it here fails with a message naming the rule instead of - surfacing a server rejection. - - :param criteria: The criteria list to check. - :param op_name: The calling operation, used in the error message. - :raises ValueError: if more than one criterion carries a textCriterion. - """ - text_count = sum(1 for criterion in criteria if criterion.WhichOneof("criterion") == "textCriterion") - if text_count > 1: - raise ValueError( - f"{op_name} accepts at most one text criterion, got {text_count}; two full-text clauses cannot be " - f"combined with AND. Combine the terms into a single text() criterion instead." - ) - - class SaveDataSetRequestParams: """ Encapsulates client parameters for a call to the saveDataSet() API method. @@ -244,7 +243,15 @@ def __init__( :param attributes: Map of key/value attributes describing the DataSet. :param modified_by: Identifier of the user or process making the change. :param dataset_id: Id of an existing DataSet to replace in full. Omit to create a new one. + :raises ValueError: if name, owner_id, or data_blocks is empty. """ + if not name: + raise ValueError("SaveDataSetRequestParams requires a non-empty name") + if not owner_id: + raise ValueError("SaveDataSetRequestParams requires a non-empty owner_id") + if not data_blocks: + raise ValueError("SaveDataSetRequestParams requires a non-empty data_blocks list") + self.name = name self.owner_id = owner_id self.data_blocks = data_blocks @@ -598,7 +605,7 @@ def query_datasets( :raises ValueError: if criteria contains more than one text criterion. """ criteria = criteria or [] - _check_at_most_one_text_criterion(criteria, "query_datasets()") + check_at_most_one_text_criterion(criteria, "query_datasets()") self.logger.info("Starting queryDataSets operation with %d criteria", len(criteria)) @@ -641,21 +648,35 @@ def iter_datasets( if not page_token: break - def get_datasets(self, ids: list[str]) -> dict[str, annotation_pb2.DataSet]: + def get_datasets(self, ids: list[str], chunk_size: int = ID_QUERY_CHUNK_SIZE) -> dict[str, annotation_pb2.DataSet]: """ - Batch-fetches DataSets by id in a single paged query, rather than one getDataSet() call apiece. + Batch-fetches DataSets by id in a small number of paged queries, rather than one getDataSet() call apiece. This exists for the common "a page of annotations needs its datasets" case: annotation records carry dataSetIds rather than the DataSets themselves, so resolving them one at a time is an N+1. Ids are deduplicated with their order preserved. An empty ids list returns {} without issuing an RPC, - because callers feed this from annotation dataSetIds lists that may legitimately be empty. Ids that resolve - to nothing are simply absent from the returned dict: a dangling dataSetIds entry is not an error. + because callers feed this from annotation dataSetIds lists that may legitimately be empty. + + Ids are fetched in chunks of chunk_size, because the id list is caller-supplied and effectively unbounded -- + it comes from the dataSetIds of a whole page of annotations -- and one criterion carrying every id becomes a + single large $in and a correspondingly large request message, which at some size exceeds the gRPC maximum. + Chunking keeps each request bounded regardless of how many ids are asked for. + + Ids that resolve to nothing are simply absent from the returned dict, because a dangling dataSetIds entry is + a normal consequence of deleting a DataSet's annotations out from under it, not an error. Note that an id + withheld for any other reason is indistinguishable from a dangling one here, so a short result is logged at + WARNING rather than passed over silently; compare len() against the ids you asked for if it matters. :param ids: DataSet ids to fetch. + :param chunk_size: Maximum number of ids to request per query. Rarely worth overriding. :return: A dict mapping DataSet id to DataSet, for those ids that resolved. + :raises ValueError: if chunk_size is not positive. :raises RuntimeError: if any page returns an error. """ + if chunk_size < 1: + raise ValueError(f"get_datasets() requires a positive chunk_size, got {chunk_size}") + if not ids: self.logger.debug("get_datasets() called with no ids; returning empty result without an RPC") return {} @@ -663,9 +684,21 @@ def get_datasets(self, ids: list[str]) -> dict[str, annotation_pb2.DataSet]: unique_ids = list(dict.fromkeys(ids)) self.logger.info("Starting get_datasets batch fetch for %d unique id(s)", len(unique_ids)) - found = {dataset.id: dataset for dataset in self.iter_datasets([DataSetQuery.ids(unique_ids)])} - - self.logger.info("get_datasets resolved %d of %d requested id(s)", len(found), len(unique_ids)) + found: dict[str, annotation_pb2.DataSet] = {} + for offset in range(0, len(unique_ids), chunk_size): + chunk = unique_ids[offset : offset + chunk_size] + for dataset in self.iter_datasets([DataSetQuery.ids(chunk)]): + found[dataset.id] = dataset + + if len(found) < len(unique_ids): + self.logger.warning( + "get_datasets resolved only %d of %d requested id(s); the rest named DataSets that do not exist " + "or were not returned", + len(found), + len(unique_ids), + ) + else: + self.logger.info("get_datasets resolved %d of %d requested id(s)", len(found), len(unique_ids)) return found # ------------------------------------------------------------------ diff --git a/src/dp_python_lib/client/query_support.py b/src/dp_python_lib/client/query_support.py new file mode 100644 index 0000000..93cb2cc --- /dev/null +++ b/src/dp_python_lib/client/query_support.py @@ -0,0 +1,32 @@ +""" +Small helpers shared by the criteria-based query clients. + +These live here rather than in any one feature client because more than one client needs them, and importing a +private name across feature modules (dataset_client importing from annotations_client, or the reverse) makes the +importing module's dependencies misleading and the imported name easy to break by accident. +""" + +from typing import Any + + +def check_at_most_one_text_criterion(criteria: list[Any], op_name: str) -> None: + """ + Rejects a criteria list carrying more than one text criterion. + + Two $text clauses cannot be ANDed -- Mongo rejects the query with "Too many text expressions" -- so the server + rejects the second one during validation. Catching it here fails with a message naming the rule instead of + surfacing a server rejection. + + This is generic over the criterion types of the different query APIs: every one of them names its oneof + "criterion" and its full-text arm "textCriterion", so the same check serves them all. + + :param criteria: The criteria list to check. + :param op_name: The calling operation, used in the error message. + :raises ValueError: if more than one criterion carries a textCriterion. + """ + text_count = sum(1 for criterion in criteria if criterion.WhichOneof("criterion") == "textCriterion") + if text_count > 1: + raise ValueError( + f"{op_name} accepts at most one text criterion, got {text_count}; two full-text clauses cannot be " + f"combined with AND. Combine the terms into a single text() criterion instead." + ) diff --git a/tests/unit/test_annotations_client.py b/tests/unit/test_annotations_client.py index 0253f67..9e24394 100644 --- a/tests/unit/test_annotations_client.py +++ b/tests/unit/test_annotations_client.py @@ -174,6 +174,32 @@ def test_build_query_request_empty_page_token_is_unset(self): self.assertEqual(assigned, [], "an empty page token must not be assigned") +class TestSaveAnnotationRequestParamsValidation(unittest.TestCase): + """ + Unit tests for the params validation. The server requires all three, so catching them here turns a round trip + into an immediate error naming the field. + """ + + def test_rejects_empty_name(self): + with self.assertRaises(ValueError) as ctx: + SaveAnnotationRequestParams(name="", owner_id="cmcchesney", dataset_ids=["ds-1"]) + self.assertIn("name", str(ctx.exception)) + + def test_rejects_empty_owner_id(self): + with self.assertRaises(ValueError) as ctx: + SaveAnnotationRequestParams(name="orbit drift", owner_id="", dataset_ids=["ds-1"]) + self.assertIn("owner_id", str(ctx.exception)) + + def test_rejects_empty_dataset_ids(self): + with self.assertRaises(ValueError) as ctx: + SaveAnnotationRequestParams(name="orbit drift", owner_id="cmcchesney", dataset_ids=[]) + self.assertIn("dataset_ids", str(ctx.exception)) + + def test_accepts_required_fields(self): + params = SaveAnnotationRequestParams(name="orbit drift", owner_id="cmcchesney", dataset_ids=["ds-1"]) + self.assertEqual(params.dataset_ids, ["ds-1"]) + + class TestAnnotationQueryHelpers(unittest.TestCase): """Unit tests for the AnnotationQuery criterion builders.""" diff --git a/tests/unit/test_dataset_client.py b/tests/unit/test_dataset_client.py index e893d11..fab40d9 100644 --- a/tests/unit/test_dataset_client.py +++ b/tests/unit/test_dataset_client.py @@ -75,6 +75,12 @@ def test_rejects_naive_datetime(self): with self.assertRaises(ValueError): data_block(datetime(2026, 7, 14), END, ["A:1"]) # noqa: DTZ001 -- naive input is the condition under test + def test_rejects_bare_string_pv_names(self): + # Without the guard this assigns one PV name per character, saving a silently wrong DataSet. + with self.assertRaises(ValueError) as ctx: + data_block(BEGIN, END, "A:1") + self.assertIn("bare string", str(ctx.exception)) + class TestDataSetClientBuildRequests(unittest.TestCase): """Unit tests for the request-building helpers (no gRPC calls).""" @@ -171,6 +177,35 @@ def test_build_query_request_empty_page_token_is_unset(self): self.assertEqual(assigned, [], "an empty page token must not be assigned") +class TestSaveDataSetRequestParamsValidation(unittest.TestCase): + """ + Unit tests for the params validation. The server requires all three, so catching them here turns a round trip + into an immediate error naming the field. + """ + + def setUp(self): + self.block = data_block(BEGIN, END, ["A:1"]) + + def test_rejects_empty_name(self): + with self.assertRaises(ValueError) as ctx: + SaveDataSetRequestParams(name="", owner_id="cmcchesney", data_blocks=[self.block]) + self.assertIn("name", str(ctx.exception)) + + def test_rejects_empty_owner_id(self): + with self.assertRaises(ValueError) as ctx: + SaveDataSetRequestParams(name="ramp study", owner_id="", data_blocks=[self.block]) + self.assertIn("owner_id", str(ctx.exception)) + + def test_rejects_empty_data_blocks(self): + with self.assertRaises(ValueError) as ctx: + SaveDataSetRequestParams(name="ramp study", owner_id="cmcchesney", data_blocks=[]) + self.assertIn("data_blocks", str(ctx.exception)) + + def test_accepts_required_fields(self): + params = SaveDataSetRequestParams(name="ramp study", owner_id="cmcchesney", data_blocks=[self.block]) + self.assertEqual(params.name, "ramp study") + + class TestDataSetQueryHelpers(unittest.TestCase): """Unit tests for the DataSetQuery criterion builders.""" @@ -677,6 +712,46 @@ def test_page_error_raises_runtime_error(self): with self.assertRaises(RuntimeError): self.client.get_datasets(["a"]) + def test_chunks_long_id_lists(self): + # The id list is caller-supplied and unbounded, so it must not become one oversized $in. + ids = [f"id-{i}" for i in range(250)] + mock_stub = Mock() + mock_stub.queryDataSets.side_effect = [self._page(ids[0:100]), self._page(ids[100:200]), self._page(ids[200:])] + self.client._stub = mock_stub + + found = self.client.get_datasets(ids) + + self.assertEqual(len(found), 250) + self.assertEqual(mock_stub.queryDataSets.call_count, 3) + sent = [list(call.args[0].criteria[0].idCriterion.ids) for call in mock_stub.queryDataSets.call_args_list] + self.assertEqual([len(chunk) for chunk in sent], [100, 100, 50]) + self.assertEqual([i for chunk in sent for i in chunk], ids, "chunking must preserve order and lose nothing") + + def test_chunk_size_is_configurable(self): + mock_stub = Mock() + mock_stub.queryDataSets.side_effect = [self._page(["a", "b"]), self._page(["c"])] + self.client._stub = mock_stub + + self.client.get_datasets(["a", "b", "c"], chunk_size=2) + + self.assertEqual(mock_stub.queryDataSets.call_count, 2) + + def test_rejects_non_positive_chunk_size(self): + with self.assertRaises(ValueError) as ctx: + self.client.get_datasets(["a"], chunk_size=0) + self.assertIn("chunk_size", str(ctx.exception)) + + def test_short_result_is_logged_at_warning(self): + # A withheld id is indistinguishable from a dangling one, so the shortfall must not pass silently. + mock_stub = Mock() + mock_stub.queryDataSets.return_value = self._page(["a"]) + self.client._stub = mock_stub + + with self.assertLogs(self.client.logger, level="WARNING") as captured: + self.client.get_datasets(["a", "missing"]) + + self.assertIn("resolved only 1 of 2", "\n".join(captured.output)) + if __name__ == "__main__": unittest.main()