iota_indexer/
test_utils.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::path::PathBuf;
6
7use diesel::{QueryableByName, connection::SimpleConnection, sql_types::BigInt};
8use iota_json_rpc_types::IotaTransactionBlockResponse;
9use iota_metrics::init_metrics;
10use tokio::task::JoinHandle;
11use tokio_util::sync::CancellationToken;
12use url::Url;
13
14use crate::{
15    IndexerMetrics,
16    config::{
17        DEFAULT_PRUNING_BATCH_SIZE, IngestionConfig, IotaNamesOptions, PruningOptions,
18        RetentionConfig,
19    },
20    db::{ConnectionPool, ConnectionPoolConfig, PoolConnection, new_connection_pool},
21    errors::IndexerError,
22    indexer::Indexer,
23    store::{PgIndexerAnalyticalStore, PgIndexerStore},
24};
25
26/// Shorter pruning delay used by the test indexer.
27const TEST_PRUNING_DELAY_MS: u64 = 1000; // 1 second
28
29/// Type to create hooks to alter initial indexer DB state in tests.
30/// Those hooks are meant to be called after DB reset (if it occurs) and before
31/// indexer is started.
32///
33/// Example:
34///
35/// ```ignore
36/// let emulate_insertion_order_set_earlier_by_optimistic_indexing: DBInitHook =
37///     Box::new(move |pg_store: &PgIndexerStore| {
38///         transactional_blocking_with_retry!(
39///             &pg_store.blocking_cp(),
40///             |conn| {
41///                 insert_or_ignore_into!(
42///                     tx_insertion_order::table,
43///                     (
44///                         tx_insertion_order::dsl::tx_digest.eq(digest.inner().to_vec()),
45///                         tx_insertion_order::dsl::insertion_order.eq(123),
46///                     ),
47///                     conn
48///                 );
49///                 Ok::<(), IndexerError>(())
50///             },
51///             Duration::from_secs(60)
52///         )
53///             .unwrap()
54///     });
55///
56/// let (_, pg_store, _) = start_simulacrum_grpc_with_write_indexer(
57///     Arc::new(sim),
58///     data_ingestion_path,
59///     None,
60///     Some("indexer_ingestion_tests_db"),
61///     Some(emulate_insertion_order_set_earlier_by_optimistic_indexing),
62/// )
63/// .await;
64/// ```
65pub type DBInitHook = Box<dyn FnOnce(&PgIndexerStore) + Send>;
66
67pub enum IndexerTypeConfig {
68    Reader {
69        reader_mode_rpc_url: String,
70    },
71    Writer {
72        retention_config: Option<RetentionConfig>,
73        pruning_delay_ms: u64,
74        pruning_batch_size: u64,
75    },
76    AnalyticalWorker,
77}
78
79impl IndexerTypeConfig {
80    pub fn reader_mode(reader_mode_rpc_url: String) -> Self {
81        Self::Reader {
82            reader_mode_rpc_url,
83        }
84    }
85
86    pub fn writer_mode(pruning_options: Option<PruningOptions>) -> Self {
87        let opts = pruning_options.unwrap_or_default();
88        Self::Writer {
89            retention_config: opts
90                .load_from_file()
91                .expect("failed to load the indexer retention configuration"),
92            pruning_delay_ms: TEST_PRUNING_DELAY_MS,
93            pruning_batch_size: opts.pruning_batch_size,
94        }
95    }
96
97    /// Writer mode with the retention config given directly, instead of loaded
98    /// from a TOML file via [`PruningOptions`].
99    pub fn writer_mode_with_retention(retention_config: Option<RetentionConfig>) -> Self {
100        Self::Writer {
101            retention_config,
102            pruning_delay_ms: TEST_PRUNING_DELAY_MS,
103            pruning_batch_size: DEFAULT_PRUNING_BATCH_SIZE,
104        }
105    }
106}
107
108pub async fn start_test_indexer(
109    db_url: String,
110    reset_db: bool,
111    db_init_hook: Option<DBInitHook>,
112    rpc_url: String,
113    reader_writer_config: IndexerTypeConfig,
114    data_ingestion_path: Option<PathBuf>,
115) -> (
116    PgIndexerStore,
117    JoinHandle<Result<(), IndexerError>>,
118    CancellationToken,
119) {
120    let token = CancellationToken::new();
121    let (store, handle) = start_test_indexer_impl(
122        db_url,
123        reset_db,
124        db_init_hook,
125        rpc_url,
126        reader_writer_config,
127        data_ingestion_path,
128        token.clone(),
129    )
130    .await;
131    (store, handle, token)
132}
133
134/// Starts an indexer reader or writer for testing depending on the
135/// `reader_writer_config`.
136///
137/// # Note
138/// For [`IndexerTypeConfig::Writer`] when `data_ingestion_path` is `Some`, the
139/// data ingestion path will be exclusively used to ingest data into the
140/// indexer. To force the indexer to sync from the fullnode via gRPC, set
141/// `data_ingestion_path` to `None` and it will use the `rpc_url` to stream
142/// checkpoints from the fullnode gRPC endpoint.
143pub async fn start_test_indexer_impl(
144    db_url: String,
145    reset_db: bool,
146    db_init_hook: Option<DBInitHook>,
147    rpc_url: String,
148    reader_writer_config: IndexerTypeConfig,
149    data_ingestion_path: Option<PathBuf>,
150    cancel: CancellationToken,
151) -> (PgIndexerStore, JoinHandle<Result<(), IndexerError>>) {
152    let store = create_pg_store(&db_url, reset_db);
153    if reset_db {
154        crate::db::reset_database(&mut store.blocking_cp().get().unwrap()).unwrap();
155    }
156    if let Some(db_init_hook) = db_init_hook {
157        db_init_hook(&store);
158    }
159
160    let registry = prometheus_filtered::Registry::default();
161    init_metrics(&registry);
162    let indexer_metrics = IndexerMetrics::new(&registry);
163
164    let handle = match reader_writer_config {
165        IndexerTypeConfig::Reader {
166            reader_mode_rpc_url,
167        } => {
168            let config = crate::config::JsonRpcConfig {
169                iota_names_options: IotaNamesOptions::default(),
170                historic_fallback_options: Default::default(),
171                rpc_address: reader_mode_rpc_url.parse().unwrap(),
172                rpc_client_url: rpc_url,
173            };
174            let pool = store.blocking_cp();
175            let store_clone = store.clone();
176            tokio::spawn(async move {
177                Indexer::start_reader(
178                    &config,
179                    store_clone,
180                    &registry,
181                    pool,
182                    indexer_metrics,
183                    CancellationToken::new(),
184                )
185                .await
186            })
187        }
188        IndexerTypeConfig::Writer {
189            retention_config,
190            pruning_delay_ms,
191            pruning_batch_size,
192        } => {
193            let fullnode_rpc_url = rpc_url.parse::<Url>().unwrap();
194            let store_clone = store.clone();
195            let mut ingestion_config = IngestionConfig::default();
196            ingestion_config.sources.remote_store_url =
197                data_ingestion_path.is_none().then_some(fullnode_rpc_url);
198            ingestion_config.sources.data_ingestion_path = data_ingestion_path;
199
200            tokio::spawn(async move {
201                Indexer::start_writer_with_config(
202                    &ingestion_config,
203                    store_clone,
204                    indexer_metrics,
205                    retention_config,
206                    pruning_delay_ms,
207                    pruning_batch_size,
208                    cancel,
209                )
210                .await
211            })
212        }
213        IndexerTypeConfig::AnalyticalWorker => {
214            let store = PgIndexerAnalyticalStore::new(store.blocking_cp());
215
216            tokio::spawn(
217                async move { Indexer::start_analytical_worker(store, indexer_metrics).await },
218            )
219        }
220    };
221
222    (store, handle)
223}
224
225/// Manage a test database for integration tests.
226pub struct TestDatabase {
227    pub url: String,
228    db_name: String,
229    connection: PoolConnection,
230    pool_config: ConnectionPoolConfig,
231}
232
233impl TestDatabase {
234    pub fn new(db_url: String) -> Self {
235        // Reduce the connection pool size to 5 for testing
236        // to prevent maxing out
237        let pool_config = ConnectionPoolConfig {
238            pool_size: 5,
239            ..Default::default()
240        };
241
242        let db_name = db_url.split('/').next_back().unwrap().into();
243        let (default_url, _) = replace_db_name(&db_url, "postgres");
244        let blocking_pool = new_connection_pool(&default_url, &pool_config).unwrap();
245        let connection = blocking_pool.get().unwrap();
246        Self {
247            url: db_url,
248            db_name,
249            connection,
250            pool_config,
251        }
252    }
253
254    /// Drop the database in the server if it exists.
255    pub fn drop_if_exists(&mut self) {
256        self.connection
257            .batch_execute(&format!("DROP DATABASE IF EXISTS {}", self.db_name))
258            .unwrap();
259    }
260
261    /// Create the database in the server.
262    pub fn create(&mut self) {
263        self.connection
264            .batch_execute(&format!("CREATE DATABASE {}", self.db_name))
265            .unwrap();
266    }
267
268    /// Drop and recreate the database in the server.
269    pub fn recreate(&mut self) {
270        self.drop_if_exists();
271        self.create();
272    }
273
274    /// Create a new connection pool to the database.
275    pub fn to_connection_pool(&self) -> ConnectionPool {
276        new_connection_pool(&self.url, &self.pool_config).unwrap()
277    }
278
279    pub fn reset_db(&mut self) {
280        crate::db::reset_database(&mut self.to_connection_pool().get().unwrap()).unwrap();
281    }
282}
283
284pub fn create_pg_store(db_url: &str, reset_database: bool) -> PgIndexerStore {
285    let registry = prometheus_filtered::Registry::default();
286    init_metrics(&registry);
287    let indexer_metrics = IndexerMetrics::new(&registry);
288
289    let mut test_db = TestDatabase::new(db_url.to_string());
290    if reset_database {
291        test_db.recreate();
292    }
293
294    PgIndexerStore::new(test_db.to_connection_pool(), indexer_metrics)
295}
296
297fn replace_db_name(db_url: &str, new_db_name: &str) -> (String, String) {
298    let pos = db_url.rfind('/').expect("unable to find / in db_url");
299    let old_db_name = &db_url[pos + 1..];
300
301    (
302        format!("{}/{}", &db_url[..pos], new_db_name),
303        old_db_name.to_string(),
304    )
305}
306
307pub async fn force_delete_database(db_url: String) {
308    // Replace the database name with the default `postgres`, which should be the
309    // last string after `/` This is necessary because you can't drop a database
310    // while being connected to it. Hence switch to the default `postgres`
311    // database to drop the active database.
312    let (default_db_url, db_name) = replace_db_name(&db_url, "postgres");
313    let mut pool_config = ConnectionPoolConfig::default();
314    pool_config.set_pool_size(1);
315
316    let blocking_pool = new_connection_pool(&default_db_url, &pool_config).unwrap();
317    blocking_pool
318        .get()
319        .unwrap()
320        .batch_execute(&format!("DROP DATABASE IF EXISTS {db_name} WITH (FORCE)"))
321        .unwrap();
322}
323
324#[derive(Clone)]
325pub struct IotaTransactionBlockResponseBuilder<'a> {
326    response: IotaTransactionBlockResponse,
327    full_response: &'a IotaTransactionBlockResponse,
328}
329
330impl<'a> IotaTransactionBlockResponseBuilder<'a> {
331    pub fn new(full_response: &'a IotaTransactionBlockResponse) -> Self {
332        Self {
333            response: IotaTransactionBlockResponse::default(),
334            full_response,
335        }
336    }
337
338    pub fn with_input(mut self) -> Self {
339        self.response = IotaTransactionBlockResponse {
340            transaction: self.full_response.transaction.clone(),
341            ..self.response
342        };
343        self
344    }
345
346    pub fn with_raw_input(mut self) -> Self {
347        self.response = IotaTransactionBlockResponse {
348            raw_transaction: self.full_response.raw_transaction.clone(),
349            ..self.response
350        };
351        self
352    }
353
354    pub fn with_effects(mut self) -> Self {
355        self.response = IotaTransactionBlockResponse {
356            effects: self.full_response.effects.clone(),
357            ..self.response
358        };
359        self
360    }
361
362    pub fn with_events(mut self) -> Self {
363        self.response = IotaTransactionBlockResponse {
364            events: self.full_response.events.clone(),
365            ..self.response
366        };
367        self
368    }
369
370    pub fn with_balance_changes(mut self) -> Self {
371        self.response = IotaTransactionBlockResponse {
372            balance_changes: self.full_response.balance_changes.clone(),
373            ..self.response
374        };
375        self
376    }
377
378    pub fn with_object_changes(mut self) -> Self {
379        self.response = IotaTransactionBlockResponse {
380            object_changes: self.full_response.object_changes.clone(),
381            ..self.response
382        };
383        self
384    }
385
386    pub fn with_input_and_changes(mut self) -> Self {
387        self.response = IotaTransactionBlockResponse {
388            transaction: self.full_response.transaction.clone(),
389            balance_changes: self.full_response.balance_changes.clone(),
390            object_changes: self.full_response.object_changes.clone(),
391            ..self.response
392        };
393        self
394    }
395
396    pub fn build(self) -> IotaTransactionBlockResponse {
397        IotaTransactionBlockResponse {
398            transaction: self.response.transaction,
399            raw_transaction: self.response.raw_transaction,
400            effects: self.response.effects,
401            events: self.response.events,
402            balance_changes: self.response.balance_changes,
403            object_changes: self.response.object_changes,
404            // Use full response for any fields that aren't showable
405            ..self.full_response.clone()
406        }
407    }
408}
409
410/// Returns a database URL for testing purposes.
411/// It uses a default user and password, and connects to a local PostgreSQL
412/// instance.
413pub fn db_url(db_name: &str) -> String {
414    format!("postgres://postgres:postgrespw@localhost:5432/{db_name}")
415}
416
417/// Represents a row count result from a SQL query.
418#[derive(QueryableByName, Debug)]
419pub struct RowCount {
420    #[diesel(sql_type = BigInt)]
421    pub cnt: i64,
422}
iota_indexer/test_utils.rs

iota_indexer/
test_utils.rs
