1use std::fmt::Write as _;
4use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::Duration;
7
8use canton_core::{Error, Result};
9use canton_daml::Contract as ContractType;
10use tokio::sync::RwLock;
11use tokio_postgres::types::ToSql;
12
13use crate::query::{Param, Predicate, Query, Sql};
14use crate::row::{Contract, Exercise};
15
16#[derive(Clone, Debug)]
22pub struct PqsClient {
23 inner: Arc<Inner>,
24}
25
26pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
30
31struct Inner {
32 config: tokio_postgres::Config,
33 tls: Tls,
34 timeout_ms: AtomicU64,
38 client: RwLock<Arc<tokio_postgres::Client>>,
41}
42
43#[allow(clippy::missing_fields_in_debug)]
46impl std::fmt::Debug for Inner {
47 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48 f.debug_struct("PqsClient")
49 .field("hosts", &self.config.get_hosts())
50 .field("dbname", &self.config.get_dbname())
51 .field("user", &self.config.get_user())
52 .field("tls", &self.tls.is_on())
53 .field("timeout_ms", &self.timeout_ms.load(Ordering::Relaxed))
54 .finish()
55 }
56}
57
58#[derive(Clone)]
60enum Tls {
61 Off,
62 #[cfg(feature = "tls")]
63 Rustls(Arc<rustls::ClientConfig>),
64}
65
66impl Tls {
67 fn is_on(&self) -> bool {
68 match self {
69 Tls::Off => false,
70 #[cfg(feature = "tls")]
71 Tls::Rustls(_) => true,
72 }
73 }
74}
75
76fn parse_config(config: &str) -> Result<tokio_postgres::Config> {
79 config.parse::<tokio_postgres::Config>().map_err(|e| {
80 Error::InvalidRequest(format!(
81 "PQS connection string is not one tokio-postgres accepts: {}",
82 canton_core::chain(&e)
83 ))
84 })
85}
86
87impl PqsClient {
88 pub async fn connect(config: &str) -> Result<Self> {
100 Self::connect_with_timeout(config, DEFAULT_TIMEOUT).await
101 }
102
103 pub async fn connect_with_timeout(config: &str, timeout: Duration) -> Result<Self> {
110 let config = parse_config(config)?;
111 Self::establish(config, Tls::Off, timeout).await
112 }
113
114 #[cfg(feature = "tls")]
122 #[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
123 pub async fn connect_tls(config: &str) -> Result<Self> {
124 Self::connect_tls_with_timeout(config, DEFAULT_TIMEOUT).await
125 }
126
127 #[cfg(feature = "tls")]
140 #[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
141 pub async fn connect_tls_with_timeout(config: &str, timeout: Duration) -> Result<Self> {
142 let mut config = parse_config(config)?;
143 config.ssl_mode(tokio_postgres::config::SslMode::Require);
144
145 let mut roots = rustls::RootCertStore::empty();
146 let native = rustls_native_certs::load_native_certs();
147 if native.certs.is_empty() {
148 return Err(Error::InvalidRequest(
149 "no native root certificates were found, so a TLS connection could not be \
150 verified"
151 .to_string(),
152 ));
153 }
154 roots.add_parsable_certificates(native.certs);
155 let client_config = rustls::ClientConfig::builder()
156 .with_root_certificates(roots)
157 .with_no_client_auth();
158 Self::establish(config, Tls::Rustls(Arc::new(client_config)), timeout).await
159 }
160
161 #[must_use]
164 pub fn with_timeout(self, timeout: Duration) -> Self {
165 self.inner.timeout_ms.store(
166 u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX),
167 Ordering::Relaxed,
168 );
169 self
170 }
171
172 fn timeout(&self) -> Duration {
173 Duration::from_millis(self.inner.timeout_ms.load(Ordering::Relaxed))
174 }
175
176 async fn establish(
177 config: tokio_postgres::Config,
178 tls: Tls,
179 timeout: Duration,
180 ) -> Result<Self> {
181 let client = open(&config, &tls, timeout).await?;
182 Ok(Self {
183 inner: Arc::new(Inner {
184 config,
185 tls,
186 timeout_ms: AtomicU64::new(u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX)),
187 client: RwLock::new(Arc::new(client)),
188 }),
189 })
190 }
191
192 async fn reconnect(&self) -> Result<()> {
198 let mut guard = self.inner.client.write().await;
199 if !guard.is_closed() {
200 return Ok(());
201 }
202 let fresh = open(&self.inner.config, &self.inner.tls, self.timeout()).await?;
203 *guard = Arc::new(fresh);
204 tracing::info!("reconnected to PQS after the store closed the connection");
205 Ok(())
206 }
207
208 pub async fn run<T>(&self, query: &Query<T>) -> Result<Vec<Contract<T>>>
214 where
215 T: ContractType + serde::de::DeserializeOwned,
216 {
217 let rows = self.query(&query.compile()).await?;
218 rows.iter().map(Contract::from_row).collect()
219 }
220
221 pub async fn active<T>(&self) -> Result<Vec<Contract<T>>>
226 where
227 T: ContractType + serde::de::DeserializeOwned,
228 {
229 self.run(&Query::<T>::active()).await
230 }
231
232 pub async fn lookup<T>(&self, contract_id: &str) -> Result<Option<Contract<T>>>
240 where
241 T: ContractType + serde::de::DeserializeOwned,
242 {
243 let sql = Sql::raw(
244 "SELECT * FROM lookup_contract($1, $2)",
245 vec![
246 Param::Text(contract_id.to_string()),
247 Param::Text(Query::<T>::qname()),
248 ],
249 );
250 let rows = self.query(&sql).await?;
251 rows.first().map(Contract::from_row).transpose()
252 }
253
254 pub async fn exercises<T, C>(
261 &self,
262 from: Option<i64>,
263 to: Option<i64>,
264 ) -> Result<Vec<Exercise<C>>>
265 where
266 T: ContractType,
267 C: canton_daml::Choice<T> + serde::de::DeserializeOwned,
268 {
269 let qname = format!("{}:{}", Query::<T>::qname(), C::NAME);
271 let rows = self.query(&exercises_sql(qname, from, to)).await?;
272 rows.iter().map(Exercise::from_row).collect()
273 }
274
275 pub async fn latest_offset(&self) -> Result<i64> {
283 let rows = self
284 .query(&Sql::raw("SELECT latest_offset()", Vec::new()))
285 .await?;
286 let row = rows.first().ok_or_else(|| {
287 Error::UnexpectedResponse("PQS answered latest_offset() with no row".to_string())
288 })?;
289 row.try_get::<_, Option<i64>>(0)
290 .map_err(|e| {
291 Error::UnexpectedResponse(format!(
292 "PQS returned a latest offset this cannot read: {e}"
293 ))
294 })?
295 .ok_or_else(|| {
296 Error::UnexpectedResponse(
297 "PQS reported no latest offset, so it has ingested nothing".to_string(),
298 )
299 })
300 }
301
302 pub async fn query(&self, sql: &Sql) -> Result<Vec<tokio_postgres::Row>> {
307 let values: Vec<Box<dyn ToSql + Sync + Send>> = sql
308 .params
309 .iter()
310 .map(|param| -> Box<dyn ToSql + Sync + Send> {
311 match param {
312 Param::Json(value) => Box::new(value.clone()),
313 Param::Text(value) | Param::Numeric(value) => Box::new(value.clone()),
318 Param::Path(segments) => Box::new(segments.clone()),
319 Param::Offset(value) => Box::new(*value),
320 }
321 })
322 .collect();
323 let borrowed: Vec<&(dyn ToSql + Sync)> = values
324 .iter()
325 .map(|value| value.as_ref() as &(dyn ToSql + Sync))
326 .collect();
327
328 tracing::debug!(sql = %sql.text, params = sql.params.len(), "PQS query");
329 let client = self.inner.client.read().await.clone();
330 match self.run_on(&client, sql, &borrowed).await {
331 Err(Attempt::Closed) => {
332 self.reconnect().await?;
333 let client = self.inner.client.read().await.clone();
334 self.run_on(&client, sql, &borrowed)
335 .await
336 .map_err(Attempt::into_error)
337 }
338 other => other.map_err(Attempt::into_error),
339 }
340 }
341
342 async fn run_on(
343 &self,
344 client: &tokio_postgres::Client,
345 sql: &Sql,
346 params: &[&(dyn ToSql + Sync)],
347 ) -> std::result::Result<Vec<tokio_postgres::Row>, Attempt> {
348 match tokio::time::timeout(self.timeout(), client.query(sql.text.as_str(), params)).await {
349 Err(_) => Err(Attempt::Failed(Error::Timeout)),
350 Ok(Ok(rows)) => Ok(rows),
351 Ok(Err(e)) if e.is_closed() => Err(Attempt::Closed),
352 Ok(Err(e)) => Err(Attempt::Failed(classify(&e))),
353 }
354 }
355}
356
357#[must_use]
363pub fn active_signed_by<T: ContractType>(party: &str) -> Query<T> {
364 Query::<T>::active().filter(Predicate::signatory(party))
365}
366
367enum Attempt {
396 Closed,
397 Failed(Error),
398}
399
400impl Attempt {
401 fn into_error(self) -> Error {
402 match self {
403 Attempt::Closed => Error::Connection(
404 "the PQS connection closed, and closed again on a fresh one".to_string(),
405 ),
406 Attempt::Failed(error) => error,
407 }
408 }
409}
410
411async fn open(
413 config: &tokio_postgres::Config,
414 tls: &Tls,
415 timeout: Duration,
416) -> Result<tokio_postgres::Client> {
417 match tls {
418 Tls::Off => {
419 let (client, connection) =
420 tokio::time::timeout(timeout, config.connect(tokio_postgres::NoTls))
421 .await
422 .map_err(|_| Error::Timeout)?
423 .map_err(|e| connect_error("cannot reach PQS", &e))?;
424 Ok(drive(client, connection))
425 }
426 #[cfg(feature = "tls")]
427 Tls::Rustls(client_config) => {
428 let connector =
429 tokio_postgres_rustls::MakeRustlsConnect::new(client_config.as_ref().clone());
430 let (client, connection) = tokio::time::timeout(timeout, config.connect(connector))
431 .await
432 .map_err(|_| Error::Timeout)?
433 .map_err(|e| connect_error("cannot reach PQS over TLS", &e))?;
434 Ok(drive(client, connection))
435 }
436 }
437}
438
439fn drive<S, T>(
440 client: tokio_postgres::Client,
441 connection: tokio_postgres::Connection<S, T>,
442) -> tokio_postgres::Client
443where
444 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
447 T: tokio_postgres::tls::TlsStream + Unpin + Send + 'static,
448{
449 tokio::spawn(async move {
450 if let Err(e) = connection.await {
451 tracing::warn!(error = %e, "the PQS connection ended");
454 }
455 });
456 client
457}
458
459fn classify(error: &tokio_postgres::Error) -> Error {
460 let Some(code) = error.code() else {
461 if is_permanent_kind(error) {
467 return Error::InvalidRequest(format!(
468 "PQS query could not be made: {}",
469 detail(error)
470 ));
471 }
472 return Error::Connection(format!("PQS query failed: {}", detail(error)));
473 };
474 if is_transient_sqlstate(code.code()) {
475 Error::Connection(format!(
479 "PQS is momentarily unavailable ({}): {}",
480 code.code(),
481 detail(error)
482 ))
483 } else {
484 Error::InvalidRequest(format!("PQS rejected the query: {}", detail(error)))
485 }
486}
487
488fn exercises_sql(qname: String, from: Option<i64>, to: Option<i64>) -> Sql {
506 let mut params = vec![Param::Text(qname)];
507 let mut text = "SELECT * FROM exercises($1".to_string();
508 if from.is_some() || to.is_some() {
509 let from_sql = from.map_or_else(
510 || "COALESCE(pruned_offset(), oldest_offset())".to_string(),
511 |value| {
512 params.push(Param::Offset(value));
513 format!("${}", params.len())
514 },
515 );
516 let to_sql = to.map_or_else(
517 || "latest_offset()".to_string(),
518 |value| {
519 params.push(Param::Offset(value));
520 format!("${}", params.len())
521 },
522 );
523 let _ = write!(text, ", {from_sql}, {to_sql}");
524 }
525 text.push(')');
526 Sql::raw(text, params)
527}
528
529fn connect_error(what: &str, error: &tokio_postgres::Error) -> Error {
530 let detail = detail(error);
531 if let Some(code) = error.code()
532 && !is_transient_sqlstate(code.code())
533 {
534 return Error::InvalidRequest(format!(
535 "{what}: the store refused the connection ({}): {detail}",
536 code.code()
537 ));
538 }
539 if is_permanent_kind(error) {
540 return Error::InvalidRequest(format!("{what}: {detail}"));
541 }
542 if detail.to_ascii_lowercase().contains("certificate") {
545 return Error::InvalidRequest(format!(
546 "{what}: its certificate could not be verified: {detail}"
547 ));
548 }
549 Error::Connection(format!("{what}: {detail}"))
550}
551
552fn is_permanent_kind(error: &tokio_postgres::Error) -> bool {
564 use std::error::Error as _;
565 const PERMANENT: [&str; 11] = [
566 "invalid connection string",
567 "invalid configuration",
568 "authentication error",
569 "error performing TLS handshake",
570 "error encoding message to server",
571 "error serializing parameter",
572 "error deserializing column",
573 "error parsing response from server",
574 "unexpected message from server",
575 "query returned an unexpected number of columns",
576 "query returned an unexpected number of rows",
577 ];
578 let kind = error.to_string();
579 PERMANENT.iter().any(|prefix| kind.starts_with(prefix))
580 || error
581 .source()
582 .is_some_and(<dyn std::error::Error>::is::<tokio_postgres::types::WrongType>)
583}
584
585fn is_transient_sqlstate(sqlstate: &str) -> bool {
586 matches!(sqlstate.get(..2), Some("08" | "40" | "53" | "55" | "57"))
587}
588
589fn detail(error: &tokio_postgres::Error) -> String {
592 canton_core::chain(error)
593}
594
595#[cfg(test)]
596mod tests {
597 use super::is_transient_sqlstate;
598
599 #[test]
605 fn a_server_that_is_momentarily_unable_is_retriable() {
606 for (code, what) in [
607 ("57P01", "the admin shut the server down — a failover"),
608 ("57P02", "the server crashed"),
609 ("57P03", "the server cannot connect yet — still starting"),
610 ("40001", "serialization failure under concurrent load"),
611 ("40P01", "deadlock detected"),
612 ("53300", "too many connections"),
613 ("53200", "out of memory"),
614 ("55P03", "lock not available"),
615 ("08006", "connection failure"),
616 ] {
617 assert!(
618 is_transient_sqlstate(code),
619 "{code} ({what}) must be retriable"
620 );
621 }
622 }
623
624 #[test]
628 fn a_statement_the_server_refused_is_the_callers_to_fix() {
629 for (code, what) in [
630 ("42601", "syntax error"),
631 ("42883", "no such function — a wrong qname"),
632 ("42703", "no such column"),
633 ("22P02", "invalid text representation — a bad cast"),
634 ("23505", "unique violation"),
635 ] {
636 assert!(
637 !is_transient_sqlstate(code),
638 "{code} ({what}) must not be retriable"
639 );
640 }
641 }
642
643 #[test]
646 fn a_malformed_sqlstate_is_not_mistaken_for_a_transient_one() {
647 for code in ["", "5", "4"] {
648 assert!(!is_transient_sqlstate(code), "{code:?} is not a class");
649 }
650 }
651}
652
653#[cfg(test)]
654#[allow(clippy::unwrap_used, clippy::expect_used)]
655mod connection_tests {
656 use super::*;
657
658 #[tokio::test]
659 async fn a_connection_string_that_does_not_parse_is_the_callers_mistake() {
660 let error = PqsClient::connect("this is not a connection string")
661 .await
662 .expect_err("refused");
663 assert!(matches!(error, Error::InvalidRequest(_)), "{error:?}");
664 assert!(!error.is_retriable(), "a typo is not transient: {error}");
665 }
666
667 #[tokio::test]
668 async fn a_store_that_never_answers_is_a_timeout_not_a_hang() {
669 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
670 let port = listener.local_addr().unwrap().port();
671 tokio::spawn(async move {
673 let _held = listener.accept().await;
674 tokio::time::sleep(Duration::from_secs(30)).await;
675 });
676 let error = PqsClient::connect_with_timeout(
677 &format!("host=127.0.0.1 port={port} user=u dbname=d"),
678 Duration::from_millis(300),
679 )
680 .await
681 .expect_err("timed out");
682 assert!(matches!(error, Error::Timeout), "{error:?}");
683 assert!(error.is_retriable());
684 }
685
686 #[cfg(feature = "tls")]
689 #[tokio::test]
690 async fn a_store_that_refuses_tls_is_refused_back_not_downgraded() {
691 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
692
693 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
694 let port = listener.local_addr().unwrap().port();
695 tokio::spawn(async move {
696 let (mut socket, _) = listener.accept().await.unwrap();
697 let mut request = [0u8; 8];
699 let _ = socket.read_exact(&mut request).await;
700 let _ = socket.write_all(b"N").await;
701 let _ = socket.shutdown().await;
702 });
703 let error = PqsClient::connect_tls_with_timeout(
704 &format!("host=127.0.0.1 port={port} user=u dbname=d"),
705 Duration::from_secs(5),
706 )
707 .await
708 .expect_err("no TLS means no connection");
709 assert!(
710 !error.is_retriable(),
711 "a store without TLS will not grow one on retry: {error}"
712 );
713 assert!(
714 format!("{error}").to_ascii_lowercase().contains("tls"),
715 "the reason names TLS: {error}"
716 );
717 }
718
719 #[test]
723 fn the_exercises_statement_binds_only_the_bounds_it_was_given() {
724 let q = || "pkg:Mod:T:Choice".to_string();
725
726 let none = exercises_sql(q(), None, None);
727 assert_eq!(none.text, "SELECT * FROM exercises($1)");
728 assert_eq!(none.params, vec![Param::Text(q())]);
729
730 let from = exercises_sql(q(), Some(10), None);
731 assert_eq!(
732 from.text,
733 "SELECT * FROM exercises($1, $2, latest_offset())"
734 );
735 assert_eq!(from.params, vec![Param::Text(q()), Param::Offset(10)]);
736
737 let to = exercises_sql(q(), None, Some(20));
738 assert_eq!(
739 to.text,
740 "SELECT * FROM exercises($1, COALESCE(pruned_offset(), oldest_offset()), $2)"
741 );
742 assert_eq!(to.params, vec![Param::Text(q()), Param::Offset(20)]);
743
744 let both = exercises_sql(q(), Some(10), Some(20));
745 assert_eq!(both.text, "SELECT * FROM exercises($1, $2, $3)");
746 assert_eq!(
747 both.params,
748 vec![Param::Text(q()), Param::Offset(10), Param::Offset(20)]
749 );
750 }
751}