iota_data_ingestion/
progress_store.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{str::FromStr, time::Duration};
6
7use anyhow::Result;
8use async_trait::async_trait;
9use aws_config::{BehaviorVersion, timeout::TimeoutConfig};
10use aws_sdk_dynamodb::{
11    Client,
12    config::{Credentials, Region},
13    error::SdkError,
14    types::AttributeValue,
15};
16use iota_data_ingestion_core::ProgressStore;
17use iota_types::messages_checkpoint::CheckpointSequenceNumber;
18
19pub struct DynamoDBProgressStore {
20    client: Client,
21    table_name: String,
22}
23
24impl DynamoDBProgressStore {
25    pub async fn new(
26        aws_access_key_id: &str,
27        aws_secret_access_key: &str,
28        aws_region: String,
29        table_name: String,
30    ) -> Self {
31        let credentials = Credentials::new(
32            aws_access_key_id,
33            aws_secret_access_key,
34            None,
35            None,
36            "dynamodb",
37        );
38        let timeout_config = TimeoutConfig::builder()
39            .operation_timeout(Duration::from_secs(3))
40            .operation_attempt_timeout(Duration::from_secs(10))
41            .connect_timeout(Duration::from_secs(3))
42            .build();
43        let http_client = aws_smithy_http_client::Builder::new()
44            .tls_provider(aws_smithy_http_client::tls::Provider::Rustls(
45                aws_smithy_http_client::tls::rustls_provider::CryptoMode::Ring,
46            ))
47            .build_https();
48        let aws_config = aws_config::defaults(BehaviorVersion::latest())
49            .http_client(http_client)
50            .credentials_provider(credentials)
51            .region(Region::new(aws_region))
52            .timeout_config(timeout_config)
53            .load()
54            .await;
55        let client = Client::new(&aws_config);
56        Self { client, table_name }
57    }
58}
59
60#[async_trait]
61impl ProgressStore for DynamoDBProgressStore {
62    type Error = anyhow::Error;
63
64    async fn load(&mut self, task_name: String) -> Result<CheckpointSequenceNumber, Self::Error> {
65        let item = self
66            .client
67            .get_item()
68            .table_name(self.table_name.clone())
69            .key("task_name", AttributeValue::S(task_name))
70            .send()
71            .await?;
72        if let Some(output) = item.item() {
73            if let AttributeValue::N(checkpoint_number) = &output["nstate"] {
74                return Ok(CheckpointSequenceNumber::from_str(checkpoint_number)?);
75            }
76        }
77        Ok(0)
78    }
79    async fn save(
80        &mut self,
81        task_name: String,
82        checkpoint_number: CheckpointSequenceNumber,
83    ) -> Result<(), Self::Error> {
84        let backoff = backoff::ExponentialBackoff::default();
85        backoff::future::retry(backoff, || async {
86            let result = self
87                .client
88                .update_item()
89                .table_name(self.table_name.clone())
90                .key("task_name", AttributeValue::S(task_name.clone()))
91                .update_expression("SET #nstate = :newState")
92                .condition_expression("#nstate < :newState")
93                .expression_attribute_names("#nstate", "nstate")
94                .expression_attribute_values(
95                    ":newState",
96                    AttributeValue::N(checkpoint_number.to_string()),
97                )
98                .send()
99                .await;
100            match result {
101                Ok(_) => Ok(()),
102                Err(SdkError::ServiceError(err))
103                    if err.err().is_conditional_check_failed_exception() =>
104                {
105                    Ok(())
106                }
107                Err(err) => Err(backoff::Error::transient(err)),
108            }
109        })
110        .await?;
111        Ok(())
112    }
113}
iota_data_ingestion/progress_store.rs

iota_data_ingestion/
progress_store.rs
