use std::collections::{BTreeSet, HashMap};
use aws_config::{BehaviorVersion, Region};
use aws_sdk_dynamodb::Client;
use aws_sdk_dynamodb::primitives::Blob;
use aws_sdk_dynamodb::types::AttributeValue;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
use tokio::sync::Mutex;
use super::error::DbError;
use super::model::{Column, QueryOutcome, Row, SchemaColumn, TablePage, Value};
#[derive(Debug, Clone, Default)]
pub struct DynamoOptions {
pub profile: Option<String>,
pub region: Option<String>,
pub endpoint_url: Option<String>,
pub read_only: bool,
pub local: bool,
}
#[derive(Debug)]
pub struct DynamoQueryOptions {
pub table: String,
pub key_condition: String,
pub index: Option<String>,
pub filter: Option<String>,
pub projection: Option<String>,
pub values: Option<serde_json::Value>,
pub names: Option<serde_json::Value>,
pub limit: usize,
pub consistent_read: bool,
pub typed_json: bool,
pub start_key: Option<serde_json::Value>,
pub scan_forward: bool,
}
#[derive(Debug)]
pub struct DynamoScanOptions {
pub table: String,
pub filter: Option<String>,
pub projection: Option<String>,
pub values: Option<serde_json::Value>,
pub names: Option<serde_json::Value>,
pub limit: usize,
pub consistent_read: bool,
pub typed_json: bool,
pub start_key: Option<serde_json::Value>,
}
pub struct DynamoSource {
client: Client,
read_only: bool,
page_keys: Mutex<HashMap<(String, usize), HashMap<String, AttributeValue>>>,
}
impl DynamoSource {
pub async fn connect(options: DynamoOptions) -> Result<Self, DbError> {
let DynamoOptions {
profile,
mut region,
mut endpoint_url,
read_only,
local,
} = options;
if local {
region.get_or_insert_with(|| "us-east-1".to_string());
endpoint_url.get_or_insert_with(|| "http://localhost:8000".to_string());
}
let mut loader = aws_config::defaults(BehaviorVersion::latest());
if !local && let Some(profile) = profile {
loader = loader.profile_name(profile);
}
if let Some(region) = region {
loader = loader.region(Region::new(region));
}
let shared = loader.load().await;
let mut config = aws_sdk_dynamodb::config::Builder::from(&shared);
if let Some(endpoint_url) = endpoint_url {
config = config.endpoint_url(endpoint_url);
}
if local {
config = config.credentials_provider(aws_sdk_dynamodb::config::Credentials::new(
"test",
"test",
None,
None,
"tuible-local",
));
}
Ok(Self {
client: Client::from_conf(config.build()),
read_only,
page_keys: Mutex::new(HashMap::new()),
})
}
pub fn is_read_only(&self) -> bool {
self.read_only
}
pub async fn list_tables(&self) -> Result<Vec<String>, DbError> {
let mut tables = Vec::new();
let mut start = None;
loop {
let output = self
.client
.list_tables()
.set_exclusive_start_table_name(start)
.send()
.await
.map_err(dynamo_error)?;
tables.extend(output.table_names().iter().cloned());
start = output.last_evaluated_table_name().map(str::to_string);
if start.is_none() {
break;
}
}
tables.sort();
Ok(tables)
}
pub async fn table_schema(&self, table: &str) -> Result<Vec<SchemaColumn>, DbError> {
let output = self
.client
.describe_table()
.table_name(table)
.send()
.await
.map_err(dynamo_error)?;
let Some(description) = output.table() else {
return Err(DbError::Dynamo(format!(
"table description missing for {table}"
)));
};
let types: HashMap<&str, &str> = description
.attribute_definitions()
.iter()
.map(|attribute| {
(
attribute.attribute_name(),
attribute.attribute_type().as_str(),
)
})
.collect();
Ok(description
.key_schema()
.iter()
.map(|key| SchemaColumn {
name: key.attribute_name().to_string(),
col_type: types
.get(key.attribute_name())
.copied()
.map_or_else(|| "UNKNOWN".to_string(), str::to_string),
notnull: true,
pk: true,
})
.collect())
}
pub async fn fetch_rows(
&self,
table: &str,
limit: i64,
offset: i64,
) -> Result<TablePage, DbError> {
let limit = dynamo_limit(limit);
let offset = offset.max(0) as usize;
let start_key = if offset == 0 {
None
} else {
self.page_keys
.lock()
.await
.get(&(table.to_string(), offset))
.cloned()
};
if offset > 0 && start_key.is_none() {
return Err(DbError::Dynamo(
"page cursor expired; refresh the table and page forward again".to_string(),
));
}
let output = self
.client
.scan()
.table_name(table)
.limit(limit)
.set_exclusive_start_key(start_key)
.send()
.await
.map_err(dynamo_error)?;
let last_key = output.last_evaluated_key().cloned();
let has_more = last_key.is_some();
if let Some(last_key) = last_key {
self.page_keys
.lock()
.await
.insert((table.to_string(), offset + limit as usize), last_key);
}
let key_names = self.key_names(table).await.unwrap_or_default();
let (columns, rows) = items_to_rows(output.items(), &key_names);
Ok(TablePage {
columns,
rows,
rowids: None,
offset,
has_more,
})
}
pub async fn execute_sql(
&self,
statement: &str,
max_rows: usize,
) -> Result<QueryOutcome, DbError> {
self.execute_partiql(statement, max_rows, None, false, None)
.await
}
pub async fn execute_partiql(
&self,
statement: &str,
max_rows: usize,
parameters: Option<&serde_json::Value>,
typed_json: bool,
next_token: Option<String>,
) -> Result<QueryOutcome, DbError> {
let is_select = Self::statement_is_read_only(statement);
if self.read_only && !is_select {
return Err(DbError::ReadOnly);
}
let parameters = parameters
.map(|parameters| json_parameters(parameters, typed_json))
.transpose()?;
let output = self
.client
.execute_statement()
.statement(statement)
.set_parameters(parameters)
.set_next_token(next_token)
.limit(usize_to_limit(max_rows))
.send()
.await
.map_err(dynamo_error)?;
if !is_select {
return Ok(QueryOutcome::Executed);
}
let mut items = output.items().to_vec();
let next_token = output.next_token().map(str::to_string);
let truncated = items.len() > max_rows || next_token.is_some();
items.truncate(max_rows);
let (columns, rows) = items_to_rows(&items, &[]);
Ok(QueryOutcome::Rows {
columns,
rows,
truncated,
next_token,
read_only: true,
})
}
pub fn statement_is_read_only(statement: &str) -> bool {
first_keyword(statement).as_deref() == Some("select")
}
pub async fn describe_table_json(&self, table: &str) -> Result<serde_json::Value, DbError> {
let output = self
.client
.describe_table()
.table_name(table)
.send()
.await
.map_err(dynamo_error)?;
let Some(description) = output.table() else {
return Err(DbError::Dynamo(format!(
"table description missing for {table}"
)));
};
let key_schema: Vec<_> = description
.key_schema()
.iter()
.map(|key| {
serde_json::json!({
"attribute": key.attribute_name(),
"key_type": key.key_type().as_str(),
})
})
.collect();
let attributes: Vec<_> = description
.attribute_definitions()
.iter()
.map(|attribute| {
serde_json::json!({
"name": attribute.attribute_name(),
"type": attribute.attribute_type().as_str(),
})
})
.collect();
let global_indexes: Vec<_> = description
.global_secondary_indexes()
.iter()
.map(|index| {
serde_json::json!({
"name": index.index_name(),
"status": index.index_status().map(|status| status.as_str()),
"item_count": index.item_count(),
"size_bytes": index.index_size_bytes(),
"key_schema": index.key_schema().iter().map(|key| serde_json::json!({
"attribute": key.attribute_name(),
"key_type": key.key_type().as_str(),
})).collect::<Vec<_>>(),
"projection": index.projection().map(|projection| serde_json::json!({
"type": projection.projection_type().map(|kind| kind.as_str()),
"non_key_attributes": projection.non_key_attributes(),
})),
})
})
.collect();
let local_indexes: Vec<_> = description
.local_secondary_indexes()
.iter()
.map(|index| {
serde_json::json!({
"name": index.index_name(),
"item_count": index.item_count(),
"size_bytes": index.index_size_bytes(),
"key_schema": index.key_schema().iter().map(|key| serde_json::json!({
"attribute": key.attribute_name(),
"key_type": key.key_type().as_str(),
})).collect::<Vec<_>>(),
"projection": index.projection().map(|projection| serde_json::json!({
"type": projection.projection_type().map(|kind| kind.as_str()),
"non_key_attributes": projection.non_key_attributes(),
})),
})
})
.collect();
Ok(serde_json::json!({
"table_name": description.table_name(),
"table_arn": description.table_arn(),
"status": description.table_status().map(|status| status.as_str()),
"billing_mode": description.billing_mode_summary()
.and_then(|summary| summary.billing_mode())
.map(|mode| mode.as_str()),
"item_count": description.item_count(),
"size_bytes": description.table_size_bytes(),
"key_schema": key_schema,
"attributes": attributes,
"global_secondary_indexes": global_indexes,
"local_secondary_indexes": local_indexes,
}))
}
pub async fn scan_json(
&self,
options: DynamoScanOptions,
) -> Result<serde_json::Value, DbError> {
let values = options
.values
.as_ref()
.map(|value| json_object_to_item(value, options.typed_json))
.transpose()?;
let names = options.names.as_ref().map(json_string_map).transpose()?;
let start_key = options
.start_key
.as_ref()
.map(|value| json_object_to_item(value, options.typed_json))
.transpose()?;
let output = self
.client
.scan()
.table_name(options.table)
.set_filter_expression(options.filter)
.set_projection_expression(options.projection)
.set_expression_attribute_values(values)
.set_expression_attribute_names(names)
.set_exclusive_start_key(start_key)
.limit(usize_to_limit(options.limit))
.consistent_read(options.consistent_read)
.send()
.await
.map_err(dynamo_error)?;
Ok(items_json(
output.items(),
output.last_evaluated_key(),
output.scanned_count(),
))
}
pub async fn query_json(
&self,
options: DynamoQueryOptions,
) -> Result<serde_json::Value, DbError> {
let values = options
.values
.as_ref()
.map(|value| json_object_to_item(value, options.typed_json))
.transpose()?;
let names = options.names.as_ref().map(json_string_map).transpose()?;
let start_key = options
.start_key
.as_ref()
.map(|value| json_object_to_item(value, options.typed_json))
.transpose()?;
let output = self
.client
.query()
.table_name(options.table)
.key_condition_expression(options.key_condition)
.set_index_name(options.index)
.set_filter_expression(options.filter)
.set_projection_expression(options.projection)
.set_expression_attribute_values(values)
.set_expression_attribute_names(names)
.set_exclusive_start_key(start_key)
.limit(usize_to_limit(options.limit))
.consistent_read(options.consistent_read)
.scan_index_forward(options.scan_forward)
.send()
.await
.map_err(dynamo_error)?;
Ok(items_json(
output.items(),
output.last_evaluated_key(),
output.scanned_count(),
))
}
pub async fn get_item_json(
&self,
table: &str,
key: &serde_json::Value,
consistent_read: bool,
typed_json: bool,
projection: Option<String>,
names: Option<&serde_json::Value>,
) -> Result<serde_json::Value, DbError> {
let names = names.map(json_string_map).transpose()?;
let output = self
.client
.get_item()
.table_name(table)
.set_key(Some(json_object_to_item(key, typed_json)?))
.consistent_read(consistent_read)
.set_projection_expression(projection)
.set_expression_attribute_names(names)
.send()
.await
.map_err(dynamo_error)?;
Ok(match output.item() {
Some(item) => serde_json::json!({
"item": item_to_json(item),
"dynamodb_item": item_to_typed_json(item),
}),
None => serde_json::json!({
"item": null,
"dynamodb_item": null,
}),
})
}
pub async fn put_item_json(
&self,
table: &str,
item: &serde_json::Value,
typed_json: bool,
condition: Option<String>,
values: Option<&serde_json::Value>,
names: Option<&serde_json::Value>,
) -> Result<(), DbError> {
self.ensure_writable()?;
let values = values
.map(|value| json_object_to_item(value, typed_json))
.transpose()?;
let names = names.map(json_string_map).transpose()?;
self.client
.put_item()
.table_name(table)
.set_item(Some(json_object_to_item(item, typed_json)?))
.set_condition_expression(condition)
.set_expression_attribute_values(values)
.set_expression_attribute_names(names)
.send()
.await
.map_err(dynamo_error)?;
Ok(())
}
pub async fn delete_item_json(
&self,
table: &str,
key: &serde_json::Value,
typed_json: bool,
condition: Option<String>,
values: Option<&serde_json::Value>,
names: Option<&serde_json::Value>,
) -> Result<(), DbError> {
self.ensure_writable()?;
let values = values
.map(|value| json_object_to_item(value, typed_json))
.transpose()?;
let names = names.map(json_string_map).transpose()?;
self.client
.delete_item()
.table_name(table)
.set_key(Some(json_object_to_item(key, typed_json)?))
.set_condition_expression(condition)
.set_expression_attribute_values(values)
.set_expression_attribute_names(names)
.send()
.await
.map_err(dynamo_error)?;
Ok(())
}
async fn key_names(&self, table: &str) -> Result<Vec<String>, DbError> {
Ok(self
.table_schema(table)
.await?
.into_iter()
.map(|column| column.name)
.collect())
}
fn ensure_writable(&self) -> Result<(), DbError> {
if self.read_only {
Err(DbError::ReadOnly)
} else {
Ok(())
}
}
}
fn dynamo_error(error: impl std::error::Error) -> DbError {
let mut message = error.to_string();
let mut source = error.source();
while let Some(cause) = source {
let cause_message = cause.to_string();
if !cause_message.is_empty() && !message.contains(&cause_message) {
message.push_str(": ");
message.push_str(&cause_message);
}
source = cause.source();
}
DbError::Dynamo(message)
}
fn dynamo_limit(limit: i64) -> i32 {
limit.clamp(1, 1_000) as i32
}
fn usize_to_limit(limit: usize) -> i32 {
limit.clamp(1, 1_000) as i32
}
fn items_to_rows(
items: &[HashMap<String, AttributeValue>],
key_names: &[String],
) -> (Vec<Column>, Vec<Row>) {
let mut names = Vec::new();
let mut seen = BTreeSet::new();
for name in key_names {
if seen.insert(name.clone()) {
names.push(name.clone());
}
}
for item in items {
for name in item.keys() {
if seen.insert(name.clone()) {
names.push(name.clone());
}
}
}
if names.len() > key_names.len() {
names[key_names.len()..].sort();
}
let rows = items
.iter()
.map(|item| Row {
values: names
.iter()
.map(|name| item.get(name).map_or(Value::Null, attribute_to_value))
.collect(),
})
.collect();
(
names.into_iter().map(|name| Column { name }).collect(),
rows,
)
}
fn attribute_to_value(value: &AttributeValue) -> Value {
match value {
AttributeValue::S(value) => Value::Text(value.clone()),
AttributeValue::N(value) => Value::Decimal(value.clone()),
AttributeValue::B(value) => Value::Bytes(value.clone().into_inner()),
AttributeValue::Bool(value) => Value::Bool(*value),
AttributeValue::Null(_) => Value::Null,
AttributeValue::M(_)
| AttributeValue::L(_)
| AttributeValue::Ss(_)
| AttributeValue::Ns(_)
| AttributeValue::Bs(_) => Value::Json(attribute_to_json(value)),
_ => Value::Text(format!("{value:?}")),
}
}
fn attribute_to_json(value: &AttributeValue) -> serde_json::Value {
match value {
AttributeValue::S(value) => value.clone().into(),
AttributeValue::N(value) => value
.parse::<serde_json::Number>()
.map(serde_json::Value::Number)
.unwrap_or_else(|_| serde_json::Value::String(value.clone())),
AttributeValue::B(value) => serde_json::json!({ "$binary": BASE64.encode(value.as_ref()) }),
AttributeValue::Bool(value) => (*value).into(),
AttributeValue::Null(_) => serde_json::Value::Null,
AttributeValue::M(values) => serde_json::Value::Object(
values
.iter()
.map(|(key, value)| (key.clone(), attribute_to_json(value)))
.collect(),
),
AttributeValue::L(values) => values.iter().map(attribute_to_json).collect(),
AttributeValue::Ss(values) => serde_json::json!({ "$string_set": values }),
AttributeValue::Ns(values) => serde_json::json!({
"$number_set": values.iter().map(|value| {
value.parse::<serde_json::Number>()
.map(serde_json::Value::Number)
.unwrap_or_else(|_| serde_json::Value::String(value.clone()))
}).collect::<Vec<_>>()
}),
AttributeValue::Bs(values) => serde_json::json!({
"$binary_set": values.iter().map(|value| BASE64.encode(value.as_ref())).collect::<Vec<_>>()
}),
_ => serde_json::Value::String(format!("{value:?}")),
}
}
fn item_to_json(item: &HashMap<String, AttributeValue>) -> serde_json::Value {
serde_json::Value::Object(
item.iter()
.map(|(key, value)| (key.clone(), attribute_to_json(value)))
.collect(),
)
}
fn item_to_typed_json(item: &HashMap<String, AttributeValue>) -> serde_json::Value {
serde_json::Value::Object(
item.iter()
.map(|(key, value)| (key.clone(), attribute_to_typed_json(value)))
.collect(),
)
}
fn attribute_to_typed_json(value: &AttributeValue) -> serde_json::Value {
match value {
AttributeValue::S(value) => serde_json::json!({ "S": value }),
AttributeValue::N(value) => serde_json::json!({ "N": value }),
AttributeValue::B(value) => serde_json::json!({ "B": BASE64.encode(value.as_ref()) }),
AttributeValue::Bool(value) => serde_json::json!({ "BOOL": value }),
AttributeValue::Null(value) => serde_json::json!({ "NULL": value }),
AttributeValue::M(values) => serde_json::json!({ "M": item_to_typed_json(values) }),
AttributeValue::L(values) => serde_json::json!({
"L": values.iter().map(attribute_to_typed_json).collect::<Vec<_>>()
}),
AttributeValue::Ss(values) => serde_json::json!({ "SS": values }),
AttributeValue::Ns(values) => serde_json::json!({ "NS": values }),
AttributeValue::Bs(values) => serde_json::json!({
"BS": values.iter().map(|value| BASE64.encode(value.as_ref())).collect::<Vec<_>>()
}),
_ => serde_json::json!({ "UNKNOWN": format!("{value:?}") }),
}
}
fn items_json(
items: &[HashMap<String, AttributeValue>],
last_key: Option<&HashMap<String, AttributeValue>>,
scanned_count: i32,
) -> serde_json::Value {
serde_json::json!({
"items": items.iter().map(item_to_json).collect::<Vec<_>>(),
"dynamodb_items": items.iter().map(item_to_typed_json).collect::<Vec<_>>(),
"count": items.len(),
"scanned_count": scanned_count,
"last_evaluated_key": last_key.map(item_to_json),
"dynamodb_last_evaluated_key": last_key.map(item_to_typed_json),
})
}
pub fn json_object_to_item(
value: &serde_json::Value,
typed_json: bool,
) -> Result<HashMap<String, AttributeValue>, DbError> {
let Some(object) = value.as_object() else {
return Err(DbError::InvalidDynamoJson(
"expected a JSON object".to_string(),
));
};
object
.iter()
.map(|(key, value)| {
let value = if typed_json {
typed_json_to_attribute(value)?
} else {
plain_json_to_attribute(value)?
};
Ok((key.clone(), value))
})
.collect()
}
fn json_parameters(
value: &serde_json::Value,
typed_json: bool,
) -> Result<Vec<AttributeValue>, DbError> {
let Some(values) = value.as_array() else {
return Err(DbError::InvalidDynamoJson(
"PartiQL parameters must be a JSON array".to_string(),
));
};
values
.iter()
.map(|value| {
if typed_json {
typed_json_to_attribute(value)
} else {
plain_json_to_attribute(value)
}
})
.collect()
}
fn plain_json_to_attribute(value: &serde_json::Value) -> Result<AttributeValue, DbError> {
Ok(match value {
serde_json::Value::Null => AttributeValue::Null(true),
serde_json::Value::Bool(value) => AttributeValue::Bool(*value),
serde_json::Value::Number(value) => AttributeValue::N(value.to_string()),
serde_json::Value::String(value) => AttributeValue::S(value.clone()),
serde_json::Value::Array(values) => AttributeValue::L(
values
.iter()
.map(plain_json_to_attribute)
.collect::<Result<_, _>>()?,
),
serde_json::Value::Object(values) => AttributeValue::M(
values
.iter()
.map(|(key, value)| Ok((key.clone(), plain_json_to_attribute(value)?)))
.collect::<Result<_, DbError>>()?,
),
})
}
fn typed_json_to_attribute(value: &serde_json::Value) -> Result<AttributeValue, DbError> {
let Some(object) = value.as_object() else {
return Err(DbError::InvalidDynamoJson(
"typed values must look like {\"S\":\"value\"}".to_string(),
));
};
if object.len() != 1 {
return Err(DbError::InvalidDynamoJson(
"typed values must contain exactly one DynamoDB type".to_string(),
));
}
let Some((kind, value)) = object.iter().next() else {
return Err(DbError::InvalidDynamoJson(
"typed value is empty".to_string(),
));
};
match kind.as_str() {
"S" => json_string(value).map(AttributeValue::S),
"N" => json_number_string(value).map(AttributeValue::N),
"B" => BASE64
.decode(json_string(value)?)
.map(Blob::new)
.map(AttributeValue::B)
.map_err(|error| DbError::InvalidDynamoJson(format!("invalid base64: {error}"))),
"BOOL" => value
.as_bool()
.map(AttributeValue::Bool)
.ok_or_else(|| DbError::InvalidDynamoJson("BOOL must contain a boolean".to_string())),
"NULL" => value
.as_bool()
.map(AttributeValue::Null)
.ok_or_else(|| DbError::InvalidDynamoJson("NULL must contain a boolean".to_string())),
"SS" => json_string_array(value).map(AttributeValue::Ss),
"NS" => json_number_string_array(value).map(AttributeValue::Ns),
"BS" => json_string_array(value).and_then(|values| {
values
.into_iter()
.map(|value| {
BASE64.decode(value).map(Blob::new).map_err(|error| {
DbError::InvalidDynamoJson(format!("invalid base64: {error}"))
})
})
.collect::<Result<Vec<_>, _>>()
.map(AttributeValue::Bs)
}),
"L" => {
let Some(values) = value.as_array() else {
return Err(DbError::InvalidDynamoJson(
"L must contain an array".to_string(),
));
};
values
.iter()
.map(typed_json_to_attribute)
.collect::<Result<Vec<_>, _>>()
.map(AttributeValue::L)
}
"M" => json_object_to_item(value, true).map(AttributeValue::M),
_ => Err(DbError::InvalidDynamoJson(format!(
"unknown DynamoDB type {kind}"
))),
}
}
fn json_string(value: &serde_json::Value) -> Result<String, DbError> {
value
.as_str()
.map(str::to_string)
.ok_or_else(|| DbError::InvalidDynamoJson("expected a string".to_string()))
}
fn json_number_string(value: &serde_json::Value) -> Result<String, DbError> {
match value {
serde_json::Value::String(value) => Ok(value.clone()),
serde_json::Value::Number(value) => Ok(value.to_string()),
_ => Err(DbError::InvalidDynamoJson(
"expected a number or numeric string".to_string(),
)),
}
}
fn json_string_array(value: &serde_json::Value) -> Result<Vec<String>, DbError> {
let Some(values) = value.as_array() else {
return Err(DbError::InvalidDynamoJson(
"expected an array of strings".to_string(),
));
};
values.iter().map(json_string).collect()
}
fn json_number_string_array(value: &serde_json::Value) -> Result<Vec<String>, DbError> {
let Some(values) = value.as_array() else {
return Err(DbError::InvalidDynamoJson(
"expected an array of numbers".to_string(),
));
};
values.iter().map(json_number_string).collect()
}
fn json_string_map(value: &serde_json::Value) -> Result<HashMap<String, String>, DbError> {
let Some(values) = value.as_object() else {
return Err(DbError::InvalidDynamoJson(
"expression names must be a JSON object".to_string(),
));
};
values
.iter()
.map(|(key, value)| Ok((key.clone(), json_string(value)?)))
.collect()
}
fn first_keyword(mut statement: &str) -> Option<String> {
loop {
statement = statement.trim_start();
if let Some(comment) = statement.strip_prefix("--") {
statement = comment.split_once('\n').map_or("", |(_, rest)| rest);
continue;
}
if let Some(comment) = statement.strip_prefix("/*") {
statement = comment.split_once("*/").map_or("", |(_, rest)| rest);
continue;
}
break;
}
let keyword: String = statement
.chars()
.take_while(|character| character.is_ascii_alphabetic())
.flat_map(char::to_lowercase)
.collect();
(!keyword.is_empty()).then_some(keyword)
}
#[cfg(test)]
mod tests {
#![allow(clippy::expect_used, clippy::unwrap_used)]
use super::*;
use std::io::{Read, Write};
use std::net::TcpListener;
fn mock_dynamodb(responses: Vec<&'static str>) -> (String, std::thread::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let address = listener.local_addr().unwrap();
let handle = std::thread::spawn(move || {
for body in responses {
let (mut stream, _) = listener.accept().unwrap();
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = stream.read(&mut buffer).unwrap();
if read == 0 {
break;
}
request.extend_from_slice(&buffer[..read]);
let Some(header_end) = request.windows(4).position(|part| part == b"\r\n\r\n")
else {
continue;
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
line.to_ascii_lowercase()
.strip_prefix("content-length:")
.and_then(|value| value.trim().parse::<usize>().ok())
})
.unwrap_or(0);
if request.len() >= header_end + 4 + content_length {
break;
}
}
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/x-amz-json-1.0\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
stream.write_all(response.as_bytes()).unwrap();
}
});
(format!("http://{address}"), handle)
}
async fn local_source(endpoint_url: String, read_only: bool) -> DynamoSource {
DynamoSource::connect(DynamoOptions {
endpoint_url: Some(endpoint_url),
read_only,
local: true,
..DynamoOptions::default()
})
.await
.unwrap()
}
#[test]
fn plain_json_preserves_nested_dynamo_types() {
let item = json_object_to_item(
&serde_json::json!({
"pk": "USER#1",
"count": 12345678901234567890123456789012345678_u128,
"active": true,
"tags": ["a", "b"],
"meta": {"nested": null},
}),
false,
)
.unwrap();
assert!(matches!(item.get("pk"), Some(AttributeValue::S(value)) if value == "USER#1"));
assert!(matches!(
item.get("active"),
Some(AttributeValue::Bool(true))
));
assert!(matches!(item.get("tags"), Some(AttributeValue::L(values)) if values.len() == 2));
}
#[test]
fn typed_json_supports_sets_and_binary() {
let item = json_object_to_item(
&serde_json::json!({
"roles": {"SS": ["admin", "author"]},
"version": {"N": "9007199254740993"},
"payload": {"B": "AQID"},
}),
true,
)
.unwrap();
assert!(matches!(item.get("roles"), Some(AttributeValue::Ss(values)) if values.len() == 2));
assert!(
matches!(item.get("payload"), Some(AttributeValue::B(value)) if value.as_ref() == [1, 2, 3])
);
}
#[test]
fn plain_json_sentinel_names_remain_regular_maps() {
let item = json_object_to_item(
&serde_json::json!({
"pk": {"$binary": "AQID"},
"roles": {"$string_set": ["admin", "author"]},
}),
false,
)
.unwrap();
assert!(matches!(item.get("pk"), Some(AttributeValue::M(_))));
assert!(matches!(item.get("roles"), Some(AttributeValue::M(_))));
}
#[test]
fn typed_output_round_trips_binary_and_sets_losslessly() {
let original = HashMap::from([
(
"payload".to_string(),
AttributeValue::B(Blob::new([1, 2, 3])),
),
(
"roles".to_string(),
AttributeValue::Ss(vec!["admin".to_string(), "author".to_string()]),
),
]);
let decoded = json_object_to_item(&item_to_typed_json(&original), true).unwrap();
assert_eq!(decoded, original);
}
#[test]
fn item_rows_put_keys_first_and_preserve_exact_numbers() {
let item = HashMap::from([
(
"payload".to_string(),
AttributeValue::M(HashMap::from([(
"ok".to_string(),
AttributeValue::Bool(true),
)])),
),
("pk".to_string(), AttributeValue::S("USER#1".to_string())),
(
"amount".to_string(),
AttributeValue::N("12345678901234567890.123".to_string()),
),
]);
let (columns, rows) = items_to_rows(&[item], &["pk".to_string()]);
assert_eq!(columns[0].name, "pk");
assert!(
rows[0]
.values
.contains(&Value::Decimal("12345678901234567890.123".to_string()))
);
assert!(
rows[0]
.values
.iter()
.any(|value| matches!(value, Value::Json(_)))
);
}
#[test]
fn partiql_read_detection_skips_comments() {
assert_eq!(
first_keyword("-- explain\n SELECT * FROM x").as_deref(),
Some("select")
);
assert_eq!(
first_keyword("/* write */ UPDATE x SET a=1").as_deref(),
Some("update")
);
}
#[tokio::test]
async fn aws_client_lists_tables_through_the_real_http_stack() {
let (endpoint, server) = mock_dynamodb(vec![r#"{"TableNames":["jobs","users"]}"#]);
let source = local_source(endpoint, true).await;
let tables = source.list_tables().await.unwrap();
server.join().unwrap();
assert_eq!(tables, ["jobs", "users"]);
}
#[tokio::test]
async fn aws_client_scans_and_decodes_items_with_key_columns_first() {
let scan = r#"{"Items":[{"active":{"BOOL":true},"count":{"N":"9007199254740993"},"pk":{"S":"USER#1"}}],"Count":1,"ScannedCount":1}"#;
let describe = r#"{"Table":{"AttributeDefinitions":[{"AttributeName":"pk","AttributeType":"S"}],"TableName":"users","KeySchema":[{"AttributeName":"pk","KeyType":"HASH"}],"TableStatus":"ACTIVE","CreationDateTime":1.0,"ProvisionedThroughput":{"NumberOfDecreasesToday":0,"ReadCapacityUnits":0,"WriteCapacityUnits":0},"TableSizeBytes":0,"ItemCount":1,"TableArn":"arn:aws:dynamodb:us-east-1:000000000000:table/users","TableId":"test"}}"#;
let (endpoint, server) = mock_dynamodb(vec![scan, describe]);
let source = local_source(endpoint, true).await;
let page = source.fetch_rows("users", 100, 0).await.unwrap();
server.join().unwrap();
assert_eq!(page.columns[0].name, "pk");
assert_eq!(page.rows[0].values[0], Value::Text("USER#1".to_string()));
assert!(
page.rows[0]
.values
.contains(&Value::Decimal("9007199254740993".to_string()))
);
assert!(page.rows[0].values.contains(&Value::Bool(true)));
}
#[tokio::test]
async fn scan_rows_survive_denied_optional_describe_table() {
let scan =
r#"{"Items":[{"pk":{"S":"USER#1"},"name":{"S":"Ada"}}],"Count":1,"ScannedCount":1}"#;
let (endpoint, server) = mock_dynamodb(vec![scan, "{}"]);
let source = local_source(endpoint, true).await;
let page = source.fetch_rows("users", 100, 0).await.unwrap();
server.join().unwrap();
assert_eq!(page.rows.len(), 1);
assert!(page.columns.iter().any(|column| column.name == "pk"));
assert!(page.columns.iter().any(|column| column.name == "name"));
}
#[tokio::test]
async fn table_pages_reuse_dynamodb_continuation_keys() {
let first_scan = r#"{"Items":[{"pk":{"S":"USER#1"}}],"Count":1,"ScannedCount":1,"LastEvaluatedKey":{"pk":{"S":"USER#1"}}}"#;
let second_scan = r#"{"Items":[{"pk":{"S":"USER#2"}}],"Count":1,"ScannedCount":1}"#;
let describe = r#"{"Table":{"AttributeDefinitions":[{"AttributeName":"pk","AttributeType":"S"}],"TableName":"users","KeySchema":[{"AttributeName":"pk","KeyType":"HASH"}],"TableStatus":"ACTIVE","CreationDateTime":1.0,"ProvisionedThroughput":{"NumberOfDecreasesToday":0,"ReadCapacityUnits":0,"WriteCapacityUnits":0},"TableSizeBytes":0,"ItemCount":2,"TableArn":"arn:aws:dynamodb:us-east-1:000000000000:table/users","TableId":"test"}}"#;
let (endpoint, server) = mock_dynamodb(vec![first_scan, describe, second_scan, describe]);
let source = local_source(endpoint, true).await;
let first = source.fetch_rows("users", 1, 0).await.unwrap();
let second = source.fetch_rows("users", 1, 1).await.unwrap();
server.join().unwrap();
assert!(first.has_more);
assert!(!second.has_more);
assert_eq!(second.rows[0].values[0], Value::Text("USER#2".to_string()));
}
#[tokio::test]
async fn partiql_returns_a_resume_token_without_fetching_another_page() {
let response = r#"{"Items":[{"pk":{"S":"USER#1"}}],"NextToken":"next-page"}"#;
let (endpoint, server) = mock_dynamodb(vec![response]);
let source = local_source(endpoint, true).await;
let outcome = source
.execute_partiql("SELECT * FROM users", 1, None, false, None)
.await
.unwrap();
server.join().unwrap();
assert!(matches!(
outcome,
QueryOutcome::Rows {
next_token: Some(token),
truncated: true,
..
} if token == "next-page"
));
}
#[tokio::test]
async fn read_only_partiql_rejects_writes_before_network_io() {
let source = local_source("http://127.0.0.1:1".to_string(), true).await;
let result = source
.execute_sql("UPDATE users SET active = true WHERE pk = 'USER#1'", 100)
.await;
assert!(matches!(result, Err(DbError::ReadOnly)));
}
}