mirror of
https://github.com/neondatabase/neon.git
synced 2026-05-19 22:20:37 +00:00
cd1d2d1996
Fixes https://github.com/neondatabase/cloud/issues/20973. This refactors `connect_raw` in order to return direct access to the delayed notices. I cannot find a way to test this with psycopg2 unfortunately, although testing it with psql does return the expected results.
63 lines
1.9 KiB
Rust
63 lines
1.9 KiB
Rust
use crate::config::SslMode;
|
|
use crate::tls::TlsConnect;
|
|
|
|
use crate::{cancel_query, client::SocketConfig, tls::MakeTlsConnect};
|
|
use crate::{cancel_query_raw, Error};
|
|
use tokio::io::{AsyncRead, AsyncWrite};
|
|
use tokio::net::TcpStream;
|
|
|
|
/// The capability to request cancellation of in-progress queries on a
|
|
/// connection.
|
|
#[derive(Clone)]
|
|
pub struct CancelToken {
|
|
pub socket_config: Option<SocketConfig>,
|
|
pub ssl_mode: SslMode,
|
|
pub process_id: i32,
|
|
pub secret_key: i32,
|
|
}
|
|
|
|
impl CancelToken {
|
|
/// Attempts to cancel the in-progress query on the connection associated
|
|
/// with this `CancelToken`.
|
|
///
|
|
/// The server provides no information about whether a cancellation attempt was successful or not. An error will
|
|
/// only be returned if the client was unable to connect to the database.
|
|
///
|
|
/// Cancellation is inherently racy. There is no guarantee that the
|
|
/// cancellation request will reach the server before the query terminates
|
|
/// normally, or that the connection associated with this token is still
|
|
/// active.
|
|
///
|
|
/// Requires the `runtime` Cargo feature (enabled by default).
|
|
pub async fn cancel_query<T>(&self, tls: T) -> Result<(), Error>
|
|
where
|
|
T: MakeTlsConnect<TcpStream>,
|
|
{
|
|
cancel_query::cancel_query(
|
|
self.socket_config.clone(),
|
|
self.ssl_mode,
|
|
tls,
|
|
self.process_id,
|
|
self.secret_key,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Like `cancel_query`, but uses a stream which is already connected to the server rather than opening a new
|
|
/// connection itself.
|
|
pub async fn cancel_query_raw<S, T>(&self, stream: S, tls: T) -> Result<(), Error>
|
|
where
|
|
S: AsyncRead + AsyncWrite + Unpin,
|
|
T: TlsConnect<S>,
|
|
{
|
|
cancel_query_raw::cancel_query_raw(
|
|
stream,
|
|
self.ssl_mode,
|
|
tls,
|
|
self.process_id,
|
|
self.secret_key,
|
|
)
|
|
.await
|
|
}
|
|
}
|