Skip to main content

toasty_driver_dynamodb/
lib.rs

1#![warn(missing_docs)]
2
3//! Toasty driver for [Amazon DynamoDB](https://aws.amazon.com/dynamodb/) using
4//! the [`aws-sdk-dynamodb`](https://docs.rs/aws-sdk-dynamodb) SDK.
5//!
6//! # Examples
7//!
8//! ```no_run
9//! # async fn example() -> toasty_core::Result<()> {
10//! use toasty_driver_dynamodb::DynamoDb;
11//!
12//! let driver = DynamoDb::from_env("dynamodb://localhost".to_string()).await?;
13//! # Ok(())
14//! # }
15//! ```
16
17mod 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/// A DynamoDB [`Driver`] backed by the AWS SDK.
52///
53/// Create one with [`DynamoDb::from_env`] to load AWS credentials and region
54/// from the environment, or [`DynamoDb::new`] / [`DynamoDb::with_sdk_config`]
55/// for manual setup.
56#[derive(Debug, Clone)]
57pub struct DynamoDb {
58    url: String,
59    client: Client,
60}
61
62impl DynamoDb {
63    /// Create driver with pre-built client (backward compatible, synchronous)
64    pub fn new(url: String, client: Client) -> Self {
65        Self { url, client }
66    }
67
68    /// Create driver loading AWS config from environment (async factory)
69    /// Reads: AWS_REGION, AWS_ENDPOINT_URL_DYNAMODB, AWS credentials, etc.
70    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    /// Create driver with custom SdkConfig (synchronous)
79    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        // Clone the shared client - cheap operation (Client uses Arc internally)
100        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        // Use shared client directly
113        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/// An open connection to DynamoDB.
147#[derive(Debug)]
148pub struct Connection {
149    /// Handle to the AWS SDK client
150    client: Client,
151    query_log: QueryLogConfig,
152}
153
154impl Connection {
155    /// Wrap an existing [`aws_sdk_dynamodb::Client`] as a Toasty connection.
156    pub fn new(client: Client) -> Self {
157        Self {
158            client,
159            query_log: QueryLogConfig::default(),
160        }
161    }
162}
163
164/// Resolves the table an operation targets, for the per-query event.
165fn 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
261/// Convert a DynamoDB AttributeValue to stmt::Value (type-inferred).
262fn 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            // Try to parse as i64 first (most common), fallback to string
269            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        // For complex types, convert to string representation
277        _ => stmt::Value::String(format!("{:?}", attr)),
278    }
279}
280
281/// Serialize a DynamoDB LastEvaluatedKey (for pagination) into stmt::Value.
282/// Format: flat record [name1, value1, name2, value2, ...]
283/// Example: { "pk": S("abc"), "sk": N("42") } → Record([String("pk"), String("abc"), String("sk"), I64(42)])
284fn 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
295/// Deserialize a stmt::Value cursor into a DynamoDB ExclusiveStartKey.
296/// Expects flat record format: [name1, value1, name2, value2, ...]
297fn 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        // Process pairs: [name, value, name, value, ...]
302        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                // DynamoDB condition expressions don't support arithmetic
388                // between operands. Arithmetic ops belong in update
389                // expressions (handled by `update_by_key.rs`), not in
390                // condition/filter expressions.
391                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            // A bare boolean column reference used as a predicate (result of
401            // `field = true` simplification) needs an explicit equality check.
402            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        // A projection into a `#[document]` column (`profile().name()`)
411        // arrives lowered as a `FuncJsonExtract` name path. DynamoDB filter
412        // expressions address nested Map attributes natively, so it renders
413        // as `#col_0.#doc_1.#doc_2`.
414        stmt::Expr::Func(stmt::ExprFunc::JsonExtract(func)) => {
415            let (path, leaf_ty) = document_path(cx, attrs, func);
416            // Like a bare bool column reference, a bool leaf in predicate
417            // position needs an explicit equality check.
418            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            // Extract the list items and create individual attribute values
445            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                    // If it's not a literal list, treat it as a single expression
453                    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            // `attribute_not_exists` takes a bare attribute path. Resolve a
461            // column alias or document path directly rather than through
462            // `ddb_expression`, which would expand a bool column/leaf to
463            // `#col = :true` — a comparison valid only in predicate position,
464            // not as a function argument. (Without this, `.is_none()` on any
465            // `Option<bool>` — including an `Option<Embed>` presence column —
466            // produces invalid syntax.)
467            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            // `Path::contains(value)` lowers to `value = ANY(col)`. On
492            // DynamoDB that's `contains(path, value)` — the standard List
493            // membership filter.
494            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            // Bool key/index fields bridge through I8 (db::Type::Integer(1) via
504            // bridge_type). The lowering wraps the I8 column ref in
505            // Cast(col_ref, Bool) when the field appears as a bare predicate
506            // (result of `field = true` simplification). In predicate position
507            // this means "is true"; the `field = false` case arrives as
508            // Not(Cast(col_ref, Bool)) and is handled by the Not arm above.
509            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
518/// Resolves a column reference to its DynamoDB attribute alias (e.g. `#col_3`),
519/// registering the underlying attribute name in `attrs`. Returns the resolved
520/// column alongside the alias so callers can inspect its storage type.
521fn 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
531/// Renders a lowered document path ([`stmt::FuncJsonExtract`]) as a DynamoDB
532/// attribute path (`#col_0.#doc_1.#doc_2`), registering each path segment as
533/// an expression attribute name. The path arrives fully resolved from the
534/// engine's document lowering — segment names and leaf type live in the node,
535/// so no schema is consulted. Returns the rendered path and the leaf type.
536fn 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    /// Placeholder per document path segment, keyed by the segment (field)
558    /// name so repeated mentions of the same key share one placeholder.
559    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    /// Registers one segment of a document path (a field name inside a Map
579    /// attribute) and returns its placeholder.
580    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}