#![warn(missing_docs)]
use std::time::{Duration, Instant};
mod batch;
pub mod error;
mod http;
mod jobs;
pub mod lock;
mod observe;
pub mod operation;
pub mod path;
mod retry;
pub mod schema;
mod spec;
pub mod stream;
pub mod trace;
mod transaction;
mod unique;
mod worker;
pub mod yson_build;
pub use crate::batch::BatchRequest;
pub use crate::error::{ClientError, RedirectRefusal, Result};
pub use crate::http::Method;
pub use crate::jobs::{JobFailure, JobInfo};
pub use crate::lock::{Lock, LockMode};
pub use crate::operation::{
Operation, OperationEvent, OperationFilter, OperationInfo, OperationList, OperationParameters,
OperationStatus,
};
pub use crate::path::{Key, RowRange, TablePath};
pub use crate::retry::{MutationId, Repeatable, RetryPolicy};
pub use crate::schema::{Column, ColumnType, SortOrder, TableRow, TableSchema};
pub use crate::spec::{
EraseSpec, MapReduceSpec, MapSpec, MergeMode, MergeSpec, OperationType, ReduceSpec,
RemoteCopySpec, SortSpec, VanillaSpec, VanillaTask,
};
pub use crate::stream::{FileReader, ResponseReader, TableReader};
pub use crate::trace::TraceContext;
pub use crate::transaction::Transaction;
pub use ytsaurus_format::DataFormat;
#[cfg(feature = "derive")]
pub use ytsaurus_helpers::TableRow;
pub use ytsaurus_skiff::{
Format as SkiffFormat, Schema as SkiffSchema, SchemaRef as SkiffSchemaRef,
WireType as SkiffWireType,
};
use crate::http::{Payload, Transport};
use ytsaurus_skiff::Decoder as SkiffDecoder;
use ytsaurus_yson::{YsonFormat, YsonNode, YsonValue, from_slice};
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(120);
const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(2);
const REPORTED_JOBS: u32 = 3;
const STDERR_EXCERPT: usize = 4096;
const DEFAULT_FILE_CACHE: &str = "//tmp/yt_wrapper/file_storage/new_cache";
const UNCACHED_UPLOAD_DIR: &str = "//tmp";
const ACCESS_DENIED: i64 = 901;
#[derive(serde::Deserialize)]
struct Envelope<T> {
value: T,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CachedFile {
pub path: String,
pub name: String,
pub uploaded: bool,
pub cached: bool,
}
enum Cached {
At(String),
Refused(ClientError),
}
#[derive(Debug, Clone)]
pub struct Client {
transport: Transport,
poll_interval: Duration,
job_diagnostics: bool,
file_cache: String,
}
impl Client {
#[must_use]
pub fn new(proxy: &str) -> Self {
Self {
transport: Transport::new(proxy, None, DEFAULT_TIMEOUT),
poll_interval: DEFAULT_POLL_INTERVAL,
job_diagnostics: true,
file_cache: DEFAULT_FILE_CACHE.to_owned(),
}
}
#[must_use]
pub fn with_token(proxy: &str, token: impl Into<String>) -> Self {
Self {
transport: Transport::new(proxy, Some(token.into()), DEFAULT_TIMEOUT),
poll_interval: DEFAULT_POLL_INTERVAL,
job_diagnostics: true,
file_cache: DEFAULT_FILE_CACHE.to_owned(),
}
}
pub fn from_env() -> Result<Self> {
Self::from_lookup(environment_value)
}
fn from_lookup(lookup: impl Fn(&str) -> Option<String>) -> Result<Self> {
let value = |name: &str| {
lookup(name)
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
};
let proxy = value("YT_PROXY").ok_or_else(|| {
ClientError::Config(
"YT_PROXY is not set; export it (for a local cluster: \
YT_PROXY=http://localhost:8000) or use Client::new"
.to_owned(),
)
})?;
let proxy = expanded_proxy(&proxy, value("YT_PROXY_SUFFIX").as_deref());
let mut client = match token_from_environment() {
Some(token) => Self::with_token(&proxy, token),
None => Self::new(&proxy),
};
if let Some(domains) = value("YT_HEAVY_PROXY_DOMAINS") {
client = client.with_heavy_proxies_under(split_domains(&domains));
}
if value("YT_HEAVY_PROXIES_ANYWHERE").is_some_and(|value| truthy(&value)) {
client = client.with_heavy_proxies_anywhere(true);
}
if let Some(cache) = value("YT_FILE_CACHE") {
client = client.with_file_cache(cache);
}
Ok(client)
}
#[must_use]
pub fn with_poll_interval(mut self, interval: Duration) -> Self {
self.poll_interval = interval;
self
}
#[must_use]
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.transport.set_timeout(timeout);
self
}
#[must_use]
pub fn with_retries(mut self, policy: RetryPolicy) -> Self {
self.transport.set_retries(policy);
self
}
#[must_use]
pub fn with_file_cache(mut self, path: impl Into<String>) -> Self {
self.file_cache = path.into();
self
}
#[must_use]
pub fn with_proxy_discovery(mut self, enabled: bool) -> Self {
self.transport.set_proxy_discovery(enabled);
self
}
#[must_use]
pub fn with_heavy_proxies_anywhere(mut self, enabled: bool) -> Self {
self.transport.set_heavy_proxies_anywhere(enabled);
self
}
#[must_use]
pub fn with_heavy_proxies_in<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.transport
.set_heavy_proxies_in(names.into_iter().map(Into::into).collect());
self
}
#[must_use]
pub fn with_heavy_proxies_under<I, S>(mut self, domains: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.transport
.set_heavy_proxies_under(domains.into_iter().map(Into::into).collect());
self
}
#[must_use]
pub fn with_hosts_timeout(mut self, timeout: Duration) -> Self {
self.transport.set_hosts_timeout(timeout);
self
}
#[must_use]
pub fn with_hosts_retry_after(mut self, after: Duration) -> Self {
self.transport.set_hosts_retry_after(after);
self
}
#[must_use]
pub fn with_host_list_refresh_interval(mut self, interval: Duration) -> Self {
self.transport.set_host_list_refresh_interval(interval);
self
}
#[must_use]
pub fn with_job_diagnostics(mut self, enabled: bool) -> Self {
self.job_diagnostics = enabled;
self
}
#[must_use]
pub fn with_transaction(mut self, id: impl Into<String>) -> Self {
self.transport.set_transaction(Some(id.into()));
self
}
#[must_use]
pub fn transaction_id(&self) -> Option<&str> {
self.transport.transaction()
}
#[must_use]
pub fn with_trace_context(mut self, context: &TraceContext) -> Self {
self.transport.set_trace(context);
self
}
#[must_use]
pub fn traceparent(&self) -> Option<&str> {
self.transport.trace()
}
#[must_use]
pub fn tracestate(&self) -> Option<&str> {
self.transport.tracestate()
}
pub fn start_transaction(&self) -> Result<Transaction> {
Transaction::start(self, transaction::DEFAULT_TRANSACTION_TIMEOUT)
}
pub fn start_transaction_with(&self, timeout: Duration) -> Result<Transaction> {
Transaction::start(self, timeout)
}
pub fn attach_transaction(&self, id: &str) -> Result<Transaction> {
Transaction::attach(self, id.to_owned())
}
pub fn ping_transaction(&self, id: &str) -> Result<()> {
transaction::ping(self, id)
}
pub fn commit_transaction(&self, id: &str) -> Result<()> {
transaction::commit_by_id(self, id)
}
pub fn abort_transaction(&self, id: &str) -> Result<()> {
transaction::abort_by_id(self, id)
}
pub fn heavy_proxy(&self) -> Result<Option<String>> {
Ok(self.transport.heavy_hosts()?.into_iter().next())
}
pub fn exists(&self, path: &str) -> Result<bool> {
let params = yson_build::map([("path", yson_build::string(path))]);
let body = self.transport.call(
Method::Get,
"exists",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
Ok(matches!(
self.value_field(&body, "value")?.node,
YsonNode::Boolean(true)
))
}
pub fn create(&self, node_type: &str, path: &str) -> Result<()> {
let params = yson_build::map([
("path", yson_build::string(path)),
("type", yson_build::string(node_type)),
("recursive", yson_build::boolean(true)),
("ignore_existing", yson_build::boolean(true)),
]);
self.transport.call(
Method::Post,
"create",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn create_table(&self, path: &str, schema: &TableSchema) -> Result<()> {
schema
.validate()
.map_err(|reason| ClientError::Config(format!("{path}: {reason}")))?;
let params = yson_build::map([
("path", yson_build::string(path)),
("type", yson_build::string("table")),
("recursive", yson_build::boolean(true)),
(
"attributes",
yson_build::map([("schema", schema.to_yson())]),
),
]);
self.transport.call(
Method::Post,
"create",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn alter_table(&self, path: &str, schema: &TableSchema) -> Result<()> {
schema
.validate()
.map_err(|reason| ClientError::Config(format!("{path}: {reason}")))?;
let params = yson_build::map([
("path", yson_build::string(path)),
("schema", schema.to_yson()),
]);
self.transport.call(
Method::Post,
"alter_table",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn table_schema(&self, path: &str) -> Result<YsonValue> {
self.get(&format!("{path}/@schema"))
}
pub fn remove(&self, path: &str) -> Result<()> {
self.remove_with(path, false, false)
}
pub fn remove_tree(&self, path: &str) -> Result<()> {
self.remove_with(path, true, true)
}
fn remove_with(&self, path: &str, recursive: bool, force: bool) -> Result<()> {
let params = yson_build::map([
("path", yson_build::string(path)),
("recursive", yson_build::boolean(recursive)),
("force", yson_build::boolean(force)),
]);
self.transport.call(
Method::Post,
"remove",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn list(&self, path: &str) -> Result<Vec<String>> {
let params = yson_build::map([("path", yson_build::string(path))]);
let body = self.transport.call(
Method::Get,
"list",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
child_names(&self.value_field(&body, "value")?, path)
}
pub fn copy(&self, source: &str, destination: &str) -> Result<()> {
self.transfer("copy", source, destination, false)
}
pub fn copy_replacing(&self, source: &str, destination: &str) -> Result<()> {
self.transfer("copy", source, destination, true)
}
pub fn move_node(&self, source: &str, destination: &str) -> Result<()> {
self.transfer("move", source, destination, false)
}
pub fn move_replacing(&self, source: &str, destination: &str) -> Result<()> {
self.transfer("move", source, destination, true)
}
fn transfer(&self, command: &str, source: &str, destination: &str, force: bool) -> Result<()> {
let params = yson_build::map([
("source_path", yson_build::string(source)),
("destination_path", yson_build::string(destination)),
("recursive", yson_build::boolean(true)),
("force", yson_build::boolean(force)),
]);
self.transport.call(
Method::Post,
command,
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn link(&self, target: &str, link_path: &str) -> Result<()> {
self.link_inner(target, link_path, false)
}
pub fn link_replacing(&self, target: &str, link_path: &str) -> Result<()> {
self.link_inner(target, link_path, true)
}
fn link_inner(&self, target: &str, link_path: &str, force: bool) -> Result<()> {
let params = yson_build::map([
("target_path", yson_build::string(target)),
("link_path", yson_build::string(link_path)),
("recursive", yson_build::boolean(true)),
("force", yson_build::boolean(force)),
]);
self.transport.call(
Method::Post,
"link",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn lock(&self, path: &str, mode: LockMode) -> Result<Lock> {
self.lock_inner(path, mode, false)
}
pub fn lock_waiting(&self, path: &str, mode: LockMode, wait_for: Duration) -> Result<Lock> {
let lock = self.lock_inner(path, mode, true)?;
let deadline = Instant::now() + wait_for;
loop {
let state = self.get(&format!("#{}/@state", lock.id))?;
if state.as_str() == Some("acquired") {
return Ok(lock);
}
if Instant::now() >= deadline {
return Err(ClientError::Config(format!(
"lock on {path}: still {} after {:.0}s — the locks ahead of it are \
still held, which can include a snapshot lock this same \
transaction took. It stays queued until this transaction ends.",
state.as_str().unwrap_or("queued"),
wait_for.as_secs_f64()
)));
}
std::thread::sleep(self.poll_interval);
}
}
fn lock_inner(&self, path: &str, mode: LockMode, waitable: bool) -> Result<Lock> {
if self.transaction_id().is_none() {
return Err(ClientError::Config(format!(
"lock {path}: a lock belongs to a transaction, and this client is not in \
one — take it through a Client::start_transaction handle. The cluster \
answers this with `A valid master transaction is required`."
)));
}
let params = yson_build::map([
("path", yson_build::string(path)),
("mode", yson_build::string(mode.as_str())),
("waitable", yson_build::boolean(waitable)),
]);
let body = self.transport.call(
Method::Post,
"lock",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
let envelope = self.strip_envelope(&body, "lock")?;
let text = |key: &str| -> Result<String> {
match &self.field_of(&envelope, key)?.node {
YsonNode::String(bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
other => Err(ClientError::Decode {
command: "lock".to_owned(),
reason: format!("{key} is not a string: {other:?}"),
}),
}
};
Ok(Lock {
id: text("lock_id")?,
node_id: text("node_id")?,
})
}
pub fn get(&self, path: &str) -> Result<YsonValue> {
let params = yson_build::map([("path", yson_build::string(path))]);
let body = self.transport.call(
Method::Get,
"get",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
self.value_field(&body, "value")
}
pub fn row_count(&self, path: &str) -> Result<i64> {
let value = self.get(&format!("{path}/@row_count"))?;
value.as_i64().ok_or_else(|| ClientError::Decode {
command: "get".to_owned(),
reason: format!("{path}/@row_count is not an integer"),
})
}
pub fn execute_batch(&self, batch: &BatchRequest) -> Result<Vec<Result<YsonValue>>> {
self.execute_batch_with(batch, None)
}
pub fn execute_batch_with(
&self,
batch: &BatchRequest,
mutation_id: Option<&MutationId>,
) -> Result<Vec<Result<YsonValue>>> {
if batch.is_empty() {
return Err(ClientError::Config(
"an empty batch is not a request worth sending: the cluster \
would answer with no results, and reporting that as success \
would call a no-op work done"
.to_owned(),
));
}
let max_part_size = batch.max_part_size();
if mutation_id.is_some() && batch.len() > max_part_size {
return Err(ClientError::Config(format!(
"a batch of {} parts is sent as several requests at {max_part_size} \
parts each, and one mutation id cannot cover them: the cluster \
derives each part's id by incrementing the batch's, so a second \
request under the same id would be answered with the first \
request's results. Raise with_max_part_size past {}, or send it \
without an id.",
batch.len(),
batch.len()
)));
}
let repeatable = batch.repeatable();
let mut results = Vec::with_capacity(batch.len());
for chunk in batch.parts().chunks(max_part_size) {
let answered = batch::render_chunk(chunk, batch.concurrency(), self.transaction_id())
.and_then(|body| {
self.transport.call_with(
Method::Post,
"execute_batch",
&yson_build::empty_map(),
Payload::Bytes(&body),
repeatable,
mutation_id,
)
})
.and_then(|answer| batch::parse_results(&answer, chunk));
match answered {
Ok(answers) => results.extend(answers),
Err(cause) if results.is_empty() => return Err(cause),
Err(cause) => {
return Err(ClientError::BatchInterrupted {
answered: results,
parts: batch.len(),
cause: Box::new(cause),
});
}
}
}
Ok(results)
}
pub fn upload_worker(&self, local: impl AsRef<std::path::Path>, remote: &str) -> Result<()> {
let local = local.as_ref();
let bytes = std::fs::read(local).map_err(|source| ClientError::Io {
path: local.display().to_string(),
source,
})?;
self.upload_executable(remote, &bytes)
}
pub fn upload_current_exe(&self, remote: &str) -> Result<()> {
let exe = std::env::current_exe().map_err(|source| ClientError::Io {
path: "the running executable".to_owned(),
source,
})?;
let bytes = std::fs::read(&exe).map_err(|source| ClientError::Io {
path: exe.display().to_string(),
source,
})?;
if let Err(reason) = worker::check_worker_binary(&bytes) {
return Err(ClientError::NotAWorker {
path: exe.display().to_string(),
reason,
});
}
self.upload_executable(remote, &bytes)
}
pub fn upload_worker_cached(&self, local: impl AsRef<std::path::Path>) -> Result<CachedFile> {
let local = local.as_ref();
let bytes = std::fs::read(local).map_err(|source| ClientError::Io {
path: local.display().to_string(),
source,
})?;
let name = local
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| "worker".to_owned());
let digest = format!("{:x}", md5::compute(&bytes));
if let Some(path) = self.file_from_cache(&digest)? {
return Ok(CachedFile {
path,
name,
uploaded: false,
cached: true,
});
}
let (path, cached) = match self.upload_into_cache(&bytes, &digest)? {
Cached::At(path) => {
self.set_attribute(&path, "executable", yson_build::boolean(true))?;
(path, true)
}
Cached::Refused(denial) => {
observe::cache_refused(&self.file_cache, &denial);
(self.upload_uncached(&digest, &bytes)?, false)
}
};
Ok(CachedFile {
path,
name,
uploaded: true,
cached,
})
}
fn upload_into_cache(&self, bytes: &[u8], digest: &str) -> Result<Cached> {
if let Err(denial) = self.create("map_node", &self.file_cache) {
return refused_or_reported(denial);
}
let staging = format!("{}/staged_{digest}_{}", self.file_cache, MutationId::new());
if let Err(denial) = self.create("file", &staging) {
return refused_or_reported(denial);
}
let cached = self
.write_file_computing_md5(&staging, bytes)
.and_then(|()| self.set_attribute(&staging, "executable", yson_build::boolean(true)))
.and_then(|()| self.put_file_to_cache(&staging, digest));
let removed = self.remove_tree(&staging);
match cached {
Ok(path) => {
removed?;
Ok(Cached::At(path))
}
Err(denial) if denied(&denial, "put_file_to_cache") => Ok(Cached::Refused(denial)),
Err(failed) => Err(failed),
}
}
fn upload_uncached(&self, digest: &str, bytes: &[u8]) -> Result<String> {
let remote = format!(
"{UNCACHED_UPLOAD_DIR}/ytsaurus_rs_worker_{digest}_{}",
MutationId::new()
);
self.upload_executable(&remote, bytes)?;
Ok(remote)
}
pub fn file_from_cache(&self, md5: &str) -> Result<Option<String>> {
let params = yson_build::map([
("md5", yson_build::string(md5)),
("cache_path", yson_build::string(&self.file_cache)),
]);
let body = self.transport.call(
Method::Get,
"get_file_from_cache",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
self.cached_path(&body, "get_file_from_cache")
}
pub fn put_file_to_cache(&self, path: &str, md5: &str) -> Result<String> {
let params = yson_build::map([
("path", yson_build::string(path)),
("md5", yson_build::string(md5)),
("cache_path", yson_build::string(&self.file_cache)),
]);
let body = self.transport.call(
Method::Post,
"put_file_to_cache",
¶ms,
Payload::None,
Repeatable::WithMutationId,
)?;
self.cached_path(&body, "put_file_to_cache")?
.ok_or_else(|| ClientError::Decode {
command: "put_file_to_cache".to_owned(),
reason: "the cluster returned no path for the cached file".to_owned(),
})
}
fn cached_path(&self, body: &[u8], command: &str) -> Result<Option<String>> {
let value = self.strip_envelope(body, command)?;
let value = match &value.node {
YsonNode::Map(_) => self.field_of(&value, "path")?,
_ => value,
};
match &value.node {
YsonNode::String(bytes) if !bytes.is_empty() => {
Ok(Some(String::from_utf8_lossy(bytes).into_owned()))
}
YsonNode::String(_) | YsonNode::Entity => Ok(None),
other => Err(ClientError::Decode {
command: command.to_owned(),
reason: format!("the cached path is not a string: {other:?}"),
}),
}
}
fn upload_executable(&self, remote: &str, bytes: &[u8]) -> Result<()> {
self.create("file", remote)?;
self.write_file(remote, bytes)?;
self.set_attribute(remote, "executable", yson_build::boolean(true))
}
pub fn write_file(&self, path: &str, contents: &[u8]) -> Result<()> {
self.write_file_inner(path, contents, false)
}
fn write_file_computing_md5(&self, path: &str, contents: &[u8]) -> Result<()> {
self.write_file_inner(path, contents, true)
}
fn write_file_inner(&self, path: &str, contents: &[u8], compute_md5: bool) -> Result<()> {
let mut params = yson_build::map([("path", yson_build::string(path))]);
if compute_md5 {
yson_build::insert(&mut params, "compute_md5", yson_build::boolean(true));
}
self.transport.call(
Method::Put,
"write_file",
¶ms,
Payload::Bytes(contents),
Repeatable::Heavy,
)?;
Ok(())
}
pub fn read_file(&self, path: &str) -> Result<Vec<u8>> {
let params = yson_build::map([("path", yson_build::string(path))]);
let body = self.transport.call(
Method::Get,
"read_file",
¶ms,
Payload::None,
Repeatable::Heavy,
)?;
let recorded = self.file_size(path)?;
if recorded != body.len() as i64 {
return Err(ClientError::Decode {
command: "read_file".to_owned(),
reason: format!(
"{path}: the cluster records {recorded} bytes but the response carried {}; \
either the stream was cut short — the proxy says so in a trailer this \
client cannot read — or the file was rewritten while it was being read",
body.len()
),
});
}
Ok(body)
}
fn file_size(&self, path: &str) -> Result<i64> {
let size = self
.get(&format!("{path}/@uncompressed_data_size"))
.map_err(|error| ClientError::Decode {
command: "read_file".to_owned(),
reason: format!(
"the file's bytes arrived, but the size they were to be checked \
against could not be read: {error}"
),
})?;
size.as_i64().ok_or_else(|| ClientError::Decode {
command: "read_file".to_owned(),
reason: format!(
"{path}/@uncompressed_data_size is not an integer: {:?}; without it the \
response cannot be checked for truncation",
size.node
),
})
}
pub fn read_file_streaming(&self, path: &str) -> Result<FileReader> {
let params = yson_build::map([("path", yson_build::string(path))]);
let body = self.transport.open(Method::Get, "read_file", ¶ms)?;
Ok(FileReader::new(body))
}
pub fn set_attribute(&self, path: &str, name: &str, value: YsonValue) -> Result<()> {
let encoded =
ytsaurus_yson::to_vec(&value, YsonFormat::Binary).map_err(|e| ClientError::Decode {
command: "set".to_owned(),
reason: format!("could not encode the attribute: {e}"),
})?;
let params = yson_build::map([
("path", yson_build::string(format!("{path}/@{name}"))),
("input_format", yson_build::binary_yson_format()),
]);
self.transport.call(
Method::Put,
"set",
¶ms,
Payload::Bytes(&encoded),
Repeatable::WithMutationId,
)?;
Ok(())
}
pub fn write_table(&self, path: impl Into<TablePath>, rows: &[u8]) -> Result<()> {
self.write_table_with_format(path, rows, &DataFormat::binary_yson())
}
pub fn write_table_with_format(
&self,
path: impl Into<TablePath>,
rows: &[u8],
format: &DataFormat,
) -> Result<()> {
let path = path.into();
match format {
DataFormat::Yson(format) => self.write_yson_table(&path, rows, *format),
DataFormat::Skiff(format) => self.write_skiff_table_impl(&path, rows, format),
_ => Err(unsupported_data_format()),
}
}
fn write_yson_table(&self, path: &TablePath, rows: &[u8], format: YsonFormat) -> Result<()> {
refuse_selection_on_write(path)?;
let params = yson_build::map([
("path", path.to_yson()),
("input_format", DataFormat::yson(format).to_yson()),
]);
self.transport.call(
Method::Put,
"write_table",
¶ms,
Payload::Bytes(rows),
Repeatable::Heavy,
)?;
Ok(())
}
pub fn write_skiff_table(
&self,
path: impl Into<TablePath>,
rows: &[u8],
format: &SkiffFormat,
) -> Result<()> {
self.write_table_with_format(path, rows, &DataFormat::skiff(format.clone()))
}
fn write_skiff_table_impl(
&self,
path: &TablePath,
rows: &[u8],
format: &SkiffFormat,
) -> Result<()> {
refuse_selection_on_write(path)?;
let path_value = skiff_table_path(path, format)?;
check_complete_skiff_stream(rows, format).map_err(|reason| ClientError::Decode {
command: "write_table".to_owned(),
reason: format!("{}: {reason}", path.as_str()),
})?;
let params = yson_build::map([("path", path_value), ("input_format", format.to_yson())]);
self.transport.call(
Method::Put,
"write_table",
¶ms,
Payload::Bytes(rows),
Repeatable::Heavy,
)?;
Ok(())
}
pub fn read_table(&self, path: impl Into<TablePath>) -> Result<Vec<u8>> {
self.read_table_with_format(path, &DataFormat::binary_yson())
}
pub fn read_table_with_format(
&self,
path: impl Into<TablePath>,
format: &DataFormat,
) -> Result<Vec<u8>> {
let path = path.into();
match format {
DataFormat::Yson(format) => self.read_yson_table(&path, *format),
DataFormat::Skiff(format) => self.read_skiff_table_impl(&path, format),
_ => Err(unsupported_data_format()),
}
}
fn read_yson_table(&self, path: &TablePath, format: YsonFormat) -> Result<Vec<u8>> {
refuse_mixed_selection_on_read(path)?;
let params = yson_build::map([
("path", path.to_yson()),
("output_format", DataFormat::yson(format).to_yson()),
]);
let body = self.transport.call(
Method::Get,
"read_table",
¶ms,
Payload::None,
Repeatable::Heavy,
)?;
check_complete_yson_fragment(&body, format).map_err(|reason| ClientError::Decode {
command: "read_table".to_owned(),
reason: format!("{path}: {reason}"),
})?;
Ok(body)
}
pub fn read_skiff_table(
&self,
path: impl Into<TablePath>,
format: &SkiffFormat,
) -> Result<Vec<u8>> {
self.read_table_with_format(path, &DataFormat::skiff(format.clone()))
}
fn read_skiff_table_impl(&self, path: &TablePath, format: &SkiffFormat) -> Result<Vec<u8>> {
refuse_mixed_selection_on_read(path)?;
let params = yson_build::map([
("path", skiff_table_path(path, format)?),
("output_format", format.to_yson()),
]);
let body = self.transport.call(
Method::Get,
"read_table",
¶ms,
Payload::None,
Repeatable::Heavy,
)?;
check_complete_skiff_stream(&body, format).map_err(|reason| ClientError::Decode {
command: "read_table".to_owned(),
reason: format!("{path}: {reason}"),
})?;
Ok(body)
}
pub fn write_table_rows<T, I>(&self, path: impl Into<TablePath>, rows: I) -> Result<()>
where
T: serde::Serialize,
I: IntoIterator<Item = T>,
{
let path = path.into();
refuse_selection_on_write(&path)?;
let params = yson_build::map([
("path", path.to_yson()),
("input_format", yson_build::binary_yson_format()),
]);
let mut stream = stream::RowStream::new(rows.into_iter());
let sent = self
.transport
.upload(Method::Put, "write_table", ¶ms, &mut stream);
if let Some(reason) = stream.failed {
return Err(ClientError::Decode {
command: "write_table".to_owned(),
reason: format!("{path}: {reason}"),
});
}
sent.map(|_| ())
}
pub fn read_table_rows<T: serde::de::DeserializeOwned>(
&self,
path: impl Into<TablePath>,
) -> Result<Vec<T>> {
let path = path.into();
decode_rows(&self.read_table(&path)?, &path.to_string())
}
pub fn get_as<T: serde::de::DeserializeOwned>(&self, path: &str) -> Result<T> {
let params = yson_build::map([("path", yson_build::string(path))]);
let body = self.transport.call(
Method::Get,
"get",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
let envelope: Envelope<T> =
from_slice(&body, YsonFormat::Text).map_err(|e| ClientError::Decode {
command: "get".to_owned(),
reason: format!(
"{path}: the answer does not fit the type asked for: {e}; body was {}",
crate::error::truncate(&String::from_utf8_lossy(&body), 200)
),
})?;
Ok(envelope.value)
}
pub fn read_table_streaming(&self, path: impl Into<TablePath>) -> Result<TableReader> {
let path = path.into();
refuse_mixed_selection_on_read(&path)?;
let params = yson_build::map([
("path", path.to_yson()),
("output_format", yson_build::binary_yson_format()),
]);
let body = self.transport.open(Method::Get, "read_table", ¶ms)?;
Ok(TableReader::new(body))
}
pub fn write_table_streaming(
&self,
path: impl Into<TablePath>,
mut rows: impl std::io::Read,
) -> Result<()> {
let path = path.into();
refuse_selection_on_write(&path)?;
let params = yson_build::map([
("path", path.to_yson()),
("input_format", yson_build::binary_yson_format()),
]);
self.transport
.upload(Method::Put, "write_table", ¶ms, &mut rows)?;
Ok(())
}
pub fn start_map(&self, spec: &MapSpec) -> Result<String> {
refuse_skiff_table_mismatch(spec.skiff_table_mismatch())?;
self.start_operation(OperationType::Map, &spec.to_yson())
}
pub fn start_map_reduce(&self, spec: &MapReduceSpec) -> Result<String> {
refuse_skiff_table_mismatch(spec.skiff_table_mismatch())?;
self.start_operation(OperationType::MapReduce, &spec.to_yson())
}
pub fn start_reduce(&self, spec: &ReduceSpec) -> Result<String> {
refuse_skiff_table_mismatch(spec.skiff_table_mismatch())?;
self.start_operation(OperationType::Reduce, &spec.to_yson())
}
pub fn start_sort(&self, spec: &SortSpec) -> Result<String> {
self.start_operation(OperationType::Sort, &spec.to_yson())
}
pub fn start_vanilla(&self, spec: &VanillaSpec) -> Result<String> {
if let Some(name) = spec.duplicate_task() {
return Err(ClientError::Config(format!(
"two vanilla tasks are both called {name:?}; a spec keys its tasks \
by name, so the second would replace the first and its jobs would \
never run"
)));
}
refuse_skiff_table_mismatch(spec.skiff_table_mismatch())?;
self.start_operation(OperationType::Vanilla, &spec.to_yson())
}
pub fn start_merge(&self, spec: &MergeSpec) -> Result<String> {
self.start_operation(OperationType::Merge, &spec.to_yson())
}
pub fn start_erase(&self, spec: &EraseSpec) -> Result<String> {
self.start_operation(OperationType::Erase, &spec.to_yson())
}
pub fn start_remote_copy(&self, spec: &RemoteCopySpec) -> Result<String> {
self.start_operation(OperationType::RemoteCopy, &spec.to_yson())
}
pub fn start_operation(&self, kind: OperationType, spec: &YsonValue) -> Result<String> {
self.start_operation_inner(kind, spec, None)
}
pub fn start_operation_with(
&self,
kind: OperationType,
spec: &YsonValue,
mutation_id: &MutationId,
) -> Result<String> {
self.start_operation_inner(kind, spec, Some(mutation_id))
}
fn start_operation_inner(
&self,
kind: OperationType,
spec: &YsonValue,
mutation_id: Option<&MutationId>,
) -> Result<String> {
let params = yson_build::map([
("operation_type", yson_build::string(kind.as_str())),
("spec", spec.clone()),
]);
let body = self.transport.call_with(
Method::Post,
"start_operation",
¶ms,
Payload::None,
Repeatable::WithMutationId,
mutation_id,
)?;
let value = self.value_field(&body, "operation_id")?;
match &value.node {
YsonNode::String(bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
other => Err(ClientError::Decode {
command: "start_operation".to_owned(),
reason: format!("operation_id is not a string: {other:?}"),
}),
}
}
pub fn abort_operation(&self, id: &str, reason: Option<&str>) -> Result<()> {
let mut params = yson_build::map([("operation_id", yson_build::string(id))]);
if let Some(reason) = reason {
yson_build::insert(&mut params, "abort_message", yson_build::string(reason));
}
self.transport.call(
Method::Post,
"abort_operation",
¶ms,
Payload::None,
Repeatable::Never,
)?;
Ok(())
}
pub fn suspend_operation(&self, id: &str, abort_running_jobs: bool) -> Result<()> {
let params = yson_build::map([
("operation_id", yson_build::string(id)),
(
"abort_running_jobs",
yson_build::boolean(abort_running_jobs),
),
]);
self.transport.call(
Method::Post,
"suspend_operation",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
Ok(())
}
pub fn resume_operation(&self, id: &str) -> Result<()> {
let params = yson_build::map([("operation_id", yson_build::string(id))]);
self.transport.call(
Method::Post,
"resume_operation",
¶ms,
Payload::None,
Repeatable::Never,
)?;
Ok(())
}
pub fn complete_operation(&self, id: &str) -> Result<()> {
let params = yson_build::map([("operation_id", yson_build::string(id))]);
self.transport.call(
Method::Post,
"complete_operation",
¶ms,
Payload::None,
Repeatable::Never,
)?;
Ok(())
}
pub fn update_operation_parameters(
&self,
id: &str,
parameters: &OperationParameters,
) -> Result<()> {
if parameters.is_empty() {
return Err(ClientError::Config(
"update_operation_parameters was given nothing to change; the \
cluster answers 200 and does nothing, so this is refused here \
instead"
.to_owned(),
));
}
let params = yson_build::map([
("operation_id", yson_build::string(id)),
("parameters", parameters.to_yson()),
]);
self.transport.call(
Method::Post,
"update_operation_parameters",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
Ok(())
}
pub fn list_operations(&self, filter: &OperationFilter) -> Result<OperationList> {
let body = self.transport.call(
Method::Get,
"list_operations",
&filter.to_yson(),
Payload::None,
Repeatable::Freely,
)?;
operation::parse_operations(&self.strip_envelope(&body, "list_operations")?)
}
pub fn list_operation_events(&self, id: &str) -> Result<Vec<OperationEvent>> {
let params = yson_build::map([("operation_id", yson_build::string(id))]);
let body = self.transport.call(
Method::Get,
"list_operation_events",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
operation::parse_events(&self.strip_envelope(&body, "list_operation_events")?)
}
#[must_use]
pub fn attach_operation(&self, id: impl Into<String>) -> Operation {
let mut id = id.into();
if id.trim().len() != id.len() {
id = id.trim().to_owned();
}
Operation::new(self.clone(), id)
}
pub fn get_operation(&self, id: &str, attributes: &[&str]) -> Result<YsonValue> {
self.get_operation_inner(
yson_build::map([("operation_id", yson_build::string(id))]),
attributes,
)
}
pub fn get_operation_by_alias(&self, alias: &str, attributes: &[&str]) -> Result<YsonValue> {
self.get_operation_inner(
yson_build::map([
("operation_alias", yson_build::string(alias)),
("include_runtime", yson_build::boolean(true)),
]),
attributes,
)
}
fn get_operation_inner(&self, params: YsonValue, attributes: &[&str]) -> Result<YsonValue> {
let body = self.get_operation_body(params, attributes)?;
self.strip_envelope(&body, "get_operation")
}
fn get_operation_body(&self, mut params: YsonValue, attributes: &[&str]) -> Result<Vec<u8>> {
if !attributes.is_empty() {
yson_build::insert(
&mut params,
"attributes",
yson_build::list(attributes.iter().map(yson_build::string)),
);
}
self.transport.call(
Method::Get,
"get_operation",
¶ms,
Payload::None,
Repeatable::Freely,
)
}
pub fn operation_state(&self, id: &str) -> Result<String> {
operation::state_of(&self.get_operation(id, &["state"])?)
}
pub fn operation_suspended(&self, id: &str) -> Result<bool> {
operation::suspended_of(&self.get_operation(id, &["suspended"])?)
}
pub fn operation_status(&self, id: &str) -> Result<OperationStatus> {
let document = self.get_operation(id, &["state", "suspended"])?;
Ok(OperationStatus {
state: operation::state_of(&document)?,
suspended: operation::suspended_of(&document)?,
})
}
pub fn custom_statistics(&self, operation_id: &str) -> Result<YsonValue> {
let all = self.job_statistics(operation_id)?;
Ok(jobs::field(&all, "custom").cloned().unwrap_or(YsonValue {
attributes: None,
node: YsonNode::Map(std::collections::BTreeMap::new()),
}))
}
pub fn job_statistics(&self, operation_id: &str) -> Result<YsonValue> {
Ok(operation::statistics_of(
&self.get_operation(operation_id, &["progress"])?,
))
}
pub fn job_statistic_sum(&self, operation_id: &str, path: &str) -> Result<Option<i64>> {
let statistics = self.job_statistics(operation_id)?;
let mut node = &statistics;
for component in path.split('/') {
match jobs::field(node, component) {
Some(next) => node = next,
None => return Ok(None),
}
}
Ok(completed_total(node))
}
pub fn statistic_sum(&self, operation_id: &str, name: &str) -> Result<Option<i64>> {
let statistics = self.custom_statistics(operation_id)?;
Ok(jobs::field(&statistics, name).and_then(completed_total))
}
pub fn wait_for_operation(&self, id: &str) -> Result<()> {
let started = Instant::now();
let mut last_reported = String::new();
loop {
let OperationStatus { state, suspended } = self.operation_status(id)?;
let reported = if suspended {
format!("{state}, suspended")
} else {
state.clone()
};
if reported != last_reported {
eprintln!(
"operation {id}: {reported} ({:.0}s)",
started.elapsed().as_secs_f64()
);
last_reported = reported;
}
match state.as_str() {
"completed" => return Ok(()),
"failed" | "aborted" => {
let quick = self.without_retries();
return Err(ClientError::OperationFailed {
id: id.to_owned(),
state,
error: quick.operation_error(id),
jobs: quick.failed_jobs(id),
});
}
_ => std::thread::sleep(self.poll_interval),
}
}
}
pub fn operation_result_error(&self, id: &str) -> Result<Option<String>> {
Ok(operation::result_error_of(
&self.get_operation(id, &["result"])?,
))
}
fn operation_error(&self, id: &str) -> Option<String> {
let body = self
.get_operation_body(
yson_build::map([("operation_id", yson_build::string(id))]),
&["result"],
)
.ok()?;
let summary = self
.strip_envelope(&body, "get_operation")
.ok()
.and_then(|document| {
jobs::field(&document, "result")
.and_then(|result| jobs::error_summary(jobs::field(result, "error")?))
});
summary.or_else(|| Some(crate::error::truncate(&String::from_utf8_lossy(&body), 600)))
}
pub fn list_jobs(
&self,
operation_id: &str,
state: Option<&str>,
limit: u32,
) -> Result<Vec<JobInfo>> {
let mut params = yson_build::map([
("operation_id", yson_build::string(operation_id)),
("limit", yson_build::int(i64::from(limit))),
]);
if let Some(state) = state {
yson_build::insert(&mut params, "state", yson_build::string(state));
}
let body = self.transport.call(
Method::Get,
"list_jobs",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
let envelope = self.strip_envelope(&body, "list_jobs")?;
Ok(jobs::parse_jobs(&self.field_of(&envelope, "jobs")?))
}
pub fn get_job(&self, operation_id: &str, job_id: &str) -> Result<JobInfo> {
let params = yson_build::map([
("operation_id", yson_build::string(operation_id)),
("job_id", yson_build::string(job_id)),
]);
let body = self.transport.call(
Method::Get,
"get_job",
¶ms,
Payload::None,
Repeatable::Freely,
)?;
let document = self.strip_envelope(&body, "get_job")?;
jobs::parse_job(&document).ok_or_else(|| ClientError::Decode {
command: "get_job".to_owned(),
reason: "the answer names no job".to_owned(),
})
}
pub fn get_job_input(&self, operation_id: &str, job_id: &str) -> Result<ResponseReader> {
let params = yson_build::map([
("operation_id", yson_build::string(operation_id)),
("job_id", yson_build::string(job_id)),
]);
let body = self.transport.open(Method::Get, "get_job_input", ¶ms)?;
Ok(ResponseReader::new(body))
}
pub fn get_job_stderr(&self, operation_id: &str, job_id: &str) -> Result<Vec<u8>> {
let params = yson_build::map([
("operation_id", yson_build::string(operation_id)),
("job_id", yson_build::string(job_id)),
]);
self.transport.call(
Method::Get,
"get_job_stderr",
¶ms,
Payload::None,
Repeatable::Heavy,
)
}
fn failed_jobs(&self, operation_id: &str) -> Vec<JobFailure> {
if !self.job_diagnostics {
return Vec::new();
}
self.list_jobs(operation_id, Some("failed"), REPORTED_JOBS)
.unwrap_or_default()
.iter()
.take(REPORTED_JOBS as usize)
.map(|job| JobFailure {
id: job.id.clone(),
address: job.address.clone(),
error: job.error.clone(),
stderr: self.stderr_excerpt(operation_id, job),
})
.collect()
}
fn stderr_excerpt(&self, operation_id: &str, job: &JobInfo) -> Option<String> {
let raw = self.get_job_stderr(operation_id, &job.id).ok()?;
if raw.is_empty() {
return None;
}
Some(crate::error::tail(
&String::from_utf8_lossy(&raw),
STDERR_EXCERPT,
))
}
pub fn raw_command(
&self,
method: Method,
command: &str,
params: &YsonValue,
payload: Option<&[u8]>,
) -> Result<Vec<u8>> {
self.raw_command_with(method, command, params, payload, Repeatable::Never, None)
}
pub fn raw_command_with(
&self,
method: Method,
command: &str,
params: &YsonValue,
payload: Option<&[u8]>,
repeatable: Repeatable,
mutation_id: Option<&MutationId>,
) -> Result<Vec<u8>> {
check_command_name(command)?;
refuse_non_dict_parameters(command, params)?;
refuse_body_on_get(method, command, payload.is_some())?;
let payload = match payload {
Some(bytes) => Payload::Bytes(bytes),
None => Payload::None,
};
self.transport
.call_with(method, command, params, payload, repeatable, mutation_id)
}
pub fn raw_command_streaming(
&self,
method: Method,
command: &str,
params: &YsonValue,
) -> Result<ResponseReader> {
check_command_name(command)?;
refuse_non_dict_parameters(command, params)?;
let body = self.transport.open(method, command, params)?;
Ok(ResponseReader::new(body))
}
pub fn raw_command_upload(
&self,
method: Method,
command: &str,
params: &YsonValue,
mut body: impl std::io::Read,
) -> Result<Vec<u8>> {
check_command_name(command)?;
refuse_non_dict_parameters(command, params)?;
refuse_body_on_get(method, command, true)?;
self.transport.upload(method, command, params, &mut body)
}
fn without_retries(&self) -> Self {
self.clone().with_retries(RetryPolicy::none())
}
fn strip_envelope(&self, body: &[u8], command: &str) -> Result<YsonValue> {
from_slice(body, YsonFormat::Text).map_err(|e| ClientError::Decode {
command: command.to_owned(),
reason: format!(
"{e}; body was {}",
crate::error::truncate(&String::from_utf8_lossy(body), 200)
),
})
}
fn field_of(&self, value: &YsonValue, key: &str) -> Result<YsonValue> {
match &value.node {
YsonNode::Map(m) => m
.get(key.as_bytes())
.cloned()
.ok_or_else(|| ClientError::Decode {
command: key.to_owned(),
reason: format!(
"response has no {key:?}; keys were {:?}",
m.keys()
.map(|k| String::from_utf8_lossy(k).into_owned())
.collect::<Vec<_>>()
),
}),
other => Err(ClientError::Decode {
command: key.to_owned(),
reason: format!("expected a dict, got {other:?}"),
}),
}
}
fn value_field(&self, body: &[u8], key: &str) -> Result<YsonValue> {
let envelope = self.strip_envelope(body, key)?;
self.field_of(&envelope, key)
}
}
fn refused_or_reported(error: ClientError) -> Result<Cached> {
if denied(&error, "create") {
return Ok(Cached::Refused(error));
}
Err(error)
}
fn denied(error: &ClientError, command: &str) -> bool {
matches!(
error,
ClientError::Cluster {
command: failed,
code,
raw,
..
} if failed == command
&& (*code == ACCESS_DENIED || retry::raw_contains_code(raw, &[ACCESS_DENIED]))
)
}
fn check_command_name(command: &str) -> Result<()> {
if command.is_empty() {
return Err(ClientError::Config(
"a raw command needs a command name, e.g. \"get_supported_features\"".to_owned(),
));
}
if let Some(bad) = command
.chars()
.find(|c| !c.is_ascii_alphanumeric() && *c != '_')
{
return Err(ClientError::Config(format!(
"{command:?} is not a command name: it contains {bad:?}, and the name \
goes into the request path as it is. A command is a bare name like \
\"get_supported_features\" — the path it acts on is a parameter."
)));
}
Ok(())
}
fn refuse_non_dict_parameters(command: &str, params: &YsonValue) -> Result<()> {
if !matches!(params.node, YsonNode::Map(_)) {
return Err(ClientError::Config(format!(
"{command}: command parameters are a YSON dict, and this is a \
{:?}. A command that takes no parameters sends `yson_build::empty_map()`.",
params.node
)));
}
Ok(())
}
fn refuse_body_on_get(method: Method, command: &str, has_body: bool) -> Result<()> {
if has_body && matches!(method, Method::Get) {
return Err(ClientError::Config(format!(
"{command}: a GET carries no request body, so the payload would be \
dropped in silence. A command with an input data stream is a PUT."
)));
}
Ok(())
}
fn environment_value(name: &str) -> Option<String> {
std::env::var(name).ok()
}
fn expanded_proxy(proxy: &str, suffix: Option<&str>) -> String {
let proxy = proxy.trim();
let Some(suffix) = suffix else {
return proxy.to_owned();
};
if proxy.contains(':') || proxy.contains('.') || proxy.contains("localhost") {
return proxy.to_owned();
}
format!("{proxy}.{}", suffix.trim_matches('.'))
}
fn split_domains(value: &str) -> Vec<String> {
value
.split([',', ' ', '\t', '\n'])
.map(str::trim)
.filter(|domain| !domain.is_empty())
.map(str::to_owned)
.collect()
}
fn truthy(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes"
)
}
fn token_from_environment() -> Option<String> {
if let Some(token) = std::env::var("YT_TOKEN").ok().and_then(clean_token) {
return Some(token);
}
if let Ok(path) = std::env::var("YT_TOKEN_PATH")
&& let Some(token) = read_token_file(std::path::Path::new(&path))
{
return Some(token);
}
let home = std::env::var("HOME")
.or_else(|_| std::env::var("USERPROFILE"))
.ok()?;
read_token_file(&std::path::Path::new(&home).join(".yt").join("token"))
}
fn read_token_file(path: &std::path::Path) -> Option<String> {
std::fs::read_to_string(path).ok().and_then(clean_token)
}
fn clean_token(raw: String) -> Option<String> {
let trimmed = raw.trim();
(!trimmed.is_empty()).then(|| trimmed.to_owned())
}
fn decode_rows<T: serde::de::DeserializeOwned>(bytes: &[u8], path: &str) -> Result<Vec<T>> {
let mut rows = Vec::new();
let mut stream = ytsaurus_yson::StreamDeserializer::<T>::new(bytes, true);
loop {
match stream.next_item() {
Ok(Some(row)) => rows.push(row),
Ok(None) => return Ok(rows),
Err(e) => {
return Err(ClientError::Decode {
command: "read_table".to_owned(),
reason: format!("{path}: row {}: {e}", rows.len()),
});
}
}
}
}
fn child_names(value: &YsonValue, path: &str) -> Result<Vec<String>> {
if matches!(
value.attr("incomplete").map(|v| &v.node),
Some(YsonNode::Boolean(true))
) {
return Err(ClientError::Decode {
command: "list".to_owned(),
reason: format!(
"{path} has more children than the cluster would list at once, so the \
answer it gave is not all of them"
),
});
}
let YsonNode::List(items) = &value.node else {
return Err(ClientError::Decode {
command: "list".to_owned(),
reason: format!("{path}: the answer is not a list: {:?}", value.node),
});
};
items
.iter()
.map(|item| match &item.node {
YsonNode::String(bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
other => Err(ClientError::Decode {
command: "list".to_owned(),
reason: format!("{path}: a child name is not a string: {other:?}"),
}),
})
.collect()
}
fn completed_total(statistic: &YsonValue) -> Option<i64> {
let by_state = jobs::field(statistic, "$").or_else(|| jobs::field(statistic, "$$"));
let Some(by_state) = by_state else {
return jobs::field(statistic, "sum").and_then(YsonValue::as_i64);
};
let completed = jobs::field(by_state, "completed")?;
let YsonNode::Map(by_type) = &completed.node else {
return None;
};
let mut total: Option<i64> = None;
for per_type in by_type.values() {
if let Some(sum) = jobs::field(per_type, "sum").and_then(YsonValue::as_i64) {
total = Some(total.unwrap_or(0) + sum);
}
}
total
}
#[cfg(test)]
fn check_complete_fragment(data: &[u8]) -> std::result::Result<(), String> {
check_complete_yson_fragment(data, YsonFormat::Binary)
}
fn check_complete_yson_fragment(
mut data: &[u8],
format: YsonFormat,
) -> std::result::Result<(), String> {
use ytsaurus_yson::{Scan, scan_value};
let total = data.len();
loop {
while data.first() == Some(&b';') || data.first().is_some_and(u8::is_ascii_whitespace) {
data = &data[1..];
}
if data.is_empty() {
return Ok(());
}
match scan_value(data, format) {
Ok(Scan::Complete { len }) => data = &data[len..],
Ok(Scan::Incomplete) => {
return Err(format!(
"the response ends inside a record — {} of {total} bytes consumed; \
the stream was cut short",
total - data.len()
));
}
Err(e) => {
return Err(format!(
"the response is not valid {format:?} YSON at byte {}: {e}",
total - data.len()
));
}
}
}
}
fn unsupported_data_format() -> ClientError {
ClientError::Config(
"this ytsaurus-client version does not support the selected data format".to_owned(),
)
}
fn refuse_skiff_table_mismatch(mismatch: Option<String>) -> Result<()> {
match mismatch {
Some(reason) => Err(ClientError::Config(reason)),
None => Ok(()),
}
}
fn refuse_selection_on_write(path: &TablePath) -> Result<()> {
match path.write_refusal() {
Some(reason) => Err(ClientError::Config(reason)),
None => Ok(()),
}
}
fn refuse_mixed_selection_on_read(path: &TablePath) -> Result<()> {
match path.read_refusal() {
Some(reason) => Err(ClientError::Config(reason)),
None => Ok(()),
}
}
fn skiff_table_path(path: &TablePath, format: &SkiffFormat) -> Result<YsonValue> {
if path.selected_columns().is_some() {
return Err(ClientError::Config(format!(
"{}: a Skiff table read's columns are its format's fields, so \
TablePath::columns cannot also apply — put the projection in the \
Skiff schema, or read YSON",
path.as_str()
)));
}
if let Some(reason) = path.selection_conflict(
true,
false,
"the Skiff format's fields become",
"the Skiff read adds",
) {
return Err(ClientError::Config(reason));
}
if format.table_schemas().len() != 1 {
return Err(ClientError::Config(format!(
"Skiff table I/O requires exactly one table schema, got {}",
format.table_schemas().len()
)));
}
let schema = format.table_schema(0).map_err(|error| {
ClientError::Config(format!(
"Skiff table I/O has an invalid table schema: {error}"
))
})?;
let columns = schema
.children
.iter()
.map(|column| {
let name = column.name.as_deref().ok_or_else(|| {
ClientError::Config("Skiff table I/O schema has an unnamed column".to_owned())
})?;
if matches!(name, "$key_switch" | "$row_index" | "$range_index") {
return Err(ClientError::Config(format!(
"Skiff table I/O schema contains job-only system column {name}"
)));
}
Ok(yson_build::string(name))
})
.collect::<Result<Vec<_>>>()?;
let mut value = path.to_yson();
value
.attributes
.get_or_insert_with(std::collections::BTreeMap::new)
.insert(b"columns".to_vec(), yson_build::list(columns));
Ok(value)
}
fn check_complete_skiff_stream(
data: &[u8],
format: &SkiffFormat,
) -> std::result::Result<(), String> {
let mut decoder = SkiffDecoder::new(data, format.clone());
while decoder
.skip_row()
.map_err(|error| format!("not a complete Skiff stream: {error}"))?
.is_some()
{}
Ok(())
}
#[cfg(test)]
mod tests {
use std::{
io::{Read, Write},
net::{TcpListener, TcpStream},
thread,
time::Duration,
};
use super::*;
use ytsaurus_skiff::{Encoder as SkiffEncoder, Schema, SchemaRef, Value, WireType};
const GET_OPERATION: &str = include_str!("../tests/fixtures/get_operation.yson");
#[test]
fn the_narrow_readers_agree_with_a_document_a_cluster_sent() {
let document = from_slice(GET_OPERATION.as_bytes(), YsonFormat::Text).expect("valid YSON");
assert_eq!(
operation::state_of(&document).expect("the capture carries a state"),
"completed"
);
assert!(
!operation::suspended_of(&document).expect("and a boolean beside it"),
"suspension is read from its own attribute, not from the state"
);
assert_eq!(
operation::result_error_of(&document),
None,
"a completed operation's code-0 error document is not a failure"
);
}
#[test]
fn job_statistics_are_read_from_under_progress() {
let document = from_slice(
br#"{"progress"={"job_statistics"={"time"={"exec"={"$$"={"completed"={"map"={"sum"=744}}}}}}}}"#,
YsonFormat::Text,
)
.expect("valid YSON");
let statistics = operation::statistics_of(&document);
assert!(
jobs::field(&statistics, "time").is_some(),
"the subtree, not the progress node that holds it: {statistics:?}"
);
let empty = from_slice(br#"{"progress"={}}"#, YsonFormat::Text).expect("valid YSON");
assert!(matches!(
operation::statistics_of(&empty).node,
YsonNode::Map(ref m) if m.is_empty()
));
}
#[test]
fn raw_parameters_that_are_not_a_dict_are_refused() {
let client = Client::new("http://localhost:8000").with_retries(RetryPolicy::none());
let not_a_dict = yson_build::list([yson_build::string("get_supported_features")]);
let refused = client.raw_command(Method::Get, "get_supported_features", ¬_a_dict, None);
assert!(
matches!(refused, Err(ClientError::Config(_))),
"a list of parameters is a mistake to report, not to panic on"
);
assert!(refuse_non_dict_parameters("c", &yson_build::empty_map()).is_ok());
}
#[test]
fn an_attached_id_is_trimmed() {
let client = Client::new("http://localhost:8000");
assert_eq!(client.attach_operation("1-2-3-4\n").id(), "1-2-3-4");
assert_eq!(client.attach_operation(" 1-2-3-4 ").id(), "1-2-3-4");
assert_eq!(client.attach_operation("1-2-3-4").id(), "1-2-3-4");
}
#[test]
fn a_get_answer_decodes_straight_into_the_type_asked_for() {
#[derive(serde::Deserialize)]
struct Node {
account: String,
#[serde(rename = "type")]
node_type: String,
}
let body = br#"{"value"={"account"="tmp";"type"="table";"chunk_count"=3}}"#;
let envelope: Envelope<Node> = from_slice(body, YsonFormat::Text).expect("decodes");
assert_eq!(envelope.value.account, "tmp");
assert_eq!(envelope.value.node_type, "table");
}
#[test]
fn an_answer_that_does_not_fit_the_type_is_an_error_rather_than_a_default() {
#[derive(serde::Deserialize)]
struct Node {
#[allow(dead_code)]
account: String,
}
let body = br#"{"value"={"type"="table"}}"#;
assert!(from_slice::<Envelope<Node>>(body, YsonFormat::Text).is_err());
}
#[test]
fn a_complete_fragment_is_accepted() {
let one = b"{\x01\x02a=\x02\x02}";
let mut two = one.to_vec();
two.push(b';');
two.extend_from_slice(one);
assert!(check_complete_fragment(b"").is_ok());
assert!(check_complete_fragment(one).is_ok());
assert!(check_complete_fragment(&two).is_ok());
}
#[test]
fn a_truncated_fragment_is_rejected() {
let full = b"{\x01\x02a=\x02\x02}";
for cut in 1..full.len() {
let err = check_complete_fragment(&full[..cut])
.expect_err("a cut record must not pass as complete");
assert!(
err.contains("cut short") || err.contains("not valid"),
"{err}"
);
}
}
fn skiff_format() -> SkiffFormat {
SkiffFormat::new(vec![SchemaRef::Inline(Schema::tuple([
Schema::named("found", WireType::Uint64),
Schema::named("rcl", WireType::String32),
]))])
.expect("a named tuple is a direct-table format")
}
#[test]
fn skiff_table_path_selects_schema_columns() {
let value = skiff_table_path(&TablePath::from("//tmp/table"), &skiff_format()).unwrap();
let rendered = ytsaurus_yson::to_string(&value, YsonFormat::Text).unwrap();
assert_eq!(rendered, r#"<columns=[found;rcl]>"//tmp/table""#);
}
#[test]
fn a_skiff_path_refuses_a_column_selection_spelled_into_its_string() {
let refused = skiff_table_path(&TablePath::from("//tmp/table{found}"), &skiff_format());
assert!(
matches!(&refused, Err(ClientError::Config(reason)) if reason.contains("already selects columns")),
"a string column selection was not refused: {refused:?}"
);
for path in [
"<columns=[found]>//tmp/table",
"<primary_medium=default>//tmp/table",
] {
let refused = skiff_table_path(&TablePath::from(path), &skiff_format());
assert!(
matches!(&refused, Err(ClientError::Config(reason)) if reason.contains("cannot tell whether")),
"{path} was not refused: {refused:?}"
);
}
let ranged = skiff_table_path(&TablePath::from("//tmp/table[#0:#2]"), &skiff_format())
.expect("a string row range is not a column selection");
assert_eq!(
ytsaurus_yson::to_string(&ranged, YsonFormat::Text).unwrap(),
r#"<columns=[found;rcl]>"//tmp/table[#0:#2]""#
);
assert!(
skiff_table_path(&TablePath::from("//tmp/table").range(0..2), &skiff_format()).is_ok()
);
assert!(skiff_table_path(&TablePath::from(r"//tmp/t\[x\]"), &skiff_format()).is_ok());
assert!(skiff_table_path(&TablePath::from(r"//tmp/t\{x\}"), &skiff_format()).is_ok());
}
#[test]
fn skiff_stream_completeness_uses_the_declared_schema() {
let schema = skiff_format().table_schema(0).unwrap().clone();
let mut encoder = SkiffEncoder::new(Vec::new(), schema).unwrap();
encoder
.write(&Value::Tuple(vec![
Value::Uint64(7),
Value::Bytes(b"ok".to_vec()),
]))
.unwrap();
let complete = encoder.into_inner().unwrap();
assert!(check_complete_skiff_stream(&complete, &skiff_format()).is_ok());
for cut in 1..complete.len() {
assert!(
check_complete_skiff_stream(&complete[..cut], &skiff_format()).is_err(),
"cut at {cut} must not pass"
);
}
}
#[test]
fn direct_skiff_table_format_rejects_multi_table_and_job_controls() {
let multiple = SkiffFormat::new(vec![
SchemaRef::Inline(Schema::tuple([Schema::named("a", WireType::Uint64)])),
SchemaRef::Inline(Schema::tuple([Schema::named("b", WireType::Uint64)])),
])
.unwrap();
assert!(matches!(
skiff_table_path(&TablePath::from("//tmp/table"), &multiple),
Err(ClientError::Config(_))
));
let job_control =
SkiffFormat::new(vec![SchemaRef::Inline(Schema::tuple([Schema::named(
"$key_switch",
WireType::Boolean,
)]))])
.unwrap();
assert!(matches!(
skiff_table_path(&TablePath::from("//tmp/table"), &job_control),
Err(ClientError::Config(_))
));
}
#[test]
fn skiff_table_calls_use_schema_format_columns_and_raw_streams() {
let schema = skiff_format().table_schema(0).unwrap().clone();
let mut encoder = SkiffEncoder::new(Vec::new(), schema).unwrap();
encoder
.write(&Value::Tuple(vec![
Value::Uint64(7),
Value::Bytes(b"ok".to_vec()),
]))
.unwrap();
let stream = encoder.into_inner().unwrap();
let (proxy, write_request) = one_request_proxy(Vec::new());
Client::new(&proxy)
.write_table_with_format("//tmp/write", &stream, &DataFormat::skiff(skiff_format()))
.unwrap();
let write_request = write_request.join().unwrap();
assert!(write_request.starts_with(b"PUT /api/v4/write_table HTTP/1.1\r\n"));
let write_headers = String::from_utf8_lossy(&write_request);
assert!(
write_headers.contains("input_format=<table_skiff_schemas="),
"{write_headers}"
);
assert!(
write_headers.contains(r#"path=<columns=[found;rcl]>"//tmp/write""#),
"{write_headers}"
);
assert!(write_request.ends_with(&stream));
let (proxy, read_request) = one_request_proxy(stream.clone());
let received = Client::new(&proxy)
.read_table_with_format("//tmp/read", &DataFormat::skiff(skiff_format()))
.unwrap();
let read_request = read_request.join().unwrap();
assert!(read_request.starts_with(b"GET /api/v4/read_table HTTP/1.1\r\n"));
let read_headers = String::from_utf8_lossy(&read_request);
assert!(
read_headers.contains("output_format=<table_skiff_schemas="),
"{read_headers}"
);
assert!(
read_headers.contains(r#"path=<columns=[found;rcl]>"//tmp/read""#),
"{read_headers}"
);
assert_eq!(received, stream);
}
#[test]
fn shared_yson_table_format_uses_the_requested_yson_encoding() {
let (proxy, request) = one_request_proxy(Vec::new());
Client::new(&proxy)
.write_table_with_format("//tmp/write", b"{value=one};", &DataFormat::text_yson())
.unwrap();
let request = request.join().unwrap();
let request = String::from_utf8_lossy(&request);
assert!(
request.contains("input_format=<format=text>yson"),
"{request}"
);
}
#[test]
fn a_raw_command_goes_where_it_says_with_the_parameters_it_was_given() {
let (proxy, request) = one_request_proxy(br#"{"value"={};}"#.to_vec());
let body = Client::new(&proxy)
.raw_command(
Method::Get,
"get_supported_features",
&yson_build::empty_map(),
None,
)
.expect("sends");
let request = request.join().unwrap();
assert!(
request.starts_with(b"GET /api/v4/get_supported_features HTTP/1.1\r\n"),
"{}",
String::from_utf8_lossy(&request)
);
let headers = String::from_utf8_lossy(&request);
assert!(headers.contains("x-yt-parameters: {}"), "{headers}");
assert_eq!(body, br#"{"value"={};}"#);
}
#[test]
fn a_raw_command_carries_its_payload_and_its_transaction() {
let (proxy, request) = one_request_proxy(Vec::new());
Client::new(&proxy)
.with_transaction("3-5d231-10001-db88")
.raw_command(
Method::Put,
"write_file",
&yson_build::map([("path", yson_build::string("//tmp/f"))]),
Some(b"payload"),
)
.expect("sends");
let request = request.join().unwrap();
let headers = String::from_utf8_lossy(&request);
assert!(
request.starts_with(b"PUT /api/v4/write_file HTTP/1.1\r\n"),
"{headers}"
);
assert!(request.ends_with(b"payload"), "{headers}");
assert!(
headers.contains(r#"transaction_id="3-5d231-10001-db88""#),
"{headers}"
);
}
#[test]
fn a_raw_command_is_sent_once_unless_the_caller_says_otherwise() {
let (proxy, request) = one_request_proxy(Vec::new());
let client = Client::new(&proxy).with_retries(RetryPolicy::none());
client
.raw_command(Method::Post, "concatenate", &yson_build::empty_map(), None)
.expect("sends");
request.join().unwrap();
}
#[test]
fn a_mutation_id_is_sent_even_when_the_command_is_not_retried() {
let id = MutationId::new().as_retry();
let (proxy, request) = one_request_proxy(Vec::new());
Client::new(&proxy)
.raw_command_with(
Method::Post,
"concatenate",
&yson_build::empty_map(),
None,
Repeatable::Never,
Some(&id),
)
.expect("sends");
let request = request.join().unwrap();
let sent = sent_parameters(&request);
assert_eq!(
parameter(&sent, "mutation_id").and_then(YsonValue::as_str),
Some(id.as_str()),
"{}",
String::from_utf8_lossy(&request)
);
assert_eq!(
parameter(&sent, "retry").map(|v| &v.node),
Some(&YsonNode::Boolean(true)),
"{}",
String::from_utf8_lossy(&request)
);
}
fn sent_parameters(request: &[u8]) -> YsonValue {
let head = String::from_utf8_lossy(request);
let line = head
.lines()
.find(|line| {
line.split_once(':')
.is_some_and(|(name, _)| name.eq_ignore_ascii_case("x-yt-parameters"))
})
.unwrap_or_else(|| panic!("no X-YT-Parameters header in:\n{head}"));
let value = line
.split_once(':')
.expect("the header has a value")
.1
.trim();
from_slice(value.as_bytes(), YsonFormat::Text)
.unwrap_or_else(|e| panic!("parameters are not text YSON ({e}): {value}"))
}
fn parameter<'a>(params: &'a YsonValue, key: &str) -> Option<&'a YsonValue> {
match ¶ms.node {
YsonNode::Map(m) => m.get(key.as_bytes()),
_ => None,
}
}
#[test]
fn a_command_name_that_would_change_the_url_is_refused() {
let client = Client::new("http://localhost:8000");
for bad in [
"",
"get/../../hosts",
"get?x=1",
"get#frag",
"get value",
"get%2f",
] {
let error = client
.raw_command(Method::Get, bad, &yson_build::empty_map(), None)
.expect_err(&format!("{bad:?} was accepted as a command name"));
assert!(matches!(error, ClientError::Config(_)), "{bad:?}: {error}");
}
assert!(check_command_name("get_supported_features").is_ok());
assert!(check_command_name("start_tx").is_ok());
assert!(check_command_name("read_table_partition2").is_ok());
}
#[test]
fn a_payload_on_a_get_is_refused_rather_than_dropped() {
let error = Client::new("http://localhost:8000")
.raw_command(
Method::Get,
"read_table",
&yson_build::empty_map(),
Some(b"x"),
)
.expect_err("a GET with a body is a mistake");
assert!(matches!(error, ClientError::Config(_)), "{error}");
assert!(refuse_body_on_get(Method::Get, "get", false).is_ok());
assert!(refuse_body_on_get(Method::Put, "write_file", true).is_ok());
assert!(refuse_body_on_get(Method::Post, "create", true).is_ok());
}
#[test]
fn read_file_refuses_a_body_it_will_not_hold() {
let (proxy, served) = one_gzip_request_proxy(vec![0_u8; 40_000]);
let mut client = Client::new(&proxy);
client.transport.set_response_limit(4_096);
let error = client
.read_file("//tmp/f")
.expect_err("40 000 bytes past a 4 096-byte ceiling");
assert!(
matches!(error, ClientError::ResponseTooLarge { limit: 4_096, .. }),
"{error:?}"
);
let message = error.to_string();
assert!(message.contains("read_file"), "{message}");
assert!(message.contains("4096"), "{message}");
assert!(message.contains("read_file_streaming"), "{message}");
let request = served.join().unwrap();
assert!(
request.starts_with(b"GET /api/v4/read_file HTTP/1.1\r\n"),
"{}",
String::from_utf8_lossy(&request)
);
}
fn one_gzip_request_proxy(payload: Vec<u8>) -> (String, thread::JoinHandle<Vec<u8>>) {
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
encoder.write_all(&payload).unwrap();
let body = encoder.finish().unwrap();
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let address = listener.local_addr().unwrap();
let task = thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(5)))
.unwrap();
let request = read_http_request(&mut stream);
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Encoding: gzip\r\nContent-Length: {}\r\n\
Connection: close\r\n\r\n",
body.len()
);
stream.write_all(response.as_bytes()).unwrap();
stream.write_all(&body).unwrap();
request
});
(format!("http://{address}"), task)
}
fn one_request_proxy(body: Vec<u8>) -> (String, thread::JoinHandle<Vec<u8>>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let address = listener.local_addr().unwrap();
let task = thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(5)))
.unwrap();
let request = read_http_request(&mut stream);
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
body.len()
);
stream.write_all(response.as_bytes()).unwrap();
stream.write_all(&body).unwrap();
request
});
(format!("http://{address}"), task)
}
fn read_http_request(stream: &mut TcpStream) -> Vec<u8> {
let mut request = Vec::new();
let mut buffer = [0; 1024];
let expected = loop {
let read = stream.read(&mut buffer).unwrap();
assert!(read != 0, "client closed before sending a complete request");
request.extend_from_slice(&buffer[..read]);
let Some(headers_end) = request.windows(4).position(|window| window == b"\r\n\r\n")
else {
continue;
};
let headers = String::from_utf8_lossy(&request[..headers_end + 4]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then_some(value.trim())
})
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(0);
break headers_end + 4 + content_length;
};
while request.len() < expected {
let read = stream.read(&mut buffer).unwrap();
assert!(read != 0, "client closed before sending its request body");
request.extend_from_slice(&buffer[..read]);
}
request
}
#[test]
fn truncation_after_a_whole_record_is_rejected() {
let one = b"{\x01\x02a=\x02\x02}";
let mut data = one.to_vec();
data.push(b';');
data.extend_from_slice(&one[..4]);
let err = check_complete_fragment(&data).expect_err("must reject");
assert!(err.contains("cut short"), "{err}");
}
#[test]
fn a_token_file_written_with_echo_still_works() {
let path = std::env::temp_dir().join(format!(
"ytsaurus-rs-token-{}-{:?}",
std::process::id(),
std::thread::current().id()
));
std::fs::write(&path, " secret-token\n").expect("writes");
assert_eq!(read_token_file(&path).as_deref(), Some("secret-token"));
std::fs::write(&path, "\n \n").expect("writes");
assert_eq!(read_token_file(&path), None, "whitespace is not a token");
std::fs::remove_file(&path).ok();
assert_eq!(
read_token_file(&path),
None,
"a missing file is no token, not an error"
);
}
#[test]
fn a_listing_is_the_names_in_the_order_given() {
let value = from_slice(br#"["t1";"t2";]"#, YsonFormat::Text).expect("valid YSON");
assert_eq!(child_names(&value, "//tmp/x").unwrap(), ["t1", "t2"]);
}
#[test]
fn a_truncated_listing_is_an_error_rather_than_a_short_list() {
let value =
from_slice(br#"<"incomplete"=%true;>["t1";]"#, YsonFormat::Text).expect("valid YSON");
let err = child_names(&value, "//tmp/x").expect_err("must not pass as a listing");
assert!(err.to_string().contains("not all of them"), "{err}");
}
const EXISTS_RESPONSE: &[u8] = br#"{"value"=%false;}"#;
#[test]
fn an_exists_answer_is_read_out_of_the_value_key() {
let client = Client::new("http://localhost:8000");
let value = client
.value_field(EXISTS_RESPONSE, "value")
.expect("the answer is an envelope around `value`");
assert!(matches!(value.node, YsonNode::Boolean(false)));
assert!(client.value_field(EXISTS_RESPONSE, "exists").is_err());
}
const CUSTOM_STATISTICS: &str = r#"{
"bytes/read" = {"$" = {completed = {map = {count=1;max=147;min=147;sum=147}}}};
"rows/read" = {"$" = {completed = {map = {count=1;max=7;min=7;sum=7}}}};
"rows/rejected" = {"$" = {completed = {map = {count=1;max=3;min=3;sum=3}}}};
}"#;
fn statistics() -> YsonValue {
from_slice(CUSTOM_STATISTICS.as_bytes(), YsonFormat::Text).expect("valid YSON")
}
#[test]
fn a_statistic_totals_over_completed_jobs() {
let all = statistics();
assert_eq!(
jobs::field(&all, "rows/rejected").and_then(completed_total),
Some(3)
);
assert_eq!(
jobs::field(&all, "bytes/read").and_then(completed_total),
Some(147)
);
assert_eq!(jobs::field(&all, "rows").and_then(completed_total), None);
}
#[test]
fn job_types_are_summed_and_other_states_are_not() {
let value = from_slice(
br#"{"$" = {
completed = {map = {sum=10}; partition_reduce = {sum=5}};
aborted = {map = {sum=99}};
}}"#,
YsonFormat::Text,
)
.expect("valid YSON");
assert_eq!(completed_total(&value), Some(15));
}
#[test]
fn a_flat_aggregate_still_yields_a_number() {
let value =
from_slice(b"{count=1;max=7;min=7;sum=7}", YsonFormat::Text).expect("valid YSON");
assert_eq!(completed_total(&value), Some(7));
}
#[test]
fn an_operation_whose_jobs_all_failed_totals_nothing() {
let value = from_slice(br#"{"$" = {failed = {map = {sum=4}}}}"#, YsonFormat::Text)
.expect("valid YSON");
assert_eq!(completed_total(&value), None);
}
#[test]
fn from_env_explains_itself_when_unconfigured() {
let err = ClientError::Config("YT_PROXY is not set".to_owned());
assert!(err.to_string().contains("YT_PROXY"));
}
fn from_environment(vars: &[(&str, &str)]) -> Result<Client> {
Client::from_lookup(|name| {
vars.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
})
}
#[test]
fn each_variable_reaches_the_setting_it_names() {
let client = from_environment(&[
("YT_PROXY", "hume"),
("YT_PROXY_SUFFIX", ".yt.example.net"),
("YT_HEAVY_PROXY_DOMAINS", "proxy-zone.net, other-zone.net"),
("YT_FILE_CACHE", "//tmp/mine/cache"),
])
.expect("YT_PROXY is set");
assert_eq!(
client.transport.configured_address(),
"https://hume.yt.example.net"
);
assert_eq!(client.file_cache, "//tmp/mine/cache");
assert_eq!(
client.transport.heavy_hosts_debug(),
r#"Under { domains: ["proxy-zone.net", "other-zone.net"], ignored: [] }"#
);
}
#[test]
fn a_machine_that_sets_nothing_gets_the_defaults() {
let bare = from_environment(&[("YT_PROXY", "http://localhost:8000")])
.expect("YT_PROXY is set")
.transport;
let new = Client::new("http://localhost:8000").transport;
assert_eq!(bare.configured_address(), new.configured_address());
assert_eq!(bare.heavy_hosts_debug(), new.heavy_hosts_debug());
assert_eq!(
from_environment(&[("YT_PROXY", "http://localhost:8000")])
.expect("YT_PROXY is set")
.file_cache,
Client::new("http://localhost:8000").file_cache
);
}
#[test]
fn the_wider_heavy_proxy_setting_wins_however_it_was_exported() {
let hosts = from_environment(&[
("YT_PROXY", "https://cluster.example.net"),
("YT_HEAVY_PROXY_DOMAINS", "proxy-zone.net"),
("YT_HEAVY_PROXIES_ANYWHERE", "1"),
])
.expect("YT_PROXY is set")
.transport
.heavy_hosts_debug();
assert!(hosts.contains("Anywhere"), "{hosts}");
let hosts = from_environment(&[
("YT_PROXY", "https://cluster.example.net"),
("YT_HEAVY_PROXY_DOMAINS", "proxy-zone.net"),
("YT_HEAVY_PROXIES_ANYWHERE", "0"),
])
.expect("YT_PROXY is set")
.transport
.heavy_hosts_debug();
assert_eq!(
hosts,
r#"Under { domains: ["proxy-zone.net"], ignored: [] }"#
);
}
#[test]
fn a_variable_set_to_nothing_is_a_variable_that_is_not_set() {
let client = from_environment(&[
("YT_PROXY", " https://cluster.example.net "),
("YT_FILE_CACHE", " "),
("YT_HEAVY_PROXY_DOMAINS", ""),
("YT_PROXY_SUFFIX", ""),
])
.expect("YT_PROXY is set");
assert_eq!(
client.transport.configured_address(),
"https://cluster.example.net",
"and a value that is set is trimmed"
);
assert_eq!(client.file_cache, Client::new("x").file_cache);
assert_eq!(client.transport.heavy_hosts_debug(), "SameDomain");
}
#[test]
fn a_proxy_set_to_nothing_is_a_proxy_that_is_not_set() {
let err = from_environment(&[("YT_PROXY", " "), ("YT_PROXY_SUFFIX", ".yt.example.net")])
.expect_err("an empty proxy is not a proxy");
assert!(err.to_string().contains("YT_PROXY is not set"), "{err}");
}
#[test]
fn a_bare_cluster_name_is_completed_only_when_a_suffix_says_so() {
assert_eq!(
expanded_proxy("hume", Some(".yt.example.net")),
"hume.yt.example.net"
);
for suffix in ["yt.example.net", "yt.example.net.", " .yt.example.net "] {
assert_eq!(
expanded_proxy("hume", Some(suffix.trim())),
"hume.yt.example.net",
"{suffix:?}"
);
}
assert_eq!(expanded_proxy("hume", None), "hume");
}
#[test]
fn a_name_that_needs_no_completing_is_left_alone() {
for proxy in [
"http://localhost:8000",
"localhost",
"hume.yt.example.net",
"cluster.example.net",
"10.0.0.7",
"hume:80",
"mylocalhostcluster",
] {
assert_eq!(expanded_proxy(proxy, Some(".yt.example.net")), proxy);
}
}
#[test]
fn domains_are_read_as_a_list_however_they_were_written() {
assert_eq!(
split_domains("proxy-zone.net, sas.proxy-zone.net"),
["proxy-zone.net", "sas.proxy-zone.net"]
);
assert_eq!(
split_domains("proxy-zone.net sas.proxy-zone.net"),
["proxy-zone.net", "sas.proxy-zone.net"]
);
assert_eq!(split_domains("proxy-zone.net,,"), ["proxy-zone.net"]);
assert!(split_domains(" , ").is_empty());
}
#[test]
fn only_the_three_spellings_of_yes_are_yes() {
for value in ["1", "true", "TRUE", "yes", " Yes "] {
assert!(truthy(value), "{value}");
}
for value in ["0", "false", "no", "on", "enabled", ""] {
assert!(!truthy(value), "{value}");
}
}
}