dynamodb_tester/
local_client.rs

1use crate::{DynamoClient, TableConfig};
2use anyhow::Result;
3use aws_sdk_dynamodb::{input::CreateTableInput, Client, Endpoint};
4use std::{path::Path, thread};
5use tokio::runtime::Runtime;
6
7pub struct LocalClient {
8    client: Option<Client>,
9    table_name: String,
10    skip_delete: bool,
11}
12
13/// helper function to create a local client easily
14pub async fn local_client(config: impl AsRef<Path>, skip_delete: bool) -> LocalClient {
15    let config = TableConfig::load_from_file(config).expect("valid config");
16    LocalClient::try_new(8000, config, skip_delete)
17        .await
18        .expect("should create")
19}
20
21#[inline(always)]
22pub async fn get_client(config: impl AsRef<Path>, skip_delete: bool) -> impl DynamoClient {
23    local_client(config, skip_delete).await
24}
25
26impl DynamoClient for LocalClient {
27    fn client(&self) -> &Client {
28        let (client, _) = self.inner();
29        client
30    }
31
32    fn table_name(&self) -> Option<&str> {
33        let (_, table_name) = self.inner();
34        Some(table_name)
35    }
36}
37
38impl LocalClient {
39    /// create a new local client
40    pub async fn try_new(port: u16, table_config: TableConfig, skip_delete: bool) -> Result<Self> {
41        let config = aws_config::load_from_env().await;
42        let mut input = CreateTableInput::try_from(table_config)?;
43        let table_name = format!("{}-{}", input.table_name.unwrap(), xid::new());
44        input.table_name = Some(table_name.clone());
45        let dynamodb_local_config = aws_sdk_dynamodb::Config::builder()
46            .region(config.region().cloned())
47            .credentials_provider(
48                config
49                    .credentials_provider()
50                    .expect("cred should exists")
51                    .clone(),
52            )
53            .endpoint_resolver(Endpoint::immutable(
54                format!("http://localhost:{port}").parse().expect("valid"),
55            ))
56            .build();
57        let client = Client::from_conf(dynamodb_local_config);
58
59        client
60            .create_table()
61            .table_name(&table_name)
62            .set_key_schema(input.key_schema)
63            .set_attribute_definitions(input.attribute_definitions)
64            .set_global_secondary_indexes(input.global_secondary_indexes)
65            .set_local_secondary_indexes(input.local_secondary_indexes)
66            .provisioned_throughput(input.provisioned_throughput.take().unwrap())
67            .send()
68            .await?;
69        Ok(Self {
70            client: Some(client),
71            table_name,
72            skip_delete,
73        })
74    }
75
76    /// retrieve the client reference and table name
77    pub fn inner(&self) -> (&Client, &str) {
78        (
79            self.client.as_ref().expect("client"),
80            self.table_name.as_str(),
81        )
82    }
83}
84
85impl Drop for LocalClient {
86    fn drop(&mut self) {
87        let client = self.client.take().expect("client");
88        let table_name = self.table_name.clone();
89        if self.skip_delete {
90            return;
91        }
92        thread::spawn(move || {
93            let rt = Runtime::new().expect("runtime");
94            rt.block_on(async move {
95                if let Err(e) = client.delete_table().table_name(&table_name).send().await {
96                    println!("failed to delete table: {:?}", e);
97                }
98            });
99        });
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[tokio::test]
108    async fn test_local_client() {
109        let config = TableConfig::load_from_file("fixtures/config.yml").unwrap();
110        let client = LocalClient::try_new(8000, config, false).await.unwrap();
111        let (client, table_name) = client.inner();
112        let resp = client
113            .describe_table()
114            .table_name(table_name)
115            .send()
116            .await
117            .unwrap();
118        assert_eq!(resp.table.and_then(|v| v.table_name).unwrap(), table_name);
119    }
120}