sui_indexer_alt_reader/
pg_reader.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use anyhow::Context;
7use anyhow::bail;
8use async_graphql::dataloader::DataLoader;
9use diesel::deserialize::FromSqlRow;
10use diesel::expression::QueryMetadata;
11use diesel::pg::Pg;
12use diesel::query_builder::Query;
13use diesel::query_builder::QueryFragment;
14use diesel::query_builder::QueryId;
15use diesel::query_dsl::CompatibleType;
16use diesel::query_dsl::methods::LimitDsl;
17use diesel_async::RunQueryDsl;
18use prometheus::Registry;
19use sui_indexer_alt_metrics::db::DbConnectionStatsCollector;
20use tracing::debug;
21use tracing::warn;
22use url::Url;
23
24use crate::metrics::DbReaderMetrics;
25
26pub use sui_pg_db as db;
27
28/// This wrapper type exists to perform error conversion between the data fetching layer and the
29/// RPC layer, metrics collection, and debug logging of database queries.
30#[derive(Clone)]
31pub struct PgReader {
32    db: Option<db::Db>,
33    metrics: Arc<DbReaderMetrics>,
34}
35
36pub struct Connection<'p> {
37    conn: db::Connection<'p>,
38    metrics: Arc<DbReaderMetrics>,
39}
40
41impl PgReader {
42    /// Create a new database reader. If `database_url` is `None`, the reader will not accept any
43    /// connection requests (they will all fail).
44    ///
45    /// `prefix` is used to prefix the metrics collected by this reader.
46    pub async fn new(
47        prefix: Option<&str>,
48        database_url: Option<Url>,
49        db_args: db::DbArgs,
50        registry: &Registry,
51    ) -> anyhow::Result<Self> {
52        let db = if let Some(database_url) = database_url {
53            let db = db::Db::for_read(database_url, db_args)
54                .await
55                .context("Failed to create database for reading")?
56                .register_metrics(prefix, registry)?;
57
58            registry
59                .register(Box::new(DbConnectionStatsCollector::new(
60                    prefix,
61                    db.clone(),
62                )))
63                .context("Failed to register database connection stats collector")?;
64
65            Some(db)
66        } else {
67            None
68        };
69
70        let metrics = DbReaderMetrics::new(prefix, registry);
71
72        Ok(Self { db, metrics })
73    }
74
75    /// Create a data loader backed by this reader.
76    pub fn as_data_loader(&self) -> DataLoader<Self> {
77        DataLoader::new(self.clone(), tokio::spawn)
78    }
79
80    /// Check if this reader has a database available.
81    pub fn has_database(&self) -> bool {
82        self.db.is_some()
83    }
84
85    /// Acquire a connection to the database. This can fail if a database has not been configured
86    /// to connect to.
87    pub async fn connect(&self) -> anyhow::Result<Connection<'_>> {
88        let Some(db) = &self.db else {
89            bail!("No database to connect to");
90        };
91
92        let conn = db
93            .connect()
94            .await
95            .context("Failed to connect to database")?;
96
97        Ok(Connection {
98            conn,
99            metrics: self.metrics.clone(),
100        })
101    }
102}
103
104impl Connection<'_> {
105    pub async fn first<'q, Q, ST, U>(&mut self, query: Q) -> anyhow::Result<U>
106    where
107        Q: LimitDsl,
108        Q::Output: Query + QueryFragment<Pg> + QueryId + Send + 'q,
109        <Q::Output as Query>::SqlType: CompatibleType<U, Pg, SqlType = ST>,
110        U: Send + FromSqlRow<ST, Pg> + 'static,
111        Pg: QueryMetadata<<Q::Output as Query>::SqlType>,
112        ST: 'static,
113    {
114        let query = query.limit(1);
115
116        self.metrics.requests_received.inc();
117        let _guard = self.metrics.latency.start_timer();
118
119        let pid = self.conn.pid;
120        let query_debug = diesel::debug_query(&query).to_string();
121        match query.get_result(&mut self.conn).await {
122            Ok(results) => {
123                self.metrics.requests_succeeded.inc();
124                debug!(pid, "{query_debug}");
125                Ok(results)
126            }
127            Err(err) => {
128                self.metrics.requests_failed.inc();
129                if is_timeout(&err) {
130                    warn!(pid, "Timed out: {query_debug}");
131                } else {
132                    warn!(pid, "Failed with '{err:?}': {query_debug}");
133                };
134                Err(err).with_context(|| format!("First error from DB request pid={pid}"))
135            }
136        }
137    }
138
139    pub async fn results<'q, Q, ST, U>(&mut self, query: Q) -> anyhow::Result<Vec<U>>
140    where
141        Q: Query + QueryFragment<Pg> + QueryId + Send + 'q,
142        Q::SqlType: CompatibleType<U, Pg, SqlType = ST>,
143        U: Send + FromSqlRow<ST, Pg> + 'static,
144        Pg: QueryMetadata<Q::SqlType>,
145        ST: 'static,
146    {
147        self.metrics.requests_received.inc();
148        let _guard = self.metrics.latency.start_timer();
149
150        let pid = self.conn.pid;
151        let query_debug = diesel::debug_query(&query).to_string();
152        match query.get_results(&mut self.conn).await {
153            Ok(results) => {
154                self.metrics.requests_succeeded.inc();
155                debug!(pid, "{query_debug}");
156                Ok(results)
157            }
158            Err(err) => {
159                self.metrics.requests_failed.inc();
160                if is_timeout(&err) {
161                    warn!(pid, "Timed out: {query_debug}");
162                } else {
163                    warn!(pid, "Failed with '{err:?}': {query_debug}");
164                };
165                Err(err).with_context(|| format!("Results error from DB request pid={pid}"))
166            }
167        }
168    }
169}
170
171/// Detect whether the error is due to a timeout.
172fn is_timeout(err: &diesel::result::Error) -> bool {
173    let diesel::result::Error::DatabaseError(_, info) = err else {
174        return false;
175    };
176
177    info.message() == "canceling statement due to statement timeout"
178}
sui_indexer_alt_reader/pg_reader.rs

sui_indexer_alt_reader/
pg_reader.rs
