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, DbUrl, 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        store
158            .execute_in_blocking_worker(move |this| {
159                db_init_hook(&this);
160                Ok(())
161            })
162            .await
163            .expect("failed to run the DB init hook");
164    }
165
166    let registry = prometheus_filtered::Registry::default();
167    init_metrics(&registry);
168    let indexer_metrics = IndexerMetrics::new(&registry);
169
170    let handle = match reader_writer_config {
171        IndexerTypeConfig::Reader {
172            reader_mode_rpc_url,
173        } => {
174            let config = crate::config::JsonRpcConfig {
175                iota_names_options: IotaNamesOptions::default(),
176                historic_fallback_options: Default::default(),
177                rpc_address: reader_mode_rpc_url.parse().unwrap(),
178                rpc_client_url: rpc_url,
179            };
180            let pool = store.blocking_cp();
181            let store_clone = store.clone();
182            tokio::spawn(async move {
183                Indexer::start_reader(
184                    &config,
185                    store_clone,
186                    &registry,
187                    pool,
188                    indexer_metrics,
189                    CancellationToken::new(),
190                )
191                .await
192            })
193        }
194        IndexerTypeConfig::Writer {
195            retention_config,
196            pruning_delay_ms,
197            pruning_batch_size,
198        } => {
199            let fullnode_rpc_url = rpc_url.parse::<Url>().unwrap();
200            let store_clone = store.clone();
201            let mut ingestion_config = IngestionConfig::default();
202            ingestion_config.sources.remote_store_url =
203                data_ingestion_path.is_none().then_some(fullnode_rpc_url);
204            ingestion_config.sources.data_ingestion_path = data_ingestion_path;
205
206            tokio::spawn(async move {
207                Indexer::start_writer_with_config(
208                    &ingestion_config,
209                    store_clone,
210                    indexer_metrics,
211                    retention_config,
212                    pruning_delay_ms,
213                    pruning_batch_size,
214                    cancel,
215                )
216                .await
217            })
218        }
219        IndexerTypeConfig::AnalyticalWorker => {
220            let store = PgIndexerAnalyticalStore::new(store.blocking_cp());
221
222            tokio::spawn(
223                async move { Indexer::start_analytical_worker(store, indexer_metrics).await },
224            )
225        }
226    };
227
228    (store, handle)
229}
230
231/// Manage a test database for integration tests.
232pub struct TestDatabase {
233    pub url: String,
234    db_name: String,
235    connection: PoolConnection,
236    pool_config: ConnectionPoolConfig,
237}
238
239impl TestDatabase {
240    pub fn new(db_url: String) -> Self {
241        // Reduce the connection pool size to 5 for testing
242        // to prevent maxing out
243        let pool_config = ConnectionPoolConfig {
244            pool_size: 5,
245            ..Default::default()
246        };
247
248        let db_name = db_url.split('/').next_back().unwrap().into();
249        let (default_url, _) = replace_db_name(&db_url, "postgres");
250        let blocking_pool = new_connection_pool(&DbUrl::from(default_url), &pool_config).unwrap();
251        let connection = blocking_pool.get().unwrap();
252        Self {
253            url: db_url,
254            db_name,
255            connection,
256            pool_config,
257        }
258    }
259
260    /// Drop the database in the server if it exists.
261    pub fn drop_if_exists(&mut self) {
262        self.connection
263            .batch_execute(&format!("DROP DATABASE IF EXISTS {}", self.db_name))
264            .unwrap();
265    }
266
267    /// Create the database in the server.
268    pub fn create(&mut self) {
269        self.connection
270            .batch_execute(&format!("CREATE DATABASE {}", self.db_name))
271            .unwrap();
272    }
273
274    /// Drop and recreate the database in the server.
275    pub fn recreate(&mut self) {
276        self.drop_if_exists();
277        self.create();
278    }
279
280    /// Create a new connection pool to the database.
281    pub fn to_connection_pool(&self) -> ConnectionPool {
282        new_connection_pool(&DbUrl::from(self.url.as_str()), &self.pool_config).unwrap()
283    }
284
285    pub fn reset_db(&mut self) {
286        crate::db::reset_database(&mut self.to_connection_pool().get().unwrap()).unwrap();
287    }
288}
289
290pub fn create_pg_store(db_url: &str, reset_database: bool) -> PgIndexerStore {
291    let registry = prometheus_filtered::Registry::default();
292    init_metrics(&registry);
293    let indexer_metrics = IndexerMetrics::new(&registry);
294
295    let mut test_db = TestDatabase::new(db_url.to_string());
296    if reset_database {
297        test_db.recreate();
298    }
299
300    PgIndexerStore::new(test_db.to_connection_pool(), indexer_metrics)
301}
302
303fn replace_db_name(db_url: &str, new_db_name: &str) -> (String, String) {
304    let pos = db_url.rfind('/').expect("unable to find / in db_url");
305    let old_db_name = &db_url[pos + 1..];
306
307    (
308        format!("{}/{}", &db_url[..pos], new_db_name),
309        old_db_name.to_string(),
310    )
311}
312
313pub async fn force_delete_database(db_url: String) {
314    // Replace the database name with the default `postgres`, which should be the
315    // last string after `/` This is necessary because you can't drop a database
316    // while being connected to it. Hence switch to the default `postgres`
317    // database to drop the active database.
318    let (default_db_url, db_name) = replace_db_name(&db_url, "postgres");
319    let mut pool_config = ConnectionPoolConfig::default();
320    pool_config.set_pool_size(1);
321
322    let blocking_pool = new_connection_pool(&DbUrl::from(default_db_url), &pool_config).unwrap();
323    blocking_pool
324        .get()
325        .unwrap()
326        .batch_execute(&format!("DROP DATABASE IF EXISTS {db_name} WITH (FORCE)"))
327        .unwrap();
328}
329
330#[derive(Clone)]
331pub struct IotaTransactionBlockResponseBuilder<'a> {
332    response: IotaTransactionBlockResponse,
333    full_response: &'a IotaTransactionBlockResponse,
334}
335
336impl<'a> IotaTransactionBlockResponseBuilder<'a> {
337    pub fn new(full_response: &'a IotaTransactionBlockResponse) -> Self {
338        Self {
339            response: IotaTransactionBlockResponse::default(),
340            full_response,
341        }
342    }
343
344    pub fn with_input(mut self) -> Self {
345        self.response = IotaTransactionBlockResponse {
346            transaction: self.full_response.transaction.clone(),
347            ..self.response
348        };
349        self
350    }
351
352    pub fn with_raw_input(mut self) -> Self {
353        self.response = IotaTransactionBlockResponse {
354            raw_transaction: self.full_response.raw_transaction.clone(),
355            ..self.response
356        };
357        self
358    }
359
360    pub fn with_effects(mut self) -> Self {
361        self.response = IotaTransactionBlockResponse {
362            effects: self.full_response.effects.clone(),
363            ..self.response
364        };
365        self
366    }
367
368    pub fn with_events(mut self) -> Self {
369        self.response = IotaTransactionBlockResponse {
370            events: self.full_response.events.clone(),
371            ..self.response
372        };
373        self
374    }
375
376    pub fn with_balance_changes(mut self) -> Self {
377        self.response = IotaTransactionBlockResponse {
378            balance_changes: self.full_response.balance_changes.clone(),
379            ..self.response
380        };
381        self
382    }
383
384    pub fn with_object_changes(mut self) -> Self {
385        self.response = IotaTransactionBlockResponse {
386            object_changes: self.full_response.object_changes.clone(),
387            ..self.response
388        };
389        self
390    }
391
392    pub fn with_input_and_changes(mut self) -> Self {
393        self.response = IotaTransactionBlockResponse {
394            transaction: self.full_response.transaction.clone(),
395            balance_changes: self.full_response.balance_changes.clone(),
396            object_changes: self.full_response.object_changes.clone(),
397            ..self.response
398        };
399        self
400    }
401
402    pub fn build(self) -> IotaTransactionBlockResponse {
403        IotaTransactionBlockResponse {
404            transaction: self.response.transaction,
405            raw_transaction: self.response.raw_transaction,
406            effects: self.response.effects,
407            events: self.response.events,
408            balance_changes: self.response.balance_changes,
409            object_changes: self.response.object_changes,
410            // Use full response for any fields that aren't showable
411            ..self.full_response.clone()
412        }
413    }
414}
415
416/// Returns a database URL for testing purposes.
417/// It uses a default user and password, and connects to a local PostgreSQL
418/// instance.
419pub fn db_url(db_name: &str) -> String {
420    format!("postgres://postgres:postgrespw@localhost:5432/{db_name}")
421}
422
423/// Represents a row count result from a SQL query.
424#[derive(QueryableByName, Debug)]
425pub struct RowCount {
426    #[diesel(sql_type = BigInt)]
427    pub cnt: i64,
428}
iota_indexer/test_utils.rs

iota_indexer/
test_utils.rs
