diff --git a/CLAUDE.md b/CLAUDE.md index bcd7696..68d6d64 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -148,14 +148,18 @@ 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 +- `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 @@ -163,6 +167,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 +523,59 @@ 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, 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. +- **`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. +- **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) 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 9abf77e..7e48e79 100644 --- a/plan/tickets/6/plan.md +++ b/plan/tickets/6/plan.md @@ -22,6 +22,20 @@ 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). +- **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 @@ -153,7 +167,11 @@ 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 | +| 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 +264,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 +285,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 +382,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..9922609 --- /dev/null +++ b/src/dp_python_lib/client/annotations_client.py @@ -0,0 +1,848 @@ +import logging +from collections.abc import Iterator + +import grpc + +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 + + +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. + + 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). + + :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. + :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 + 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..372bb23 --- /dev/null +++ b/src/dp_python_lib/client/dataset_client.py @@ -0,0 +1,761 @@ +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.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, + 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 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. + + 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. + + 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 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") + + 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 + + +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. + :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 + 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], chunk_size: int = ID_QUERY_CHUNK_SIZE) -> dict[str, annotation_pb2.DataSet]: + """ + 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 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 {} + + unique_ids = list(dict.fromkeys(ids)) + self.logger.info("Starting get_datasets batch fetch for %d unique id(s)", 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 + + # ------------------------------------------------------------------ + # 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/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/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() 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..9e24394 --- /dev/null +++ b/tests/unit/test_annotations_client.py @@ -0,0 +1,741 @@ +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 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.""" + + 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..fab40d9 --- /dev/null +++ b/tests/unit/test_dataset_client.py @@ -0,0 +1,757 @@ +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 + + 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).""" + + 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 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.""" + + 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"]) + + 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() 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()