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 parsed = toasty_core::driver::ConnectionUrl::parse(&url)?;
77 parsed.validate_authority()?;
78
79 if !parsed.has_scheme("dynamodb") {
80 return Err(toasty_core::Error::invalid_connection_url(format!(
81 "connection URL does not have a `dynamodb` scheme; url={url}"
82 )));
83 }
84
85 let sdk_config = aws_config::defaults(BehaviorVersion::latest()).load().await;
86 let client = Client::new(&sdk_config);
87 Ok(Self::new(url, client))
88 }
89
90 pub fn with_sdk_config(url: String, sdk_config: &aws_config::SdkConfig) -> Self {
92 let client = Client::new(sdk_config);
93 Self::new(url, client)
94 }
95}
96
97#[async_trait]
98impl Driver for DynamoDb {
99 fn url(&self) -> Cow<'_, str> {
100 Cow::Borrowed(&self.url)
101 }
102
103 fn capability(&self) -> &'static Capability {
104 &Capability::DYNAMODB
105 }
106
107 async fn connect(
108 &self,
109 cx: &ConnectContext,
110 ) -> toasty_core::Result<Box<dyn toasty_core::driver::Connection>> {
111 let mut connection = Connection::new(self.client.clone());
113 connection.query_log = cx.query_log;
114 Ok(Box::new(connection))
115 }
116
117 fn generate_migration(&self, _schema_diff: &diff::Schema<'_>) -> Migration {
118 unimplemented!(
119 "DynamoDB migrations are not yet supported. DynamoDB schema changes require manual table updates through the AWS console or SDK."
120 )
121 }
122
123 async fn reset_db(&self) -> toasty_core::Result<()> {
124 let mut exclusive_start_table_name = None;
126 loop {
127 let mut req = self.client.list_tables();
128 if let Some(start) = &exclusive_start_table_name {
129 req = req.exclusive_start_table_name(start);
130 }
131
132 let resp = req
133 .send()
134 .await
135 .map_err(toasty_core::Error::driver_operation_failed)?;
136
137 if let Some(table_names) = &resp.table_names {
138 for table_name in table_names {
139 self.client
140 .delete_table()
141 .table_name(table_name)
142 .send()
143 .await
144 .map_err(toasty_core::Error::driver_operation_failed)?;
145 }
146 }
147
148 exclusive_start_table_name = resp.last_evaluated_table_name;
149 if exclusive_start_table_name.is_none() {
150 break;
151 }
152 }
153
154 Ok(())
155 }
156}
157
158#[derive(Debug)]
160pub struct Connection {
161 client: Client,
163 query_log: QueryLogConfig,
164}
165
166impl Connection {
167 pub fn new(client: Client) -> Self {
169 Self {
170 client,
171 query_log: QueryLogConfig::default(),
172 }
173 }
174}
175
176fn op_table_name<'a>(schema: &'a Schema, op: &Operation) -> Option<&'a str> {
178 let table_id = match op {
179 Operation::GetByKey(op) => op.table,
180 Operation::QueryPk(op) => op.table,
181 Operation::DeleteByKey(op) => op.table,
182 Operation::UpdateByKey(op) => op.table,
183 Operation::Upsert(op) => op.stmt.target.as_table_unwrap().table,
184 Operation::FindPkByIndex(op) => op.table,
185 Operation::Scan(op) => op.table,
186 _ => return None,
187 };
188 Some(&schema.db.table(table_id).name)
189}
190
191#[async_trait]
192impl toasty_core::driver::Connection for Connection {
193 async fn exec(&mut self, schema: &Arc<Schema>, op: Operation) -> Result<ExecResponse> {
194 let log = QueryLog::operation(
195 &self.query_log,
196 "dynamodb",
197 op.name(),
198 op_table_name(schema, &op),
199 );
200 let result = self.exec2(schema, op).await;
201 log.finish(&result);
202 result
203 }
204
205 async fn push_schema(&mut self, schema: &Schema) -> Result<()> {
206 for table in &schema.db.tables {
207 tracing::debug!(table = %table.name, "creating table");
208 self.create_table(&schema.db, table, true).await?;
209 }
210 Ok(())
211 }
212
213 async fn applied_migrations(
214 &mut self,
215 ) -> Result<Vec<toasty_core::schema::db::AppliedMigration>> {
216 todo!("DynamoDB migrations are not yet implemented")
217 }
218
219 async fn apply_migration(
220 &mut self,
221 _id: u64,
222 _name: &str,
223 _migration: &toasty_core::schema::db::Migration,
224 ) -> Result<()> {
225 todo!("DynamoDB migrations are not yet implemented")
226 }
227}
228
229impl Connection {
230 async fn exec2(&mut self, schema: &Arc<Schema>, op: Operation) -> Result<ExecResponse> {
231 match op {
232 Operation::GetByKey(op) => self.exec_get_by_key(schema, op).await,
233 Operation::QueryPk(op) => self.exec_query_pk(schema, op).await,
234 Operation::DeleteByKey(op) => self.exec_delete_by_key(&schema.db, op).await,
235 Operation::UpdateByKey(op) => self.exec_update_by_key(&schema.db, op).await,
236 Operation::Upsert(op) => self.exec_upsert(&schema.db, op).await,
237 Operation::FindPkByIndex(op) => self.exec_find_pk_by_index(schema, op).await,
238 Operation::Insert(op) => {
239 self.exec_insert(&schema.db, op.stmt.into_insert_unwrap())
240 .await
241 }
242 Operation::QuerySql(op) => todo!("op={:#?}", op.stmt),
243 Operation::Scan(op) => self.exec_scan(schema, op).await,
244 Operation::RawSql(_) => Err(Error::unsupported_feature(
245 "raw SQL is only supported by SQL drivers",
246 )),
247 Operation::Transaction(_) => Err(Error::unsupported_feature(
248 "transactions are not supported by the DynamoDB driver",
249 )),
250 }
251 }
252}
253
254fn ddb_key(table: &Table, key: &stmt::Value) -> HashMap<String, AttributeValue> {
255 let mut ret = HashMap::new();
256
257 for (index, column) in table.primary_key_columns().enumerate() {
258 let value = match key {
259 stmt::Value::Record(record) => &record[index],
260 value => value,
261 };
262
263 ret.insert(column.name.clone(), Value::from(value.clone()).to_ddb());
264 }
265
266 ret
267}
268
269fn attr_value_to_stmt_value(attr: &AttributeValue) -> stmt::Value {
271 use AttributeValue as AV;
272
273 match attr {
274 AV::S(s) => stmt::Value::String(s.clone()),
275 AV::N(n) => {
276 n.parse::<i64>()
278 .map(stmt::Value::I64)
279 .unwrap_or_else(|_| stmt::Value::String(n.clone()))
280 }
281 AV::Bool(b) => stmt::Value::Bool(*b),
282 AV::B(bytes) => stmt::Value::Bytes(bytes.clone().into_inner()),
283 AV::Null(_) => stmt::Value::Null,
284 _ => stmt::Value::String(format!("{:?}", attr)),
286 }
287}
288
289fn serialize_ddb_cursor(last_key: &HashMap<String, AttributeValue>) -> stmt::Value {
293 let mut fields = Vec::with_capacity(last_key.len() * 2);
294
295 for (name, attr_value) in last_key {
296 fields.push(stmt::Value::String(name.clone()));
297 fields.push(attr_value_to_stmt_value(attr_value));
298 }
299
300 stmt::Value::Record(stmt::ValueRecord::from_vec(fields))
301}
302
303fn deserialize_ddb_cursor(cursor: &stmt::Value) -> HashMap<String, AttributeValue> {
306 let mut ret = HashMap::new();
307
308 if let stmt::Value::Record(fields) = cursor {
309 for chunk in fields.chunks(2) {
311 if chunk.len() == 2
312 && let (stmt::Value::String(name), value) = (&chunk[0], &chunk[1])
313 {
314 ret.insert(name.clone(), Value::from(value.clone()).to_ddb());
315 }
316 }
317 }
318
319 ret
320}
321
322fn ddb_key_schema(
323 partition_columns: &[&Column],
324 range_columns: &[&Column],
325) -> Vec<KeySchemaElement> {
326 let mut ks = vec![];
327
328 for col in partition_columns {
329 ks.push(
330 KeySchemaElement::builder()
331 .attribute_name(&col.name)
332 .key_type(KeyType::Hash)
333 .build()
334 .unwrap(),
335 );
336 }
337
338 for col in range_columns {
339 ks.push(
340 KeySchemaElement::builder()
341 .attribute_name(&col.name)
342 .key_type(KeyType::Range)
343 .build()
344 .unwrap(),
345 );
346 }
347
348 ks
349}
350
351fn item_to_record<'a, 'stmt>(
352 item: &HashMap<String, AttributeValue>,
353 columns: impl Iterator<Item = &'a Column>,
354) -> Result<stmt::ValueRecord> {
355 Ok(stmt::ValueRecord::from_vec(
356 columns
357 .map(|column| {
358 if let Some(value) = item.get(&column.name) {
359 Value::from_ddb(&column.ty, value).into_inner()
360 } else {
361 stmt::Value::Null
362 }
363 })
364 .collect(),
365 ))
366}
367
368fn ddb_expression(
369 cx: &ExprContext<'_, db::Schema>,
370 attrs: &mut ExprAttrs,
371 primary: bool,
372 expr: &stmt::Expr,
373) -> String {
374 match expr {
375 stmt::Expr::Between(expr_between) => {
376 let field = ddb_expression(cx, attrs, primary, &expr_between.expr);
377 let low = ddb_expression(cx, attrs, primary, &expr_between.low);
378 let high = ddb_expression(cx, attrs, primary, &expr_between.high);
379 format!("{field} BETWEEN {low} AND {high}")
380 }
381 stmt::Expr::BinaryOp(expr_binary_op) => {
382 let lhs = ddb_expression(cx, attrs, primary, &expr_binary_op.lhs);
383 let rhs = ddb_expression(cx, attrs, primary, &expr_binary_op.rhs);
384
385 match expr_binary_op.op {
386 stmt::BinaryOp::Eq => format!("{lhs} = {rhs}"),
387 stmt::BinaryOp::Ne if primary => {
388 todo!("!= conditions on primary key not supported")
389 }
390 stmt::BinaryOp::Ne => format!("{lhs} <> {rhs}"),
391 stmt::BinaryOp::Gt => format!("{lhs} > {rhs}"),
392 stmt::BinaryOp::Ge => format!("{lhs} >= {rhs}"),
393 stmt::BinaryOp::Lt => format!("{lhs} < {rhs}"),
394 stmt::BinaryOp::Le => format!("{lhs} <= {rhs}"),
395 stmt::BinaryOp::Add | stmt::BinaryOp::Sub => {
400 todo!(
401 "arithmetic operators in DynamoDB condition expressions are not supported"
402 )
403 }
404 }
405 }
406 stmt::Expr::Reference(expr_reference) => {
407 let (column, col_alias) = column_alias(cx, attrs, expr_reference);
408 if column.ty.is_bool() {
411 let true_val = attrs.ddb_value(aws_sdk_dynamodb::types::AttributeValue::Bool(true));
412 format!("{col_alias} = {true_val}")
413 } else {
414 col_alias
415 }
416 }
417 stmt::Expr::Value(val) => attrs.value(val),
418 stmt::Expr::Func(stmt::ExprFunc::JsonExtract(func)) => {
423 let (path, leaf_ty) = document_path(cx, attrs, func);
424 if leaf_ty.is_bool() {
427 let true_val = attrs.ddb_value(aws_sdk_dynamodb::types::AttributeValue::Bool(true));
428 format!("{path} = {true_val}")
429 } else {
430 path
431 }
432 }
433 stmt::Expr::And(expr_and) => {
434 let operands = expr_and
435 .operands
436 .iter()
437 .map(|operand| ddb_expression(cx, attrs, primary, operand))
438 .collect::<Vec<_>>();
439 operands.join(" AND ")
440 }
441 stmt::Expr::Or(expr_or) => {
442 let operands = expr_or
443 .operands
444 .iter()
445 .map(|operand| ddb_expression(cx, attrs, primary, operand))
446 .collect::<Vec<_>>();
447 format!("({})", operands.join(" OR "))
448 }
449 stmt::Expr::InList(in_list) => {
450 let expr = ddb_expression(cx, attrs, primary, &in_list.expr);
451
452 let items = match &*in_list.list {
454 stmt::Expr::Value(stmt::Value::List(vals)) => vals
455 .iter()
456 .map(|val| attrs.value(val))
457 .collect::<Vec<_>>()
458 .join(", "),
459 _ => {
460 ddb_expression(cx, attrs, primary, &in_list.list)
462 }
463 };
464
465 format!("{expr} IN ({items})")
466 }
467 stmt::Expr::IsNull(expr_is_null) => {
468 let inner = match &*expr_is_null.expr {
476 stmt::Expr::Reference(expr_reference) => column_alias(cx, attrs, expr_reference).1,
477 stmt::Expr::Func(stmt::ExprFunc::JsonExtract(func)) => {
478 document_path(cx, attrs, func).0
479 }
480 other => ddb_expression(cx, attrs, primary, other),
481 };
482 format!("attribute_not_exists({inner})")
483 }
484 stmt::Expr::Not(expr_not) => {
485 let inner = ddb_expression(cx, attrs, primary, &expr_not.expr);
486 format!("(NOT {inner})")
487 }
488 stmt::Expr::StartsWith(expr_starts_with) => {
489 let expr = ddb_expression(cx, attrs, primary, &expr_starts_with.expr);
490 let prefix = ddb_expression(cx, attrs, primary, &expr_starts_with.prefix);
491 format!("begins_with({expr}, {prefix})")
492 }
493 stmt::Expr::Like(_) => {
494 panic!(
495 "LIKE is not supported by the DynamoDB driver; use starts_with for prefix matching"
496 )
497 }
498 stmt::Expr::AnyOp(any) if matches!(any.op, stmt::BinaryOp::Eq) => {
499 let value = ddb_expression(cx, attrs, primary, &any.lhs);
503 let path = ddb_expression(cx, attrs, primary, &any.rhs);
504 format!("contains({path}, {value})")
505 }
506 stmt::Expr::Length(expr) => {
507 let inner = ddb_expression(cx, attrs, primary, &expr.expr);
508 format!("size({inner})")
509 }
510 stmt::Expr::Cast(expr_cast) if expr_cast.ty == stmt::Type::Bool => {
511 let col_alias = ddb_expression(cx, attrs, primary, &expr_cast.expr);
518 let true_val =
519 attrs.ddb_value(aws_sdk_dynamodb::types::AttributeValue::N("1".to_string()));
520 format!("{col_alias} = {true_val}")
521 }
522 _ => todo!("FILTER = {:#?}", expr),
523 }
524}
525
526fn column_alias<'a>(
530 cx: &ExprContext<'a, db::Schema>,
531 attrs: &mut ExprAttrs,
532 expr_reference: &stmt::ExprReference,
533) -> (&'a Column, String) {
534 let column = cx.resolve_expr_reference(expr_reference).as_column_unwrap();
535 let alias = attrs.column(column).to_string();
536 (column, alias)
537}
538
539fn document_path<'a>(
545 cx: &ExprContext<'a, db::Schema>,
546 attrs: &mut ExprAttrs,
547 func: &'a stmt::FuncJsonExtract,
548) -> (String, &'a stmt::Type) {
549 let stmt::Expr::Reference(expr_reference) = func.base.as_ref() else {
550 todo!("document path base must be a column reference; func={func:#?}")
551 };
552 let (_, mut path) = column_alias(cx, attrs, expr_reference);
553
554 for name in &func.path {
555 path.push('.');
556 path.push_str(attrs.document_segment(name));
557 }
558
559 (path, &func.ty)
560}
561
562#[derive(Default)]
563struct ExprAttrs {
564 columns: HashMap<ColumnId, String>,
565 document_segments: HashMap<String, String>,
568 attr_names: HashMap<String, String>,
569 attr_values: HashMap<String, AttributeValue>,
570}
571
572impl ExprAttrs {
573 fn column(&mut self, column: &Column) -> &str {
574 use std::collections::hash_map::Entry;
575
576 match self.columns.entry(column.id) {
577 Entry::Vacant(e) => {
578 let name = format!("#col_{}", column.id.index);
579 self.attr_names.insert(name.clone(), column.name.clone());
580 e.insert(name)
581 }
582 Entry::Occupied(e) => e.into_mut(),
583 }
584 }
585
586 fn document_segment(&mut self, name: &str) -> &str {
589 use std::collections::hash_map::Entry;
590
591 match self.document_segments.entry(name.to_owned()) {
592 Entry::Vacant(e) => {
593 let placeholder = format!("#doc_{}", self.attr_names.len());
594 self.attr_names.insert(placeholder.clone(), name.to_owned());
595 e.insert(placeholder)
596 }
597 Entry::Occupied(e) => e.into_mut(),
598 }
599 }
600
601 fn value(&mut self, val: &stmt::Value) -> String {
602 self.ddb_value(Value::from(val.clone()).to_ddb())
603 }
604
605 fn ddb_value(&mut self, val: AttributeValue) -> String {
606 let i = self.attr_values.len();
607 let name = format!(":v_{i}");
608 self.attr_values.insert(name.clone(), val);
609 name
610 }
611}