use ytsaurus_yson::{YsonNode, YsonValue};
use crate::error::{ClientError, Result};
use crate::jobs::{JobInfo, field, text};
use crate::stream::ResponseReader;
use crate::{Client, yson_build};
#[derive(Debug, Clone)]
pub struct Operation {
client: Client,
id: String,
}
impl Operation {
pub(crate) fn new(client: Client, id: String) -> Self {
Self { client, id }
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
#[must_use]
pub fn client(&self) -> &Client {
&self.client
}
pub fn get(&self, attributes: &[&str]) -> Result<YsonValue> {
self.client.get_operation(&self.id, attributes)
}
pub fn state(&self) -> Result<String> {
self.client.operation_state(&self.id)
}
pub fn suspended(&self) -> Result<bool> {
self.client.operation_suspended(&self.id)
}
pub fn status(&self) -> Result<OperationStatus> {
self.client.operation_status(&self.id)
}
pub fn wait(&self) -> Result<()> {
self.client.wait_for_operation(&self.id)
}
pub fn abort(&self, reason: Option<&str>) -> Result<()> {
self.client.abort_operation(&self.id, reason)
}
pub fn suspend(&self, abort_running_jobs: bool) -> Result<()> {
self.client.suspend_operation(&self.id, abort_running_jobs)
}
pub fn resume(&self) -> Result<()> {
self.client.resume_operation(&self.id)
}
pub fn complete(&self) -> Result<()> {
self.client.complete_operation(&self.id)
}
pub fn update_parameters(&self, parameters: &OperationParameters) -> Result<()> {
self.client
.update_operation_parameters(&self.id, parameters)
}
pub fn error(&self) -> Result<Option<String>> {
self.client.operation_result_error(&self.id)
}
pub fn jobs(&self, state: Option<&str>, limit: u32) -> Result<Vec<JobInfo>> {
self.client.list_jobs(&self.id, state, limit)
}
pub fn job(&self, job_id: &str) -> Result<JobInfo> {
self.client.get_job(&self.id, job_id)
}
pub fn job_input(&self, job_id: &str) -> Result<ResponseReader> {
self.client.get_job_input(&self.id, job_id)
}
pub fn job_stderr(&self, job_id: &str) -> Result<Vec<u8>> {
self.client.get_job_stderr(&self.id, job_id)
}
pub fn events(&self) -> Result<Vec<OperationEvent>> {
self.client.list_operation_events(&self.id)
}
pub fn statistics(&self) -> Result<YsonValue> {
self.client.job_statistics(&self.id)
}
pub fn custom_statistics(&self) -> Result<YsonValue> {
self.client.custom_statistics(&self.id)
}
pub fn statistic_sum(&self, name: &str) -> Result<Option<i64>> {
self.client.statistic_sum(&self.id, name)
}
pub fn job_statistic_sum(&self, path: &str) -> Result<Option<i64>> {
self.client.job_statistic_sum(&self.id, path)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OperationInfo {
pub id: String,
pub kind: String,
pub state: String,
pub user: Option<String>,
pub start_time: Option<String>,
pub finish_time: Option<String>,
pub suspended: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OperationStatus {
pub state: String,
pub suspended: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OperationList {
pub operations: Vec<OperationInfo>,
pub incomplete: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OperationEvent {
pub event_type: String,
pub timestamp: Option<String>,
pub incarnation: Option<String>,
}
#[derive(Debug, Clone)]
pub struct OperationFilter {
params: YsonValue,
}
impl Default for OperationFilter {
fn default() -> Self {
Self::new()
}
}
impl OperationFilter {
#[must_use]
pub fn new() -> Self {
Self {
params: yson_build::empty_map(),
}
}
fn set(mut self, key: &str, value: YsonValue) -> Self {
yson_build::insert(&mut self.params, key, value);
self
}
#[must_use]
pub fn with_user(self, user: impl AsRef<str>) -> Self {
self.set("user", yson_build::string(user.as_ref()))
}
#[must_use]
pub fn with_state(self, state: impl AsRef<str>) -> Self {
self.set("state", yson_build::string(state.as_ref()))
}
#[must_use]
pub fn with_kind(self, kind: crate::OperationType) -> Self {
self.set("type", yson_build::string(kind.as_str()))
}
#[must_use]
pub fn with_pool(self, pool: impl AsRef<str>) -> Self {
self.set("pool", yson_build::string(pool.as_ref()))
}
#[must_use]
pub fn with_pool_tree(self, tree: impl AsRef<str>) -> Self {
self.set("pool_tree", yson_build::string(tree.as_ref()))
}
#[must_use]
pub fn with_substring(self, text: impl AsRef<str>) -> Self {
self.set("filter", yson_build::string(text.as_ref()))
}
#[must_use]
pub fn with_from_time(self, time: impl AsRef<str>) -> Self {
self.set("from_time", yson_build::string(time.as_ref()))
}
#[must_use]
pub fn with_to_time(self, time: impl AsRef<str>) -> Self {
self.set("to_time", yson_build::string(time.as_ref()))
}
#[must_use]
pub fn with_failed_jobs(self, with_failed_jobs: bool) -> Self {
self.set("with_failed_jobs", yson_build::boolean(with_failed_jobs))
}
#[must_use]
pub fn with_archive(self, include: bool) -> Self {
self.set("include_archive", yson_build::boolean(include))
}
#[must_use]
pub fn with_limit(self, limit: u32) -> Self {
self.set("limit", yson_build::int(i64::from(limit)))
}
#[must_use]
pub fn with_raw(self, key: impl AsRef<str>, value: YsonValue) -> Self {
self.set(key.as_ref(), value)
}
#[must_use]
pub fn to_yson(&self) -> YsonValue {
self.params.clone()
}
}
#[derive(Debug, Clone)]
pub struct OperationParameters {
params: YsonValue,
}
impl Default for OperationParameters {
fn default() -> Self {
Self::new()
}
}
impl OperationParameters {
#[must_use]
pub fn new() -> Self {
Self {
params: yson_build::empty_map(),
}
}
fn set(mut self, key: &str, value: YsonValue) -> Self {
yson_build::insert(&mut self.params, key, value);
self
}
#[must_use]
pub fn with_pool(self, pool: impl AsRef<str>) -> Self {
self.set("pool", yson_build::string(pool.as_ref()))
}
#[must_use]
pub fn with_weight(self, weight: f64) -> Self {
self.set("weight", yson_build::double(weight))
}
#[must_use]
pub fn with_pool_in_tree(mut self, tree: impl AsRef<str>, pool: impl AsRef<str>) -> Self {
let mut trees = map_or_empty(tree_options(&self.params));
let mut options = map_or_empty(field(&trees, tree.as_ref()));
yson_build::insert(&mut options, "pool", yson_build::string(pool.as_ref()));
yson_build::insert(&mut trees, tree.as_ref(), options);
yson_build::insert(&mut self.params, "scheduling_options_per_pool_tree", trees);
self
}
#[must_use]
pub fn with_raw(self, key: impl AsRef<str>, value: YsonValue) -> Self {
self.set(key.as_ref(), value)
}
#[must_use]
pub fn is_empty(&self) -> bool {
match &self.params.node {
YsonNode::Map(m) => m.is_empty(),
_ => true,
}
}
#[must_use]
pub fn to_yson(&self) -> YsonValue {
self.params.clone()
}
}
fn tree_options(params: &YsonValue) -> Option<&YsonValue> {
field(params, "scheduling_options_per_pool_tree")
}
fn map_or_empty(value: Option<&YsonValue>) -> YsonValue {
match value {
Some(existing) if matches!(existing.node, YsonNode::Map(_)) => existing.clone(),
_ => yson_build::empty_map(),
}
}
pub(crate) fn parse_operations(response: &YsonValue) -> Result<OperationList> {
let Some(YsonNode::List(items)) = field(response, "operations").map(|ops| &ops.node) else {
return Err(ClientError::Decode {
command: "list_operations".to_owned(),
reason: format!(
"the answer carries no `operations` list: {}",
crate::error::truncate(&format!("{:?}", response.node), 300)
),
});
};
Ok(OperationList {
operations: items.iter().filter_map(parse_operation).collect(),
incomplete: flag(field(response, "incomplete")).unwrap_or(false),
})
}
fn parse_operation(operation: &YsonValue) -> Option<OperationInfo> {
let id = text(field(operation, "id")?)?;
Some(OperationInfo {
id,
kind: field(operation, "type")
.and_then(text)
.or_else(|| field(operation, "operation_type").and_then(text))
.unwrap_or_default(),
state: field(operation, "state").and_then(text).unwrap_or_default(),
user: field(operation, "authenticated_user").and_then(text),
start_time: field(operation, "start_time").and_then(text),
finish_time: field(operation, "finish_time").and_then(text),
suspended: flag(field(operation, "suspended")).unwrap_or(false),
})
}
pub(crate) fn parse_events(response: &YsonValue) -> Result<Vec<OperationEvent>> {
let items = match &response.node {
YsonNode::List(items) => items,
_ => match field(response, "events").map(|events| &events.node) {
Some(YsonNode::List(items)) => items,
_ => {
return Err(ClientError::Decode {
command: "list_operation_events".to_owned(),
reason: format!(
"expected a list of events, or a dict holding one under \
`events`: {}",
crate::error::truncate(&format!("{:?}", response.node), 300)
),
});
}
},
};
Ok(items.iter().filter_map(parse_event).collect())
}
fn parse_event(event: &YsonValue) -> Option<OperationEvent> {
Some(OperationEvent {
event_type: text(field(event, "event_type")?)?,
timestamp: field(event, "timestamp").and_then(text),
incarnation: field(event, "incarnation").and_then(text),
})
}
pub(crate) fn flag(value: Option<&YsonValue>) -> Option<bool> {
match value?.node {
YsonNode::Boolean(b) => Some(b),
_ => None,
}
}
pub(crate) fn state_of(document: &YsonValue) -> Result<String> {
match field(document, "state").map(|state| &state.node) {
Some(YsonNode::String(bytes)) => Ok(String::from_utf8_lossy(bytes).into_owned()),
other => Err(ClientError::Decode {
command: "get_operation".to_owned(),
reason: format!("state is missing or not a string: {other:?}"),
}),
}
}
pub(crate) fn suspended_of(document: &YsonValue) -> Result<bool> {
match field(document, "suspended") {
None => Ok(false),
Some(value) => flag(Some(value)).ok_or_else(|| ClientError::Decode {
command: "get_operation".to_owned(),
reason: format!("suspended is not a boolean: {:?}", value.node),
}),
}
}
pub(crate) fn result_error_of(document: &YsonValue) -> Option<String> {
let error = field(document, "result").and_then(|result| field(result, "error"))?;
if field(error, "code").and_then(YsonValue::as_i64) == Some(0) {
return None;
}
crate::jobs::error_summary(error)
}
pub(crate) fn statistics_of(document: &YsonValue) -> YsonValue {
field(document, "progress")
.and_then(|progress| field(progress, "job_statistics"))
.cloned()
.unwrap_or_else(yson_build::empty_map)
}
#[cfg(test)]
mod tests {
use super::*;
use ytsaurus_yson::{YsonFormat, from_slice, to_string};
fn parse(text: &str) -> YsonValue {
from_slice(text.as_bytes(), YsonFormat::Text).expect("valid YSON")
}
fn rendered(value: &YsonValue) -> String {
to_string(value, YsonFormat::Text).expect("encodes")
}
const LIST_OPERATIONS: &str = include_str!("../tests/fixtures/list_operations.yson");
fn operations(text: &str) -> OperationList {
parse_operations(&parse(text)).expect("a well-formed listing")
}
#[test]
fn reads_a_list_captured_from_a_cluster() {
let list = operations(LIST_OPERATIONS);
assert_eq!(list.operations.len(), 2);
assert!(!list.incomplete);
let running = &list.operations[0];
assert_eq!(running.id, "4f5a087b-aac92287-103e8-a74d2331");
assert_eq!(running.kind, "vanilla");
assert_eq!(running.state, "running");
assert_eq!(running.user.as_deref(), Some("root"));
assert!(running.start_time.is_some());
assert_eq!(
running.finish_time, None,
"an operation that has not finished has no finish time, and that \
must stay distinguishable from a time of zero"
);
let finished = &list.operations[1];
assert_eq!(finished.state, "completed");
assert!(finished.finish_time.is_some());
}
#[test]
fn suspension_is_read_from_its_own_field() {
let list =
operations(r#"{"operations"=[{"id"="a-b-c-d";"state"="running";"suspended"=%true}]}"#);
assert_eq!(list.operations[0].state, "running");
assert!(list.operations[0].suspended);
}
#[test]
fn an_operation_without_an_id_is_dropped() {
let list = operations(
r#"{"operations"=[{"state"="running"};{"id"="a-b-c-d"}];"incomplete"=%true}"#,
);
assert_eq!(list.operations.len(), 1);
assert_eq!(list.operations[0].id, "a-b-c-d");
assert!(list.incomplete, "a truncated listing must say so");
}
#[test]
fn a_response_without_an_operation_list_is_an_error() {
assert!(parse_operations(&parse(r#"{"operations"=#}"#)).is_err());
assert!(parse_operations(&parse(r#""not a dict""#)).is_err());
assert!(
parse_operations(&parse(r#"{"operations"=[]}"#)).is_ok(),
"an empty list is a cluster with nothing running, and stays Ok"
);
}
#[test]
fn the_type_falls_back_when_it_is_present_but_unreadable() {
let list =
operations(r#"{"operations"=[{"id"="a-b-c-d";"type"=#;"operation_type"="map"}]}"#);
assert_eq!(list.operations[0].kind, "map");
}
#[test]
fn reads_the_documented_event_list() {
let events = parse_events(&parse(
r#"[
{"timestamp"="2026-08-06T09:21:23.534387Z";"event_type"="started_running"};
{"timestamp"="2026-08-06T09:22:00.000000Z";"event_type"="incarnation_started";
"incarnation"="8fd0b4a1-…"};
]"#,
))
.expect("a bare list is the shape the cluster sent");
assert_eq!(events.len(), 2);
assert_eq!(events[0].event_type, "started_running");
assert_eq!(events[0].incarnation, None);
assert_eq!(events[1].incarnation.as_deref(), Some("8fd0b4a1-…"));
}
#[test]
fn an_enveloped_event_list_is_read_rather_than_dropped() {
let events = parse_events(&parse(r#"{"events"=[{"event_type"="started_running"}]}"#))
.expect("the enveloped shape is accepted too");
assert_eq!(events.len(), 1, "an envelope must not read as no events");
}
#[test]
fn an_empty_event_list_is_not_a_failure() {
assert!(
parse_events(&parse("[]"))
.expect("empty is fine")
.is_empty()
);
}
#[test]
fn an_event_answer_of_neither_shape_is_an_error() {
assert!(parse_events(&parse(r#"{"event_list"=[]}"#)).is_err());
assert!(parse_events(&parse(r#""not a list""#)).is_err());
}
#[test]
fn a_filter_renders_the_keys_the_command_expects() {
let filter = OperationFilter::new()
.with_user("robot")
.with_state("running")
.with_kind(crate::OperationType::Merge)
.with_limit(7);
assert_eq!(
rendered(&filter.to_yson()),
"{limit=7;state=running;type=merge;user=robot}"
);
}
#[test]
fn setting_a_filter_twice_replaces_it() {
let out = rendered(&OperationFilter::new().with_limit(1).with_limit(2).to_yson());
assert_eq!(out, "{limit=2}");
}
#[test]
fn parameters_render_pool_and_weight() {
let out = rendered(
&OperationParameters::new()
.with_pool("fast")
.with_weight(2.5)
.to_yson(),
);
assert_eq!(
out, "{pool=fast;weight=2.5}",
"a weight is a double, and 2.5 must not arrive as an int: {out}"
);
}
#[test]
fn a_pool_can_be_set_for_one_tree_at_a_time() {
let out = rendered(
&OperationParameters::new()
.with_pool_in_tree("default", "fast")
.with_pool_in_tree("gpu", "research")
.to_yson(),
);
assert_eq!(
out, "{scheduling_options_per_pool_tree={default={pool=fast};gpu={pool=research}}}",
"the second tree must not replace the first: {out}"
);
}
#[test]
fn a_pool_is_added_to_what_the_tree_already_carries() {
let out = rendered(
&OperationParameters::new()
.with_raw(
"scheduling_options_per_pool_tree",
yson_build::map([(
"default",
yson_build::map([("weight", yson_build::double(3.0))]),
)]),
)
.with_pool_in_tree("default", "fast")
.to_yson(),
);
assert_eq!(
out, "{scheduling_options_per_pool_tree={default={pool=fast;weight=3.0}}}",
"the weight the caller set must survive: {out}"
);
}
#[test]
fn a_tree_option_that_is_not_a_dict_is_replaced_rather_than_panicked_on() {
let out = rendered(
&OperationParameters::new()
.with_raw(
"scheduling_options_per_pool_tree",
yson_build::string("oops"),
)
.with_pool_in_tree("default", "fast")
.to_yson(),
);
assert_eq!(
out,
"{scheduling_options_per_pool_tree={default={pool=fast}}}"
);
}
#[test]
fn an_empty_update_is_recognisable() {
assert!(OperationParameters::new().is_empty());
assert!(!OperationParameters::new().with_weight(1.0).is_empty());
}
}