use std::time::Duration;
use bytes::Bytes;
use crate::connection::Connection;
use crate::error::{Error, Result};
use crate::guid::Guid;
use crate::proto;
use crate::wire::{self, MaybeRow, Row};
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);
pub const DEFAULT_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(15);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(i32)]
pub enum TransactionType {
Master = 0,
Tablet = 1,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(i32)]
pub enum Atomicity {
Full = 0,
None = 1,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(i32)]
pub enum RowModificationType {
Write = 0,
Delete = 1,
WriteAndLock = 3,
}
pub type Timestamp = u64;
pub const LATEST_TIMESTAMP: Timestamp = 0x3fff_ffff_ffff_ff01;
#[derive(Debug)]
pub struct Client {
connection: Connection,
timeout: Duration,
}
impl Client {
pub async fn connect(address: &str) -> Result<Self> {
Self::builder(address).connect().await
}
pub fn builder(address: &str) -> ClientBuilder {
ClientBuilder {
address: address.to_owned(),
token: None,
timeout: DEFAULT_TIMEOUT,
}
}
pub fn connection(&self) -> &Connection {
&self.connection
}
pub async fn discover_proxies(&self, role: Option<&str>) -> Result<Vec<String>> {
crate::connection::discover_proxies(&self.connection, role, Some(self.timeout)).await
}
pub async fn start_transaction(
&self,
transaction_type: TransactionType,
options: StartTransactionOptions,
) -> Result<Transaction<'_>> {
let request = start_transaction_request(transaction_type, &options);
let (response, _) = self
.connection
.invoke::<proto::api::TRspStartTransaction>(
"StartTransaction",
&request,
Vec::new(),
Some(self.timeout),
"TRspStartTransaction",
)
.await?;
Ok(Transaction {
client: self,
id: Guid::from_proto(&response.id),
start_timestamp: response.start_timestamp,
finished: false,
})
}
pub async fn lookup_rows(
&self,
path: &str,
columns: &[&str],
keys: &[Row],
options: LookupOptions<'_>,
) -> Result<Vec<MaybeRow>> {
let keys: Vec<MaybeRow> = keys.iter().cloned().map(Some).collect();
let request = lookup_request(path, columns, &options);
let (response, attachments) = self
.connection
.invoke::<proto::api::TRspLookupRows>(
"LookupRows",
&request,
vec![wire::encode_rowset(&keys)?],
Some(self.timeout),
"TRspLookupRows",
)
.await?;
decode_rowset_attachments(&attachments, Some(&response.rowset_descriptor))
}
pub async fn lookup_rows_with_columns(
&self,
path: &str,
columns: &[&str],
keys: &[Row],
options: LookupOptions<'_>,
) -> Result<(Vec<MaybeRow>, Vec<String>)> {
let keys: Vec<MaybeRow> = keys.iter().cloned().map(Some).collect();
let request = lookup_request(path, columns, &options);
let (response, attachments) = self
.connection
.invoke::<proto::api::TRspLookupRows>(
"LookupRows",
&request,
vec![wire::encode_rowset(&keys)?],
Some(self.timeout),
"TRspLookupRows",
)
.await?;
let descriptor = &response.rowset_descriptor;
let rows = decode_rowset_attachments(&attachments, Some(descriptor))?;
Ok((rows, descriptor_column_names(descriptor)))
}
pub async fn select_rows(&self, query: &str, options: SelectOptions) -> Result<Vec<MaybeRow>> {
Ok(self.select_rows_with_columns(query, options).await?.0)
}
pub async fn select_rows_with_columns(
&self,
query: &str,
options: SelectOptions,
) -> Result<(Vec<MaybeRow>, Vec<String>)> {
let request = select_request(query, &options);
let (response, attachments) = self
.connection
.invoke::<proto::api::TRspSelectRows>(
"SelectRows",
&request,
Vec::new(),
Some(self.timeout),
"TRspSelectRows",
)
.await?;
let descriptor = &response.rowset_descriptor;
let rows = decode_rowset_attachments(&attachments, Some(descriptor))?;
Ok((rows, descriptor_column_names(descriptor)))
}
}
#[derive(Debug, Clone)]
pub struct ClientBuilder {
address: String,
token: Option<String>,
timeout: Duration,
}
impl ClientBuilder {
pub fn token(mut self, token: impl Into<String>) -> Self {
self.token = Some(token.into());
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
pub async fn connect(self) -> Result<Client> {
let connection = Connection::connect(&self.address, self.token).await?;
Ok(Client {
connection,
timeout: self.timeout,
})
}
}
#[derive(Debug, Clone)]
pub struct StartTransactionOptions {
pub timeout: Duration,
pub atomicity: Atomicity,
pub parent_id: Option<Guid>,
}
impl Default for StartTransactionOptions {
fn default() -> Self {
Self {
timeout: DEFAULT_TRANSACTION_TIMEOUT,
atomicity: Atomicity::Full,
parent_id: None,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct LookupOptions<'a> {
pub timestamp: Option<Timestamp>,
pub column_filter: Vec<&'a str>,
}
#[derive(Debug, Clone, Default)]
pub struct SelectOptions {
pub timestamp: Option<Timestamp>,
pub output_row_limit: Option<u64>,
}
#[derive(Debug)]
pub struct Transaction<'a> {
client: &'a Client,
id: Guid,
start_timestamp: Timestamp,
finished: bool,
}
impl Transaction<'_> {
pub fn id(&self) -> Guid {
self.id
}
pub fn start_timestamp(&self) -> Timestamp {
self.start_timestamp
}
pub async fn ping(&self) -> Result<()> {
let request = proto::api::TReqPingTransaction {
transaction_id: self.id.to_proto(),
..Default::default()
};
self.client
.connection
.invoke::<proto::api::TRspPingTransaction>(
"PingTransaction",
&request,
Vec::new(),
Some(self.client.timeout),
"TRspPingTransaction",
)
.await?;
Ok(())
}
pub async fn commit(mut self) -> Result<()> {
let request = proto::api::TReqCommitTransaction {
transaction_id: self.id.to_proto(),
..Default::default()
};
self.client
.connection
.invoke::<proto::api::TRspCommitTransaction>(
"CommitTransaction",
&request,
Vec::new(),
Some(self.client.timeout),
"TRspCommitTransaction",
)
.await?;
self.finished = true;
Ok(())
}
pub async fn abort(mut self) -> Result<()> {
let request = proto::api::TReqAbortTransaction {
transaction_id: self.id.to_proto(),
..Default::default()
};
self.client
.connection
.invoke::<proto::api::TRspAbortTransaction>(
"AbortTransaction",
&request,
Vec::new(),
Some(self.client.timeout),
"TRspAbortTransaction",
)
.await?;
self.finished = true;
Ok(())
}
pub fn is_finished(&self) -> bool {
self.finished
}
pub async fn lookup_rows(
&self,
path: &str,
columns: &[&str],
keys: &[Row],
mut options: LookupOptions<'_>,
) -> Result<Vec<MaybeRow>> {
options.timestamp = Some(self.start_timestamp);
self.client.lookup_rows(path, columns, keys, options).await
}
pub async fn lookup_rows_with_columns(
&self,
path: &str,
columns: &[&str],
keys: &[Row],
mut options: LookupOptions<'_>,
) -> Result<(Vec<MaybeRow>, Vec<String>)> {
options.timestamp = Some(self.start_timestamp);
self.client
.lookup_rows_with_columns(path, columns, keys, options)
.await
}
pub async fn select_rows_with_columns(
&self,
query: &str,
mut options: SelectOptions,
) -> Result<(Vec<MaybeRow>, Vec<String>)> {
options.timestamp = Some(self.start_timestamp);
self.client.select_rows_with_columns(query, options).await
}
pub async fn select_rows(
&self,
query: &str,
mut options: SelectOptions,
) -> Result<Vec<MaybeRow>> {
options.timestamp = Some(self.start_timestamp);
self.client.select_rows(query, options).await
}
pub async fn insert_rows(&self, path: &str, columns: &[&str], rows: &[Row]) -> Result<()> {
self.modify_rows(path, columns, rows, RowModificationType::Write)
.await
}
pub async fn delete_rows(&self, path: &str, columns: &[&str], keys: &[Row]) -> Result<()> {
self.modify_rows(path, columns, keys, RowModificationType::Delete)
.await
}
pub async fn modify_rows(
&self,
path: &str,
columns: &[&str],
rows: &[Row],
modification: RowModificationType,
) -> Result<()> {
let owned: Vec<MaybeRow> = rows.iter().cloned().map(Some).collect();
let request = modify_request(self.id, path, columns, rows.len(), modification);
self.client
.connection
.invoke::<proto::api::TRspModifyRows>(
"ModifyRows",
&request,
vec![wire::encode_rowset(&owned)?],
Some(self.client.timeout),
"TRspModifyRows",
)
.await?;
Ok(())
}
}
fn start_transaction_request(
transaction_type: TransactionType,
options: &StartTransactionOptions,
) -> proto::api::TReqStartTransaction {
proto::api::TReqStartTransaction {
r#type: transaction_type as i32,
timeout: Some(options.timeout.as_micros() as i64),
sticky: Some(transaction_type == TransactionType::Tablet),
atomicity: Some(options.atomicity as i32),
parent_id: options.parent_id.map(Guid::to_proto),
..Default::default()
}
}
fn lookup_request(
path: &str,
columns: &[&str],
options: &LookupOptions<'_>,
) -> proto::api::TReqLookupRows {
proto::api::TReqLookupRows {
path: path.as_bytes().to_vec(),
rowset_descriptor: name_table_descriptor(columns),
timestamp: options.timestamp,
keep_missing_rows: Some(true),
columns: options
.column_filter
.iter()
.map(|column| (*column).to_owned())
.collect(),
..Default::default()
}
}
fn select_request(query: &str, options: &SelectOptions) -> proto::api::TReqSelectRows {
proto::api::TReqSelectRows {
query: query.to_owned(),
timestamp: options.timestamp,
output_row_limit: options.output_row_limit,
..Default::default()
}
}
fn modify_request(
transaction_id: Guid,
path: &str,
columns: &[&str],
row_count: usize,
modification: RowModificationType,
) -> proto::api::TReqModifyRows {
proto::api::TReqModifyRows {
transaction_id: transaction_id.to_proto(),
path: path.as_bytes().to_vec(),
rowset_descriptor: name_table_descriptor(columns),
row_modification_types: vec![modification as i32; row_count],
..Default::default()
}
}
fn descriptor_column_names(descriptor: &proto::api::TRowsetDescriptor) -> Vec<String> {
descriptor
.name_table_entries
.iter()
.map(|entry| entry.name.clone().unwrap_or_default())
.collect()
}
fn name_table_descriptor(columns: &[&str]) -> proto::api::TRowsetDescriptor {
proto::api::TRowsetDescriptor {
wire_format_version: Some(CURRENT_WIRE_FORMAT_VERSION),
rowset_kind: Some(proto::api::ERowsetKind::RkUnversioned as i32),
name_table_entries: columns
.iter()
.map(|name| proto::api::t_rowset_descriptor::TNameTableEntry {
name: Some((*name).to_owned()),
..Default::default()
})
.collect(),
..Default::default()
}
}
const CURRENT_WIRE_FORMAT_VERSION: i32 = 1;
fn decode_rowset_attachments(
attachments: &[Bytes],
descriptor: Option<&proto::api::TRowsetDescriptor>,
) -> Result<Vec<MaybeRow>> {
if let Some(descriptor) = descriptor
&& let Some(version) = descriptor.wire_format_version
&& version != CURRENT_WIRE_FORMAT_VERSION
{
return Err(Error::Protocol(format!(
"the proxy replied with wire format version {version}, and this client speaks {CURRENT_WIRE_FORMAT_VERSION}"
)));
}
let merged = match attachments {
[] => return Ok(Vec::new()),
[single] => single.clone(),
many => {
let mut merged =
bytes::BytesMut::with_capacity(many.iter().map(Bytes::len).sum::<usize>());
for attachment in many {
merged.extend_from_slice(attachment);
}
merged.freeze()
}
};
Ok(wire::decode_rowset(&merged)?)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::wire::{UnversionedValue, Value};
use prost::Message as _;
#[test]
fn enum_values_match_the_proto() {
assert_eq!(
RowModificationType::Write as i32,
proto::api::ERowModificationType::RmtWrite as i32
);
assert_eq!(
RowModificationType::Delete as i32,
proto::api::ERowModificationType::RmtDelete as i32
);
assert_eq!(
RowModificationType::WriteAndLock as i32,
proto::api::ERowModificationType::RmtModify as i32
);
assert_eq!(
TransactionType::Master as i32,
proto::api::ETransactionType::TtMaster as i32
);
assert_eq!(
TransactionType::Tablet as i32,
proto::api::ETransactionType::TtTablet as i32
);
assert_eq!(Atomicity::Full as i32, proto::api::EAtomicity::AFull as i32);
assert_eq!(Atomicity::None as i32, proto::api::EAtomicity::ANone as i32);
}
#[test]
fn enum_values_match_the_documented_numbers() {
assert_eq!(TransactionType::Master as i32, 0);
assert_eq!(TransactionType::Tablet as i32, 1);
assert_eq!(Atomicity::Full as i32, 0);
assert_eq!(Atomicity::None as i32, 1);
assert_eq!(RowModificationType::Write as i32, 0);
assert_eq!(RowModificationType::Delete as i32, 1);
assert_eq!(RowModificationType::WriteAndLock as i32, 3);
}
#[test]
fn the_latest_timestamp_sentinel_is_the_proto_default() {
assert_ne!(LATEST_TIMESTAMP, 0, "zero is NullTimestamp, not 'latest'");
let mut buffer = Vec::new();
proto::api::TReqLookupRows {
path: b"//tmp/t".to_vec(),
timestamp: None,
..Default::default()
}
.encode(&mut buffer)
.unwrap();
let parsed = proto::api::TReqLookupRows::decode(&buffer[..]).unwrap();
assert_eq!(
parsed.timestamp.unwrap_or(LATEST_TIMESTAMP),
LATEST_TIMESTAMP
);
assert_eq!(LATEST_TIMESTAMP, 0x3fff_ffff_ffff_ff01);
}
#[test]
fn lookup_asks_for_what_it_promises() {
let request = lookup_request(
"//tmp/table",
&["key"],
&LookupOptions {
timestamp: None,
column_filter: vec!["key", "value"],
},
);
assert_eq!(request.path, b"//tmp/table".to_vec());
assert_eq!(request.columns, ["key", "value"]);
assert_eq!(
request.keep_missing_rows,
Some(true),
"without this a missing key shortens the answer and misaligns the rest"
);
assert_eq!(
request.timestamp, None,
"omitted means the proto default, which is the latest committed data; sending 0 would ask for NullTimestamp instead"
);
assert_eq!(
request
.rowset_descriptor
.name_table_entries
.iter()
.map(|entry| entry.name.clone().unwrap())
.collect::<Vec<_>>(),
["key"],
"the descriptor names the key columns the attachment carries"
);
}
#[test]
fn a_lookup_in_a_transaction_reads_at_its_start_timestamp() {
let request = lookup_request(
"//tmp/table",
&["key"],
&LookupOptions {
timestamp: Some(1234),
column_filter: Vec::new(),
},
);
assert_eq!(request.timestamp, Some(1234));
assert!(
request.columns.is_empty(),
"an empty filter means every column"
);
}
#[test]
fn select_carries_the_query_timestamp_and_output_limit() {
let request = select_request(
"* from [//tmp/t]",
&SelectOptions {
timestamp: Some(99),
output_row_limit: Some(10),
},
);
assert_eq!(request.query, "* from [//tmp/t]");
assert_eq!(request.timestamp, Some(99));
assert_eq!(request.output_row_limit, Some(10));
}
#[test]
fn modify_names_the_transaction_and_one_type_per_row() {
let transaction = Guid::random();
let request = modify_request(
transaction,
"//tmp/table",
&["key", "value"],
3,
RowModificationType::Delete,
);
assert_eq!(Guid::from_proto(&request.transaction_id), transaction);
assert_eq!(request.path, b"//tmp/table".to_vec());
assert_eq!(
request.row_modification_types,
vec![RowModificationType::Delete as i32; 3],
"one entry per row, parallel to the rows in the attachment"
);
assert!(
request.row_legacy_read_locks.is_empty()
&& request.row_legacy_locks.is_empty()
&& request.row_locks.is_empty(),
"the lock arrays are all-or-nothing per request; a partially filled one breaks the server's one-per-row invariant"
);
}
#[test]
fn only_a_tablet_transaction_is_sticky() {
let options = StartTransactionOptions::default();
let tablet = start_transaction_request(TransactionType::Tablet, &options);
assert_eq!(tablet.r#type, 1);
assert_eq!(
tablet.sticky,
Some(true),
"a tablet tx belongs to one proxy"
);
assert_eq!(
tablet.timeout,
Some(options.timeout.as_micros() as i64),
"microseconds, not milliseconds"
);
let master = start_transaction_request(TransactionType::Master, &options);
assert_eq!(master.r#type, 0);
assert_eq!(master.sticky, Some(false));
}
#[test]
fn the_descriptor_numbers_columns_in_order() {
let descriptor = name_table_descriptor(&["key", "value", "extra"]);
assert_eq!(descriptor.wire_format_version, Some(1));
assert_eq!(
descriptor.rowset_kind,
Some(proto::api::ERowsetKind::RkUnversioned as i32)
);
let names: Vec<_> = descriptor
.name_table_entries
.iter()
.map(|entry| entry.name.clone().unwrap())
.collect();
assert_eq!(names, ["key", "value", "extra"]);
}
#[test]
fn attachments_are_concatenated_before_decoding() {
let rows = vec![
Some(vec![UnversionedValue::new(0, Value::Int64(1))]),
Some(vec![UnversionedValue::new(0, Value::Int64(2))]),
];
let encoded = wire::encode_rowset(&rows).unwrap();
let split = encoded.len() / 2;
let attachments = vec![encoded.slice(0..split), encoded.slice(split..)];
let decoded = decode_rowset_attachments(&attachments, None).unwrap();
assert_eq!(decoded, rows);
}
#[test]
fn no_attachments_means_no_rows() {
assert_eq!(decode_rowset_attachments(&[], None).unwrap(), Vec::new());
}
#[test]
fn an_unknown_wire_format_version_is_refused() {
let descriptor = proto::api::TRowsetDescriptor {
wire_format_version: Some(99),
..Default::default()
};
let error = decode_rowset_attachments(&[], Some(&descriptor)).unwrap_err();
assert!(
error.to_string().contains("wire format version 99"),
"unexpected error: {error}"
);
}
}