use std::sync::Arc;
use arrow::array::{ArrayRef, StringBuilder, UInt64Builder};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::error::ArrowError;
use arrow::record_batch::RecordBatch;
use polyc_proto::proto::polychrome::persona::v1::SpendPolicy;
#[allow(unused_imports)]
use super::wallets::PersonaWalletRow;
#[derive(Debug, Clone)]
pub(crate) struct PersonaSpendPolicyRow {
pub persona_id: String,
pub limit: String,
pub period_secs: u64,
pub max_lifetime_secs: u64,
pub allowed_hosts: String,
pub updated_at_ms: u64,
pub set_by: String,
}
impl PersonaSpendPolicyRow {
#[must_use]
pub(crate) fn new(persona_id: &str, policy: &SpendPolicy) -> Self {
Self {
persona_id: persona_id.to_string(),
limit: policy.limit.clone(),
period_secs: policy.period_secs,
max_lifetime_secs: policy.max_lifetime_secs,
allowed_hosts: serde_json::to_string(&policy.allowed_hosts)
.unwrap_or_else(|_| "[]".to_string()),
updated_at_ms: policy.updated_at_ms,
set_by: policy.set_by.clone(),
}
}
}
#[must_use]
pub(crate) fn persona_spend_policies_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("persona_id", DataType::Utf8, false),
Field::new("limit", DataType::Utf8, false),
Field::new("period_secs", DataType::UInt64, false),
Field::new("max_lifetime_secs", DataType::UInt64, false),
Field::new("allowed_hosts", DataType::Utf8, false),
Field::new("updated_at_ms", DataType::UInt64, false),
Field::new("set_by", DataType::Utf8, false),
]))
}
pub(crate) fn decode_persona_spend_policies_batch(
rows: &[PersonaSpendPolicyRow],
) -> Result<RecordBatch, ArrowError> {
let mut persona_id_b = StringBuilder::with_capacity(rows.len(), rows.len() * 36);
let mut limit_b = StringBuilder::with_capacity(rows.len(), rows.len() * 8);
let mut period_secs_b = UInt64Builder::with_capacity(rows.len());
let mut max_lifetime_secs_b = UInt64Builder::with_capacity(rows.len());
let mut allowed_hosts_b = StringBuilder::with_capacity(rows.len(), rows.len() * 32);
let mut updated_at_ms_b = UInt64Builder::with_capacity(rows.len());
let mut set_by_b = StringBuilder::with_capacity(rows.len(), rows.len() * 36);
for row in rows {
persona_id_b.append_value(&row.persona_id);
limit_b.append_value(&row.limit);
period_secs_b.append_value(row.period_secs);
max_lifetime_secs_b.append_value(row.max_lifetime_secs);
allowed_hosts_b.append_value(&row.allowed_hosts);
updated_at_ms_b.append_value(row.updated_at_ms);
set_by_b.append_value(&row.set_by);
}
let columns: Vec<ArrayRef> = vec![
Arc::new(persona_id_b.finish()),
Arc::new(limit_b.finish()),
Arc::new(period_secs_b.finish()),
Arc::new(max_lifetime_secs_b.finish()),
Arc::new(allowed_hosts_b.finish()),
Arc::new(updated_at_ms_b.finish()),
Arc::new(set_by_b.finish()),
];
RecordBatch::try_new(persona_spend_policies_schema(), columns)
}
#[cfg(test)]
mod tests {
use arrow::array::Array as _;
use super::*;
#[test]
fn persona_spend_policy_round_trips_to_its_row() {
let policy = SpendPolicy {
persona_id: "persona-1".to_string(),
limit: "5".to_string(),
period_secs: 86_400,
max_lifetime_secs: 604_800,
allowed_hosts: vec!["api.example.com".to_string(), "pay.example.com".to_string()],
updated_at_ms: 3_000,
set_by: "admin-1".to_string(),
..Default::default()
};
let row = PersonaSpendPolicyRow::new("persona-1", &policy);
assert_eq!(
row.allowed_hosts,
r#"["api.example.com","pay.example.com"]"#
);
let batch = decode_persona_spend_policies_batch(&[row]).expect("batch build");
assert_eq!(batch.num_rows(), 1);
assert_eq!(batch.schema(), persona_spend_policies_schema());
let limit = batch
.column(1)
.as_any()
.downcast_ref::<arrow::array::StringArray>()
.unwrap()
.value(0);
assert_eq!(limit, "5");
}
#[test]
fn persona_spend_policies_schema_shape() {
let schema = persona_spend_policies_schema();
let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
assert_eq!(
names,
vec![
"persona_id",
"limit",
"period_secs",
"max_lifetime_secs",
"allowed_hosts",
"updated_at_ms",
"set_by",
]
);
for field in schema.fields() {
assert!(!field.is_nullable(), "{} must be non-null", field.name());
}
}
}