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        // The URL does not name the endpoint — that comes from the ambient AWS
74        // config — but validate it anyway, so a typo fails here instead of
75        // silently connecting to whatever the environment points at.
76        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    /// Create driver with custom SdkConfig (synchronous)
91    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        // Clone the shared client - cheap operation (Client uses Arc internally)
112        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        // Use shared client directly
125        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/// An open connection to DynamoDB.
159#[derive(Debug)]
160pub struct Connection {
161    /// Handle to the AWS SDK client
162    client: Client,
163    query_log: QueryLogConfig,
164}
165
166impl Connection {
167    /// Wrap an existing [`aws_sdk_dynamodb::Client`] as a Toasty connection.
168    pub fn new(client: Client) -> Self {
169        Self {
170            client,
171            query_log: QueryLogConfig::default(),
172        }
173    }
174}
175
176/// Resolves the table an operation targets, for the per-query event.
177fn 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
269/// Convert a DynamoDB AttributeValue to stmt::Value (type-inferred).
270fn 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            // Try to parse as i64 first (most common), fallback to string
277            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        // For complex types, convert to string representation
285        _ => stmt::Value::String(format!("{:?}", attr)),
286    }
287}
288
289/// Serialize a DynamoDB LastEvaluatedKey (for pagination) into stmt::Value.
290/// Format: flat record [name1, value1, name2, value2, ...]
291/// Example: { "pk": S("abc"), "sk": N("42") } → Record([String("pk"), String("abc"), String("sk"), I64(42)])
292fn 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
303/// Deserialize a stmt::Value cursor into a DynamoDB ExclusiveStartKey.
304/// Expects flat record format: [name1, value1, name2, value2, ...]
305fn 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        // Process pairs: [name, value, name, value, ...]
310        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                // DynamoDB condition expressions don't support arithmetic
396                // between operands. Arithmetic ops belong in update
397                // expressions (handled by `update_by_key.rs`), not in
398                // condition/filter expressions.
399                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            // A bare boolean column reference used as a predicate (result of
409            // `field = true` simplification) needs an explicit equality check.
410            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        // A projection into a `#[document]` column (`profile().name()`)
419        // arrives lowered as a `FuncJsonExtract` name path. DynamoDB filter
420        // expressions address nested Map attributes natively, so it renders
421        // as `#col_0.#doc_1.#doc_2`.
422        stmt::Expr::Func(stmt::ExprFunc::JsonExtract(func)) => {
423            let (path, leaf_ty) = document_path(cx, attrs, func);
424            // Like a bare bool column reference, a bool leaf in predicate
425            // position needs an explicit equality check.
426            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            // Extract the list items and create individual attribute values
453            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                    // If it's not a literal list, treat it as a single expression
461                    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            // `attribute_not_exists` takes a bare attribute path. Resolve a
469            // column alias or document path directly rather than through
470            // `ddb_expression`, which would expand a bool column/leaf to
471            // `#col = :true` — a comparison valid only in predicate position,
472            // not as a function argument. (Without this, `.is_none()` on any
473            // `Option<bool>` — including an `Option<Embed>` presence column —
474            // produces invalid syntax.)
475            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            // `Path::contains(value)` lowers to `value = ANY(col)`. On
500            // DynamoDB that's `contains(path, value)` — the standard List
501            // membership filter.
502            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            // Bool key/index fields bridge through I8 (db::Type::Integer(1) via
512            // bridge_type). The lowering wraps the I8 column ref in
513            // Cast(col_ref, Bool) when the field appears as a bare predicate
514            // (result of `field = true` simplification). In predicate position
515            // this means "is true"; the `field = false` case arrives as
516            // Not(Cast(col_ref, Bool)) and is handled by the Not arm above.
517            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
526/// Resolves a column reference to its DynamoDB attribute alias (e.g. `#col_3`),
527/// registering the underlying attribute name in `attrs`. Returns the resolved
528/// column alongside the alias so callers can inspect its storage type.
529fn 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
539/// Renders a lowered document path ([`stmt::FuncJsonExtract`]) as a DynamoDB
540/// attribute path (`#col_0.#doc_1.#doc_2`), registering each path segment as
541/// an expression attribute name. The path arrives fully resolved from the
542/// engine's document lowering — segment names and leaf type live in the node,
543/// so no schema is consulted. Returns the rendered path and the leaf type.
544fn 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    /// Placeholder per document path segment, keyed by the segment (field)
566    /// name so repeated mentions of the same key share one placeholder.
567    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    /// Registers one segment of a document path (a field name inside a Map
587    /// attribute) and returns its placeholder.
588    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}