diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 5696f3051..a05248e90 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -144,12 +144,14 @@ impl Error { /// It's not a supported feature so we should tell the client so. pub(crate) fn is_unsupported(&self) -> bool { + use crate::backend::pool::Error as PoolError; matches!( self, Self::Aggregate(AggregateError::UnsupportedAggregation { .. }) | Self::MultiShard(MultiShardError::Aggregate( AggregateError::UnsupportedAggregation { .. } )) + | Self::Pool(PoolError::PubSubDisabled) ) } diff --git a/pgdog/src/frontend/client/query_engine/pub_sub.rs b/pgdog/src/frontend/client/query_engine/pub_sub.rs index 8cb149160..933e7d28e 100644 --- a/pgdog/src/frontend/client/query_engine/pub_sub.rs +++ b/pgdog/src/frontend/client/query_engine/pub_sub.rs @@ -1,4 +1,6 @@ use super::*; +use crate::backend::Error as BackendError; +use crate::backend::pool::Error as PoolError; use crate::frontend::client::query_engine::fake::FakeResponse; use crate::net::ProtocolMessage; @@ -10,7 +12,10 @@ impl QueryEngine { channel: &str, shard: Shard, ) -> Result<(), Error> { - self.backend.listen(channel, shard).await?; + let result = self.backend.listen(channel, shard).await; + if self.pub_sub_error(context, result).await? { + return Ok(()); + } self.fake_command_response(context, client_messages, &FakeResponse::command("LISTEN")) .await?; @@ -25,13 +30,20 @@ impl QueryEngine { payload: &str, shard: &Shard, ) -> Result<(), Error> { - if context.in_transaction() { + let result = if !context.in_transaction() { + // Send immediately if not in transaction + self.backend.notify(channel, payload, shard.clone()).await + } else if config().pub_sub_enabled() { // Buffer the NOTIFY command if we're in a transaction self.notify_buffer .add(channel.to_string(), payload.to_string(), shard.clone()); + Ok(()) } else { - // Send immediately if not in transaction - self.backend.notify(channel, payload, shard.clone()).await?; + // Fail now: the buffer is flushed after COMMIT has already succeeded. + Err(BackendError::Pool(PoolError::PubSubDisabled)) + }; + if self.pub_sub_error(context, result).await? { + return Ok(()); } self.fake_command_response(context, client_messages, &FakeResponse::command("NOTIFY")) .await?; @@ -58,4 +70,26 @@ impl QueryEngine { } Ok(()) } + + /// Report an unsupported pub/sub request to the client as an `ERROR`, + /// keeping the connection open. Returns `true` if an error was sent. + async fn pub_sub_error( + &mut self, + context: &mut QueryEngineContext<'_>, + result: Result<(), BackendError>, + ) -> Result { + match result { + Ok(()) => Ok(false), + Err(err) if err.is_unsupported() => { + if let Some(transaction) = context.transaction.as_mut() { + transaction.set_error(); + } + let mut error = ErrorResponse::unsupported(&err); + error.hint = Some("set pub_sub_channel_size to enable LISTEN/NOTIFY".into()); + self.send_error_response(context, error).await?; + Ok(true) + } + Err(err) => Err(err.into()), + } + } } diff --git a/pgdog/src/frontend/client/query_engine/test/pub_sub.rs b/pgdog/src/frontend/client/query_engine/test/pub_sub.rs index 7dc72d706..08ffa3f46 100644 --- a/pgdog/src/frontend/client/query_engine/test/pub_sub.rs +++ b/pgdog/src/frontend/client/query_engine/test/pub_sub.rs @@ -1,12 +1,79 @@ use std::io::ErrorKind; +use std::time::Duration; + +use tokio::time::timeout; use crate::{ expect_message, - net::{BindComplete, CommandComplete, Parameters, ParseComplete}, + net::{BindComplete, CommandComplete, ErrorResponse, Parameters, ParseComplete, ReadyForQuery}, }; use super::prelude::*; +#[tokio::test] +async fn pub_sub_disabled_keeps_connection_usable() { + let mut client = SpawnedClient::new_sharded(Parameters::default()).await; + + timeout(Duration::from_secs(5), async { + for extended in [false, true] { + for (begin, query) in [ + (false, "LISTEN pgdog_disabled"), + (false, "NOTIFY pgdog_disabled"), + (true, "LISTEN pgdog_disabled"), + (true, "NOTIFY pgdog_disabled"), + ] { + if begin { + client.send(Query::new("BEGIN")).await; + client.read_until('Z').await; + } + + if extended { + client.send(Parse::new_anonymous(query)).await; + client.send(Bind::new_statement("")).await; + client.send(Execute::new()).await; + client.send(Sync).await; + } else { + client.send(Query::new(query)).await; + } + let error = expect_message!(client.read().await, ErrorResponse); + assert_eq!(error.severity, "ERROR", "{query}"); + assert_eq!(error.code, "0A000", "{query}"); + assert_eq!(error.message, "pub/sub disabled", "{query}"); + assert!( + error + .hint + .as_deref() + .is_some_and(|hint| hint.contains("pub_sub_channel_size")), + "{query}" + ); + assert_eq!( + expect_message!(client.read().await, ReadyForQuery).status, + if begin { 'E' } else { 'I' }, + "{query}" + ); + + if begin { + client.send(Query::new("ROLLBACK")).await; + assert_eq!( + expect_message!(client.read().await, CommandComplete).command(), + "ROLLBACK" + ); + expect_message!(client.read().await, ReadyForQuery); + } + } + } + + client.send(Query::new("SELECT 1")).await; + let messages = client.read_until('Z').await; + assert!( + messages.iter().any(|message| message.code() == 'D'), + "connection still serves queries" + ); + }) + .await + .expect("pub/sub disabled error or subsequent query hung"); +} + #[tokio::test] async fn extended_unlisten_flush_does_not_emit_ready_for_query() { let mut client = TestClient::new_sharded(Parameters::default()).await;