Line data    Source code

       1              : use serde::{Deserialize, Serialize};
       2              : use tokio::io::{AsyncRead, AsyncWrite};
       3              : use tokio::net::TcpStream;
       4              : 
       5              : use crate::client::SocketConfig;
       6              : use crate::config::SslMode;
       7              : use crate::tls::{MakeTlsConnect, TlsConnect};
       8              : use crate::{Error, cancel_query, cancel_query_raw};
       9              : 
      10              : /// A cancellation token that allows easy cancellation of a query.
      11              : #[derive(Clone)]
      12              : pub struct CancelToken {
      13              :     pub socket_config: SocketConfig,
      14              :     pub raw: RawCancelToken,
      15              : }
      16              : 
      17              : /// A raw cancellation token that allows cancellation of a query, given a fresh connection to postgres.
      18            0 : #[derive(Debug, Clone, Serialize, Deserialize)]
      19              : pub struct RawCancelToken {
      20              :     pub ssl_mode: SslMode,
      21              :     pub process_id: i32,
      22              :     pub secret_key: i32,
      23              : }
      24              : 
      25              : impl CancelToken {
      26              :     /// Attempts to cancel the in-progress query on the connection associated
      27              :     /// with this `CancelToken`.
      28              :     ///
      29              :     /// The server provides no information about whether a cancellation attempt was successful or not. An error will
      30              :     /// only be returned if the client was unable to connect to the database.
      31              :     ///
      32              :     /// Cancellation is inherently racy. There is no guarantee that the
      33              :     /// cancellation request will reach the server before the query terminates
      34              :     /// normally, or that the connection associated with this token is still
      35              :     /// active.
      36              :     ///
      37              :     /// Requires the `runtime` Cargo feature (enabled by default).
      38            0 :     pub async fn cancel_query<T>(&self, tls: T) -> Result<(), Error>
      39            0 :     where
      40            0 :         T: MakeTlsConnect<TcpStream>,
      41            0 :     {
      42            0 :         cancel_query::cancel_query(
      43            0 :             self.socket_config.clone(),
      44            0 :             self.raw.ssl_mode,
      45            0 :             tls,
      46            0 :             self.raw.process_id,
      47            0 :             self.raw.secret_key,
      48            0 :         )
      49            0 :         .await
      50            0 :     }
      51              : }
      52              : 
      53              : impl RawCancelToken {
      54              :     /// Like `cancel_query`, but uses a stream which is already connected to the server rather than opening a new
      55              :     /// connection itself.
      56            0 :     pub async fn cancel_query_raw<S, T>(&self, stream: S, tls: T) -> Result<(), Error>
      57            0 :     where
      58            0 :         S: AsyncRead + AsyncWrite + Unpin,
      59            0 :         T: TlsConnect<S>,
      60            0 :     {
      61            0 :         cancel_query_raw::cancel_query_raw(
      62            0 :             stream,
      63            0 :             self.ssl_mode,
      64            0 :             tls,
      65            0 :             self.process_id,
      66            0 :             self.secret_key,
      67            0 :         )
      68            0 :         .await
      69            0 :     }
      70              : }