1#![warn(missing_docs)]
2
3mod op;
18mod r#type;
19mod value;
20
21pub(crate) use r#type::TypeExt;
22pub(crate) use value::Value;
23
24use async_trait::async_trait;
25use toasty_core::{
26 Error, Result, Schema,
27 driver::{
28 Capability, ConnectContext, Driver, ExecResponse, QueryLogConfig, log::QueryLog,
29 operation::Operation,
30 },
31 schema::{
32 db::{self, Column, ColumnId, Migration, Table},
33 diff,
34 },
35 stmt::{self, ExprContext},
36};
37
38use aws_sdk_dynamodb::{
39 Client,
40 error::SdkError,
41 operation::transact_write_items::TransactWriteItemsError,
42 operation::update_item::UpdateItemError,
43 types::{
44 AttributeDefinition, AttributeValue, BillingMode, Delete, GlobalSecondaryIndex,
45 KeySchemaElement, KeyType, KeysAndAttributes, Projection, ProjectionType, Put, PutRequest,
46 ReturnValuesOnConditionCheckFailure, TransactWriteItem, Update, WriteRequest,
47 },
48};
49use std::{borrow::Cow, collections::HashMap, sync::Arc};
50
51#[derive(Debug, Clone)]
57pub struct DynamoDb {
58 url: String,
59 client: Client,
60}
61
62impl DynamoDb {
63 pub fn new(url: String, client: Client) -> Self {
65 Self { url, client }
66 }
67
68 pub async fn from_env(url: String) -> Result<Self> {
71 use aws_config::BehaviorVersion;
72
73 let sdk_config = aws_config::defaults(BehaviorVersion::latest()).load().await;
74 let client = Client::new(&sdk_config);
75 Ok(Self::new(url, client))
76 }
77
78 pub fn with_sdk_config(url: String, sdk_config: &aws_config::SdkConfig) -> Self {
80 let client = Client::new(sdk_config);
81 Self::new(url, client)
82 }
83}
84
85#[async_trait]
86impl Driver for DynamoDb {
87 fn url(&self) -> Cow<'_, str> {
88 Cow::Borrowed(&self.url)
89 }
90
91 fn capability(&self) -> &'static Capability {
92 &Capability::DYNAMODB
93 }
94
95 async fn connect(
96 &self,
97 cx: &ConnectContext,
98 ) -> toasty_core::Result<Box<dyn toasty_core::driver::Connection>> {
99 let mut connection = Connection::new(self.client.clone());
101 connection.query_log = cx.query_log;
102 Ok(Box::new(connection))
103 }
104
105 fn generate_migration(&self, _schema_diff: &diff::Schema<'_>) -> Migration {
106 unimplemented!(
107 "DynamoDB migrations are not yet supported. DynamoDB schema changes require manual table updates through the AWS console or SDK."
108 )
109 }
110
111 async fn reset_db(&self) -> toasty_core::Result<()> {
112 let mut exclusive_start_table_name = None;
114 loop {
115 let mut req = self.client.list_tables();
116 if let Some(start) = &exclusive_start_table_name {
117 req = req.exclusive_start_table_name(start);
118 }
119
120 let resp = req
121 .send()
122 .await
123 .map_err(toasty_core::Error::driver_operation_failed)?;
124
125 if let Some(table_names) = &resp.table_names {
126 for table_name in table_names {
127 self.client
128 .delete_table()
129 .table_name(table_name)
130 .send()
131 .await
132 .map_err(toasty_core::Error::driver_operation_failed)?;
133 }
134 }
135
136 exclusive_start_table_name = resp.last_evaluated_table_name;
137 if exclusive_start_table_name.is_none() {
138 break;
139 }
140 }
141
142 Ok(())
143 }
144}
145
146#[derive(Debug)]
148pub struct Connection {
149 client: Client,
151 query_log: QueryLogConfig,
152}
153
154impl Connection {
155 pub fn new(client: Client) -> Self {
157 Self {
158 client,
159 query_log: QueryLogConfig::default(),
160 }
161 }
162}
163
164fn op_table_name<'a>(schema: &'a Schema, op: &Operation) -> Option<&'a str> {
166 let table_id = match op {
167 Operation::GetByKey(op) => op.table,
168 Operation::QueryPk(op) => op.table,
169 Operation::DeleteByKey(op) => op.table,
170 Operation::UpdateByKey(op) => op.table,
171 Operation::FindPkByIndex(op) => op.table,
172 Operation::Scan(op) => op.table,
173 _ => return None,
174 };
175 Some(&schema.db.table(table_id).name)
176}
177
178#[async_trait]
179impl toasty_core::driver::Connection for Connection {
180 async fn exec(&mut self, schema: &Arc<Schema>, op: Operation) -> Result<ExecResponse> {
181 let log = QueryLog::operation(
182 &self.query_log,
183 "dynamodb",
184 op.name(),
185 op_table_name(schema, &op),
186 );
187 let result = self.exec2(schema, op).await;
188 log.finish(&result);
189 result
190 }
191
192 async fn push_schema(&mut self, schema: &Schema) -> Result<()> {
193 for table in &schema.db.tables {
194 tracing::debug!(table = %table.name, "creating table");
195 self.create_table(&schema.db, table, true).await?;
196 }
197 Ok(())
198 }
199
200 async fn applied_migrations(
201 &mut self,
202 ) -> Result<Vec<toasty_core::schema::db::AppliedMigration>> {
203 todo!("DynamoDB migrations are not yet implemented")
204 }
205
206 async fn apply_migration(
207 &mut self,
208 _id: u64,
209 _name: &str,
210 _migration: &toasty_core::schema::db::Migration,
211 ) -> Result<()> {
212 todo!("DynamoDB migrations are not yet implemented")
213 }
214}
215
216impl Connection {
217 async fn exec2(&mut self, schema: &Arc<Schema>, op: Operation) -> Result<ExecResponse> {
218 match op {
219 Operation::GetByKey(op) => self.exec_get_by_key(schema, op).await,
220 Operation::QueryPk(op) => self.exec_query_pk(schema, op).await,
221 Operation::DeleteByKey(op) => self.exec_delete_by_key(&schema.db, op).await,
222 Operation::UpdateByKey(op) => self.exec_update_by_key(&schema.db, op).await,
223 Operation::FindPkByIndex(op) => self.exec_find_pk_by_index(schema, op).await,
224 Operation::QuerySql(op) => {
225 assert!(
226 op.last_insert_id_hack.is_none(),
227 "last_insert_id_hack is MySQL-specific and should not be set for DynamoDB"
228 );
229 match op.stmt {
230 stmt::Statement::Insert(insert) => self.exec_insert(&schema.db, insert).await,
231 _ => todo!("op={:#?}", op.stmt),
232 }
233 }
234 Operation::Scan(op) => self.exec_scan(schema, op).await,
235 Operation::RawSql(_) => Err(Error::unsupported_feature(
236 "raw SQL is only supported by SQL drivers",
237 )),
238 Operation::Transaction(_) => Err(Error::unsupported_feature(
239 "transactions are not supported by the DynamoDB driver",
240 )),
241 _ => todo!("op={op:#?}"),
242 }
243 }
244}
245
246fn ddb_key(table: &Table, key: &stmt::Value) -> HashMap<String, AttributeValue> {
247 let mut ret = HashMap::new();
248
249 for (index, column) in table.primary_key_columns().enumerate() {
250 let value = match key {
251 stmt::Value::Record(record) => &record[index],
252 value => value,
253 };
254
255 ret.insert(column.name.clone(), Value::from(value.clone()).to_ddb());
256 }
257
258 ret
259}
260
261fn attr_value_to_stmt_value(attr: &AttributeValue) -> stmt::Value {
263 use AttributeValue as AV;
264
265 match attr {
266 AV::S(s) => stmt::Value::String(s.clone()),
267 AV::N(n) => {
268 n.parse::<i64>()
270 .map(stmt::Value::I64)
271 .unwrap_or_else(|_| stmt::Value::String(n.clone()))
272 }
273 AV::Bool(b) => stmt::Value::Bool(*b),
274 AV::B(bytes) => stmt::Value::Bytes(bytes.clone().into_inner()),
275 AV::Null(_) => stmt::Value::Null,
276 _ => stmt::Value::String(format!("{:?}", attr)),
278 }
279}
280
281fn serialize_ddb_cursor(last_key: &HashMap<String, AttributeValue>) -> stmt::Value {
285 let mut fields = Vec::with_capacity(last_key.len() * 2);
286
287 for (name, attr_value) in last_key {
288 fields.push(stmt::Value::String(name.clone()));
289 fields.push(attr_value_to_stmt_value(attr_value));
290 }
291
292 stmt::Value::Record(stmt::ValueRecord::from_vec(fields))
293}
294
295fn deserialize_ddb_cursor(cursor: &stmt::Value) -> HashMap<String, AttributeValue> {
298 let mut ret = HashMap::new();
299
300 if let stmt::Value::Record(fields) = cursor {
301 for chunk in fields.chunks(2) {
303 if chunk.len() == 2
304 && let (stmt::Value::String(name), value) = (&chunk[0], &chunk[1])
305 {
306 ret.insert(name.clone(), Value::from(value.clone()).to_ddb());
307 }
308 }
309 }
310
311 ret
312}
313
314fn ddb_key_schema(
315 partition_columns: &[&Column],
316 range_columns: &[&Column],
317) -> Vec<KeySchemaElement> {
318 let mut ks = vec![];
319
320 for col in partition_columns {
321 ks.push(
322 KeySchemaElement::builder()
323 .attribute_name(&col.name)
324 .key_type(KeyType::Hash)
325 .build()
326 .unwrap(),
327 );
328 }
329
330 for col in range_columns {
331 ks.push(
332 KeySchemaElement::builder()
333 .attribute_name(&col.name)
334 .key_type(KeyType::Range)
335 .build()
336 .unwrap(),
337 );
338 }
339
340 ks
341}
342
343fn item_to_record<'a, 'stmt>(
344 item: &HashMap<String, AttributeValue>,
345 columns: impl Iterator<Item = &'a Column>,
346) -> Result<stmt::ValueRecord> {
347 Ok(stmt::ValueRecord::from_vec(
348 columns
349 .map(|column| {
350 if let Some(value) = item.get(&column.name) {
351 Value::from_ddb(&column.ty, value).into_inner()
352 } else {
353 stmt::Value::Null
354 }
355 })
356 .collect(),
357 ))
358}
359
360fn ddb_expression(
361 cx: &ExprContext<'_, db::Schema>,
362 attrs: &mut ExprAttrs,
363 primary: bool,
364 expr: &stmt::Expr,
365) -> String {
366 match expr {
367 stmt::Expr::Between(expr_between) => {
368 let field = ddb_expression(cx, attrs, primary, &expr_between.expr);
369 let low = ddb_expression(cx, attrs, primary, &expr_between.low);
370 let high = ddb_expression(cx, attrs, primary, &expr_between.high);
371 format!("{field} BETWEEN {low} AND {high}")
372 }
373 stmt::Expr::BinaryOp(expr_binary_op) => {
374 let lhs = ddb_expression(cx, attrs, primary, &expr_binary_op.lhs);
375 let rhs = ddb_expression(cx, attrs, primary, &expr_binary_op.rhs);
376
377 match expr_binary_op.op {
378 stmt::BinaryOp::Eq => format!("{lhs} = {rhs}"),
379 stmt::BinaryOp::Ne if primary => {
380 todo!("!= conditions on primary key not supported")
381 }
382 stmt::BinaryOp::Ne => format!("{lhs} <> {rhs}"),
383 stmt::BinaryOp::Gt => format!("{lhs} > {rhs}"),
384 stmt::BinaryOp::Ge => format!("{lhs} >= {rhs}"),
385 stmt::BinaryOp::Lt => format!("{lhs} < {rhs}"),
386 stmt::BinaryOp::Le => format!("{lhs} <= {rhs}"),
387 stmt::BinaryOp::Add | stmt::BinaryOp::Sub => {
392 todo!(
393 "arithmetic operators in DynamoDB condition expressions are not supported"
394 )
395 }
396 }
397 }
398 stmt::Expr::Reference(expr_reference) => {
399 let (column, col_alias) = column_alias(cx, attrs, expr_reference);
400 if column.ty.is_bool() {
403 let true_val = attrs.ddb_value(aws_sdk_dynamodb::types::AttributeValue::Bool(true));
404 format!("{col_alias} = {true_val}")
405 } else {
406 col_alias
407 }
408 }
409 stmt::Expr::Value(val) => attrs.value(val),
410 stmt::Expr::Func(stmt::ExprFunc::JsonExtract(func)) => {
415 let (path, leaf_ty) = document_path(cx, attrs, func);
416 if leaf_ty.is_bool() {
419 let true_val = attrs.ddb_value(aws_sdk_dynamodb::types::AttributeValue::Bool(true));
420 format!("{path} = {true_val}")
421 } else {
422 path
423 }
424 }
425 stmt::Expr::And(expr_and) => {
426 let operands = expr_and
427 .operands
428 .iter()
429 .map(|operand| ddb_expression(cx, attrs, primary, operand))
430 .collect::<Vec<_>>();
431 operands.join(" AND ")
432 }
433 stmt::Expr::Or(expr_or) => {
434 let operands = expr_or
435 .operands
436 .iter()
437 .map(|operand| ddb_expression(cx, attrs, primary, operand))
438 .collect::<Vec<_>>();
439 format!("({})", operands.join(" OR "))
440 }
441 stmt::Expr::InList(in_list) => {
442 let expr = ddb_expression(cx, attrs, primary, &in_list.expr);
443
444 let items = match &*in_list.list {
446 stmt::Expr::Value(stmt::Value::List(vals)) => vals
447 .iter()
448 .map(|val| attrs.value(val))
449 .collect::<Vec<_>>()
450 .join(", "),
451 _ => {
452 ddb_expression(cx, attrs, primary, &in_list.list)
454 }
455 };
456
457 format!("{expr} IN ({items})")
458 }
459 stmt::Expr::IsNull(expr_is_null) => {
460 let inner = match &*expr_is_null.expr {
468 stmt::Expr::Reference(expr_reference) => column_alias(cx, attrs, expr_reference).1,
469 stmt::Expr::Func(stmt::ExprFunc::JsonExtract(func)) => {
470 document_path(cx, attrs, func).0
471 }
472 other => ddb_expression(cx, attrs, primary, other),
473 };
474 format!("attribute_not_exists({inner})")
475 }
476 stmt::Expr::Not(expr_not) => {
477 let inner = ddb_expression(cx, attrs, primary, &expr_not.expr);
478 format!("(NOT {inner})")
479 }
480 stmt::Expr::StartsWith(expr_starts_with) => {
481 let expr = ddb_expression(cx, attrs, primary, &expr_starts_with.expr);
482 let prefix = ddb_expression(cx, attrs, primary, &expr_starts_with.prefix);
483 format!("begins_with({expr}, {prefix})")
484 }
485 stmt::Expr::Like(_) => {
486 panic!(
487 "LIKE is not supported by the DynamoDB driver; use starts_with for prefix matching"
488 )
489 }
490 stmt::Expr::AnyOp(any) if matches!(any.op, stmt::BinaryOp::Eq) => {
491 let value = ddb_expression(cx, attrs, primary, &any.lhs);
495 let path = ddb_expression(cx, attrs, primary, &any.rhs);
496 format!("contains({path}, {value})")
497 }
498 stmt::Expr::Length(expr) => {
499 let inner = ddb_expression(cx, attrs, primary, &expr.expr);
500 format!("size({inner})")
501 }
502 stmt::Expr::Cast(expr_cast) if expr_cast.ty == stmt::Type::Bool => {
503 let col_alias = ddb_expression(cx, attrs, primary, &expr_cast.expr);
510 let true_val =
511 attrs.ddb_value(aws_sdk_dynamodb::types::AttributeValue::N("1".to_string()));
512 format!("{col_alias} = {true_val}")
513 }
514 _ => todo!("FILTER = {:#?}", expr),
515 }
516}
517
518fn column_alias<'a>(
522 cx: &ExprContext<'a, db::Schema>,
523 attrs: &mut ExprAttrs,
524 expr_reference: &stmt::ExprReference,
525) -> (&'a Column, String) {
526 let column = cx.resolve_expr_reference(expr_reference).as_column_unwrap();
527 let alias = attrs.column(column).to_string();
528 (column, alias)
529}
530
531fn document_path<'a>(
537 cx: &ExprContext<'a, db::Schema>,
538 attrs: &mut ExprAttrs,
539 func: &'a stmt::FuncJsonExtract,
540) -> (String, &'a stmt::Type) {
541 let stmt::Expr::Reference(expr_reference) = func.base.as_ref() else {
542 todo!("document path base must be a column reference; func={func:#?}")
543 };
544 let (_, mut path) = column_alias(cx, attrs, expr_reference);
545
546 for name in &func.path {
547 path.push('.');
548 path.push_str(attrs.document_segment(name));
549 }
550
551 (path, &func.ty)
552}
553
554#[derive(Default)]
555struct ExprAttrs {
556 columns: HashMap<ColumnId, String>,
557 document_segments: HashMap<String, String>,
560 attr_names: HashMap<String, String>,
561 attr_values: HashMap<String, AttributeValue>,
562}
563
564impl ExprAttrs {
565 fn column(&mut self, column: &Column) -> &str {
566 use std::collections::hash_map::Entry;
567
568 match self.columns.entry(column.id) {
569 Entry::Vacant(e) => {
570 let name = format!("#col_{}", column.id.index);
571 self.attr_names.insert(name.clone(), column.name.clone());
572 e.insert(name)
573 }
574 Entry::Occupied(e) => e.into_mut(),
575 }
576 }
577
578 fn document_segment(&mut self, name: &str) -> &str {
581 use std::collections::hash_map::Entry;
582
583 match self.document_segments.entry(name.to_owned()) {
584 Entry::Vacant(e) => {
585 let placeholder = format!("#doc_{}", self.attr_names.len());
586 self.attr_names.insert(placeholder.clone(), name.to_owned());
587 e.insert(placeholder)
588 }
589 Entry::Occupied(e) => e.into_mut(),
590 }
591 }
592
593 fn value(&mut self, val: &stmt::Value) -> String {
594 self.ddb_value(Value::from(val.clone()).to_ddb())
595 }
596
597 fn ddb_value(&mut self, val: AttributeValue) -> String {
598 let i = self.attr_values.len();
599 let name = format!(":v_{i}");
600 self.attr_values.insert(name.clone(), val);
601 name
602 }
603}