use std::collections::HashMap;
use std::fmt::Display;
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use ts_rs::TS;
use crate::assert_table_api;
use crate::client::{Client, Features};
use crate::config::{Expressions, ViewConfigUpdate};
use crate::proto::make_table_req::MakeTableOptions;
use crate::proto::make_table_req::make_table_options::MakeTableType;
use crate::proto::request::ClientReq;
use crate::proto::response::ClientResp;
use crate::proto::*;
use crate::table_data::UpdateData;
use crate::utils::*;
use crate::view::{View, ViewSource};
pub type Schema = HashMap<String, ColumnType>;
#[derive(Clone, Copy, Debug, Serialize, Deserialize, TS)]
pub enum TableReadFormat {
#[serde(rename = "csv")]
Csv,
#[serde(rename = "json")]
JsonString,
#[serde(rename = "columns")]
ColumnsString,
#[serde(rename = "arrow")]
Arrow,
#[serde(rename = "ndjson")]
Ndjson,
}
impl TableReadFormat {
pub fn parse(value: Option<String>) -> Result<Option<Self>, String> {
Ok(match value.as_deref() {
Some("csv") => Some(TableReadFormat::Csv),
Some("json") => Some(TableReadFormat::JsonString),
Some("columns") => Some(TableReadFormat::ColumnsString),
Some("arrow") => Some(TableReadFormat::Arrow),
Some("ndjson") => Some(TableReadFormat::Ndjson),
None => None,
Some(x) => return Err(format!("Unknown format \"{x}\"")),
})
}
}
#[derive(Clone, Debug, Default, Serialize, Deserialize, TS)]
pub struct TableInitOptions {
#[serde(default)]
#[ts(optional)]
pub name: Option<String>,
#[serde(default)]
#[ts(optional)]
pub format: Option<TableReadFormat>,
#[serde(default)]
#[ts(optional)]
pub index: Option<String>,
#[serde(default)]
#[ts(optional)]
pub limit: Option<u32>,
#[serde(default)]
#[ts(optional)]
pub page_to_disk: Option<bool>,
#[serde(default)]
#[ts(optional)]
pub list_flatten: Option<crate::proto::ListFlatten>,
#[serde(default)]
#[ts(optional)]
pub schema: Option<IndexMap<String, ColumnType>>,
}
impl TableInitOptions {
pub fn set_name<D: Display>(&mut self, name: D) {
self.name = Some(format!("{name}"))
}
}
impl TryFrom<TableOptions> for MakeTableOptions {
type Error = ClientError;
fn try_from(value: TableOptions) -> Result<Self, Self::Error> {
let page_to_disk = value.page_to_disk;
let list_flatten = value.list_flatten.map(|x| x as i32);
let view_schema = value.schema.clone().map(|schema| crate::proto::Schema {
schema: schema
.into_iter()
.map(|(name, r#type)| crate::proto::schema::KeyTypePair {
name,
r#type: r#type as i32,
})
.collect(),
});
Ok(MakeTableOptions {
page_to_disk,
list_flatten,
view_schema,
make_table_type: match value {
TableOptions {
index: Some(_),
limit: Some(_),
..
} => Err(ClientError::BadTableOptions)?,
TableOptions {
index: Some(index), ..
} => Some(MakeTableType::MakeIndexTable(index)),
TableOptions {
limit: Some(limit), ..
} => Some(MakeTableType::MakeLimitTable(limit)),
_ => None,
},
})
}
}
#[derive(Clone, Debug)]
pub(crate) struct TableOptions {
pub index: Option<String>,
pub limit: Option<u32>,
pub page_to_disk: Option<bool>,
pub list_flatten: Option<crate::proto::ListFlatten>,
pub schema: Option<IndexMap<String, ColumnType>>,
}
#[derive(Clone)]
pub(crate) struct ViewBinding {
pub view: View,
pub update_token: u32,
pub remove_token: Option<u32>,
}
impl From<TableInitOptions> for TableOptions {
fn from(value: TableInitOptions) -> Self {
TableOptions {
index: value.index,
limit: value.limit,
page_to_disk: value.page_to_disk,
list_flatten: value.list_flatten,
schema: value.schema,
}
}
}
#[derive(Clone, Debug, Default, Serialize, Deserialize, TS)]
pub struct JoinOptions {
#[serde(default)]
#[ts(optional)]
pub join_type: Option<crate::proto::JoinType>,
#[serde(default)]
#[ts(optional)]
pub name: Option<String>,
#[serde(default)]
#[ts(optional)]
pub right_on: Option<String>,
}
#[derive(Clone, Debug, Default, Deserialize, TS)]
pub struct DeleteOptions {
pub lazy: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize, TS)]
pub struct UpdateOptions {
pub port_id: Option<u32>,
pub format: Option<TableReadFormat>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ExprValidationResult {
pub expression_schema: Schema,
pub errors: HashMap<String, ExpressionError>,
pub expression_alias: HashMap<String, String>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Description {
pub expression_schema: Schema,
pub view_schema: Schema,
}
#[derive(Clone, Debug, PartialEq)]
pub enum DescribeError {
Expressions {
expression_schema: Schema,
errors: HashMap<String, ExpressionError>,
},
Config(String),
}
impl Display for DescribeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Expressions { errors, .. } => {
let mut names = errors.keys().collect::<Vec<_>>();
names.sort();
write!(f, "Invalid expressions: ")?;
for (i, name) in names.iter().enumerate() {
if i > 0 {
write!(f, "; ")?;
}
write!(f, "{}: {}", name, errors[*name].error_message)?;
}
Ok(())
},
Self::Config(msg) => write!(f, "{}", msg),
}
}
}
impl std::error::Error for DescribeError {}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum DescribeVerdict {
Ok {
expression_schema: Schema,
view_schema: Schema,
},
Expressions {
#[serde(default)]
expression_schema: Schema,
expression_errors: HashMap<String, ExpressionError>,
},
Config {
config_error: String,
},
}
impl From<Result<Description, DescribeError>> for DescribeVerdict {
fn from(value: Result<Description, DescribeError>) -> Self {
match value {
Ok(d) => Self::Ok {
expression_schema: d.expression_schema,
view_schema: d.view_schema,
},
Err(DescribeError::Expressions {
expression_schema,
errors,
}) => Self::Expressions {
expression_schema,
expression_errors: errors,
},
Err(DescribeError::Config(config_error)) => Self::Config { config_error },
}
}
}
impl TryFrom<DescribeVerdict> for Result<Description, DescribeError> {
type Error = ClientError;
fn try_from(value: DescribeVerdict) -> Result<Self, ClientError> {
Ok(match value {
DescribeVerdict::Ok {
expression_schema,
view_schema,
} => Ok(Description {
expression_schema,
view_schema,
}),
DescribeVerdict::Expressions {
expression_schema,
expression_errors,
} => {
if expression_errors.is_empty() {
return Err(ClientError::Unknown(
"Describe verdict has an empty `expression_errors`".to_string(),
));
}
Err(DescribeError::Expressions {
expression_schema,
errors: expression_errors,
})
},
DescribeVerdict::Config { config_error } => Err(DescribeError::Config(config_error)),
})
}
}
fn decode_schema(schema: HashMap<String, i32>) -> ClientResult<Schema> {
schema
.into_iter()
.map(|(name, ty)| {
ColumnType::try_from(ty)
.map(|ty| (name, ty))
.map_err(|e| ClientError::Unknown(e.to_string()))
})
.collect()
}
#[derive(Clone)]
pub struct Table {
name: String,
client: Client,
options: TableOptions,
pub(crate) view_binding: Option<ViewBinding>,
}
assert_table_api!(Table);
impl PartialEq for Table {
fn eq(&self, other: &Self) -> bool {
self.name == other.name && self.client == other.client
}
}
impl Table {
pub(crate) fn new(name: String, client: Client, options: TableOptions) -> Self {
Table {
name,
client,
options,
view_binding: None,
}
}
fn client_message(&self, req: ClientReq) -> Request {
Request {
msg_id: self.client.gen_id(),
entity_id: self.name.clone(),
client_req: Some(req),
}
}
pub fn get_client(&self) -> Client {
self.client.clone()
}
pub async fn get_features(&self) -> ClientResult<Features> {
self.client.get_features().await
}
pub fn get_index(&self) -> Option<String> {
self.options.index.as_ref().map(|index| index.to_owned())
}
pub fn get_limit(&self) -> Option<u32> {
self.options.limit.as_ref().map(|limit| *limit)
}
pub fn get_name(&self) -> &str {
self.name.as_str()
}
pub async fn clear(&self) -> ClientResult<()> {
self.replace(UpdateData::JsonRows("[]".to_owned())).await
}
pub async fn delete(&self, options: DeleteOptions) -> ClientResult<()> {
if let Some(binding) = &self.view_binding {
binding.view.remove_update(binding.update_token).await?;
if let Some(token) = binding.remove_token {
binding.view.remove_remove(token).await?;
}
}
let msg = self.client_message(ClientReq::TableDeleteReq(TableDeleteReq {
is_immediate: !options.lazy,
}));
match self.client.oneshot(&msg).await? {
ClientResp::TableDeleteResp(_) => Ok(()),
resp => Err(resp.into()),
}
}
pub async fn columns(&self) -> ClientResult<Vec<String>> {
let msg = self.client_message(ClientReq::TableSchemaReq(TableSchemaReq {}));
match self.client.oneshot(&msg).await? {
ClientResp::TableSchemaResp(TableSchemaResp { schema }) => Ok(schema
.map(|x| x.schema.into_iter().map(|x| x.name.to_owned()).collect())
.unwrap()),
resp => Err(resp.into()),
}
}
pub async fn size(&self) -> ClientResult<usize> {
let msg = self.client_message(ClientReq::TableSizeReq(TableSizeReq {}));
match self.client.oneshot(&msg).await? {
ClientResp::TableSizeResp(TableSizeResp { size }) => Ok(size as usize),
resp => Err(resp.into()),
}
}
pub async fn schema(&self) -> ClientResult<Schema> {
let msg = self.client_message(ClientReq::TableSchemaReq(TableSchemaReq {}));
match self.client.oneshot(&msg).await? {
ClientResp::TableSchemaResp(TableSchemaResp { schema }) => Ok(schema
.map(|x| {
x.schema
.into_iter()
.map(|x| (x.name, ColumnType::try_from(x.r#type).unwrap()))
.collect()
})
.unwrap()),
resp => Err(resp.into()),
}
}
pub async fn make_port(&self) -> ClientResult<i32> {
let msg = self.client_message(ClientReq::TableMakePortReq(TableMakePortReq {}));
match self.client.oneshot(&msg).await? {
ClientResp::TableMakePortResp(TableMakePortResp { port_id }) => Ok(port_id as i32),
_ => Err(ClientError::Unknown("make_port".to_string())),
}
}
pub async fn on_delete(
&self,
on_delete: Box<dyn Fn() + Send + Sync + 'static>,
) -> ClientResult<u32> {
let callback = move |resp: Response| match resp.client_resp {
Some(ClientResp::TableOnDeleteResp(_)) => {
on_delete();
Ok(())
},
resp => Err(resp.into()),
};
let msg = self.client_message(ClientReq::TableOnDeleteReq(TableOnDeleteReq {}));
self.client.subscribe_once(&msg, Box::new(callback)).await?;
Ok(msg.msg_id)
}
pub async fn remove_delete(&self, callback_id: u32) -> ClientResult<()> {
let msg = self.client_message(ClientReq::TableRemoveDeleteReq(TableRemoveDeleteReq {
id: callback_id,
}));
match self.client.oneshot(&msg).await? {
ClientResp::TableRemoveDeleteResp(_) => Ok(()),
resp => Err(resp.into()),
}
}
pub async fn remove(&self, input: UpdateData) -> ClientResult<()> {
let msg = self.client_message(ClientReq::TableRemoveReq(TableRemoveReq {
data: Some(input.into()),
}));
match self.client.oneshot(&msg).await? {
ClientResp::TableRemoveResp(_) => Ok(()),
resp => Err(resp.into()),
}
}
pub async fn replace(&self, input: UpdateData) -> ClientResult<()> {
let msg = self.client_message(ClientReq::TableReplaceReq(TableReplaceReq {
data: Some(input.into()),
}));
match self.client.oneshot(&msg).await? {
ClientResp::TableReplaceResp(_) => Ok(()),
resp => Err(resp.into()),
}
}
pub async fn update(&self, input: UpdateData, options: UpdateOptions) -> ClientResult<()> {
let msg = self.client_message(ClientReq::TableUpdateReq(TableUpdateReq {
data: Some(input.into()),
port_id: options.port_id.unwrap_or(0),
}));
match self.client.oneshot(&msg).await? {
ClientResp::TableUpdateResp(_) => Ok(()),
resp => Err(resp.into()),
}
}
pub async fn describe(
&self,
config: ViewConfigUpdate,
) -> ClientResult<Result<Description, DescribeError>> {
let msg = self.client_message(ClientReq::TableDescribeReq(TableDescribeReq {
config: Some(config.into()),
}));
let resp = match self.client.oneshot(&msg).await? {
ClientResp::TableDescribeResp(resp) => resp,
resp => return Err(resp.into()),
};
let expression_schema = decode_schema(resp.expression_schema)?;
let has_errors = !resp.expression_errors.is_empty();
Ok(match (resp.result, has_errors) {
(Some(table_describe_resp::Result::View(view)), false) => Ok(Description {
expression_schema,
view_schema: decode_schema(view.schema)?,
}),
(Some(table_describe_resp::Result::ConfigError(msg)), false) => {
Err(DescribeError::Config(msg))
},
(None, true) => Err(DescribeError::Expressions {
expression_schema,
errors: resp.expression_errors,
}),
(result, _) => {
return Err(ClientError::Unknown(format!(
"Malformed describe response: result={:?} expression_errors={:?}",
result.is_some(),
resp.expression_errors.keys().collect::<Vec<_>>()
)));
},
})
}
pub async fn validate_expressions(
&self,
expressions: Expressions,
) -> ClientResult<ExprValidationResult> {
let expression_alias = expressions.0.clone();
let config = ViewConfigUpdate {
expressions: Some(expressions),
columns: Some(vec![]),
..ViewConfigUpdate::default()
};
match self.describe(config).await? {
Ok(d) => Ok(ExprValidationResult {
expression_schema: d.expression_schema,
errors: HashMap::new(),
expression_alias,
}),
Err(DescribeError::Expressions {
expression_schema,
errors,
}) => Ok(ExprValidationResult {
expression_schema,
errors,
expression_alias,
}),
Err(DescribeError::Config(msg)) => Err(ClientError::Internal(msg)),
}
}
pub async fn view(&self, config: Option<ViewConfigUpdate>) -> ClientResult<View> {
let view_name = randid();
let msg = Request {
msg_id: self.client.gen_id(),
entity_id: self.name.clone(),
client_req: ClientReq::TableMakeViewReq(TableMakeViewReq {
view_id: view_name.clone(),
config: config.map(|x| x.into()),
})
.into(),
};
match self.client.oneshot(&msg).await? {
ClientResp::TableMakeViewResp(TableMakeViewResp { view_id })
if view_id == view_name =>
{
Ok(View::new_with_source(
view_name,
self.client.clone(),
ViewSource {
options: self.options.clone(),
},
))
},
resp => Err(resp.into()),
}
}
}