Apply patch to upgrade pg_wire

Author: mamcx

Cargo.lock

@@ -215,7 +215,7 @@ paste = "1.0"
 percent-encoding = "2.3"
 petgraph = { version = "0.6.5", default-features = false }
 pin-project-lite = "0.2.9"
-pgwire = { version = "0.28.0", features = ["server-api"] }
+pgwire = { version = "0.32", features = ["server-api"] }
 postgres-types = "0.2.5"
 pretty_assertions = { version = "1.4", features = ["unstable"] }
 proc-macro2 = "1.0"

Cargo.toml

@@ -518,7 +518,7 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate>(
     // so, unless you are the owner, this will fail.
 
     let (database_identity, db_name) = match &name_or_identity {
-        Some(noa) => match noa.try_resolve(&ctx).await? {
+        Some(noa) => match noa.try_resolve(&ctx).await.map_err(log_and_500)? {
             Ok(resolved) => (resolved, noa.name()),
             Err(name) => {
                 // `name_or_identity` was a `NameOrIdentity::Name`, but no record

crates/client-api/src/routes/database.rs

@@ -97,18 +97,21 @@ impl NameOrIdentity {
     pub async fn try_resolve(
         &self,
         ctx: &(impl ControlStateReadAccess + ?Sized),
-    ) -> axum::response::Result<Result<Identity, &DatabaseName>> {
+    ) -> anyhow::Result<Result<Identity, &DatabaseName>> {
         Ok(match self {
             Self::Identity(identity) => Ok(Identity::from(*identity)),
-            Self::Name(name) => ctx.lookup_identity(name.as_ref()).map_err(log_and_500)?.ok_or(name),
+            Self::Name(name) => ctx.lookup_identity(name.as_ref())?.ok_or(name),
         })
     }
 
     /// A variant of [`Self::try_resolve()`] which maps to a 404 (Not Found)
     /// response if `self` is a [`NameOrIdentity::Name`] for which no
     /// corresponding [`Identity`] is found in the SpacetimeDB DNS.
     pub async fn resolve(&self, ctx: &(impl ControlStateReadAccess + ?Sized)) -> axum::response::Result<Identity> {
-        self.try_resolve(ctx).await?.map_err(|_| StatusCode::NOT_FOUND.into())
+        self.try_resolve(ctx)
+            .await
+            .map_err(log_and_500)?
+            .map_err(|_| StatusCode::NOT_FOUND.into())
     }
 }
 

crates/client-api/src/util.rs

@@ -12,11 +12,10 @@ use pgwire::api::auth::{
     finish_authentication, save_startup_parameters_to_metadata, DefaultServerParameterProvider, LoginInfo,
     StartupHandler,
 };
-use pgwire::api::copy::NoopCopyHandler;
 use pgwire::api::portal::Format;
-use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
+use pgwire::api::query::SimpleQueryHandler;
 use pgwire::api::results::{DataRowEncoder, FieldInfo, QueryResponse, Response, Tag};
-use pgwire::api::{ClientInfo, NoopErrorHandler, METADATA_DATABASE};
+use pgwire::api::{ClientInfo, METADATA_DATABASE};
 use pgwire::api::{PgWireConnectionState, PgWireServerHandlers};
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use pgwire::messages::data::DataRow;
@@ -145,12 +144,12 @@ async fn response<T>(res: axum::response::Result<T>, database: &str) -> Result<T
 }
 
 struct PgSpacetimeDB<T> {
-    ctx: Arc<T>,
+    ctx: T,
     cached: Mutex<Option<Metadata>>,
     parameter_provider: DefaultServerParameterProvider,
 }
 
-impl<T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate> PgSpacetimeDB<T> {
+impl<T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Clone> PgSpacetimeDB<T> {
     async fn exe_sql<'a>(&self, query: String) -> PgWireResult<Vec<Response<'a>>> {
         let params = self.cached.lock().await.clone().unwrap();
         let db = SqlParams {
@@ -283,13 +282,11 @@ impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDel
                 self.cached.lock().await.clone_from(&Some(metadata));
                 finish_authentication(client, &self.parameter_provider).await?;
             }
-            PgWireFrontendMessage::SslRequest(ssl) => {
-                if ssl.is_some() {
-                    let err = PgError::SSLNotSupported;
-                    log::error!("{err}");
-                    let err = ErrorInfo::new("FATAL".to_owned(), "28P01".to_owned(), err.to_string());
-                    return close_client(client, err).await;
-                }
+            PgWireFrontendMessage::SslRequest(_) => {
+                let err = PgError::SSLNotSupported;
+                log::error!("{err}");
+                let err = ErrorInfo::new("FATAL".to_owned(), "28P01".to_owned(), err.to_string());
+                return close_client(client, err).await;
             }
             // The other messages are for features not supported by SpacetimeDB, that are rejected by the parser.
             _ => {
@@ -301,10 +298,10 @@ impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDel
 }
 
 #[async_trait]
-impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate> SimpleQueryHandler
+impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Clone> SimpleQueryHandler
     for PgSpacetimeDB<T>
 {
-    async fn do_query<'a, C>(&self, _client: &mut C, query: &'a str) -> PgWireResult<Vec<Response<'a>>>
+    async fn do_query<'a, C>(&self, _client: &mut C, query: &str) -> PgWireResult<Vec<Response<'a>>>
     where
         C: ClientInfo + Unpin + Send + Sync,
     {
@@ -313,12 +310,12 @@ impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDel
 }
 
 #[derive(Clone)]
-struct PgSpacetimeDBFactory<T> {
+pub struct PgSpacetimeDBFactory<T> {
     handler: Arc<PgSpacetimeDB<T>>,
 }
 
 impl<T> PgSpacetimeDBFactory<T> {
-    pub fn new(ctx: Arc<T>) -> Self {
+    pub fn new(ctx: T) -> Self {
         let mut parameter_provider = DefaultServerParameterProvider::default();
         parameter_provider.server_version = format!("spacetime {}", spacetimedb_lib_version());
 
@@ -333,39 +330,23 @@ impl<T> PgSpacetimeDBFactory<T> {
     }
 }
 
-impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate> PgWireServerHandlers
+impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Clone> PgWireServerHandlers
     for PgSpacetimeDBFactory<T>
 {
-    type StartupHandler = PgSpacetimeDB<T>;
-    type SimpleQueryHandler = PgSpacetimeDB<T>;
-    type ExtendedQueryHandler = PlaceholderExtendedQueryHandler;
-    type CopyHandler = NoopCopyHandler;
-    type ErrorHandler = NoopErrorHandler;
-
-    fn simple_query_handler(&self) -> Arc<Self::SimpleQueryHandler> {
+    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
         self.handler.clone()
     }
 
-    fn extended_query_handler(&self) -> Arc<Self::ExtendedQueryHandler> {
-        Arc::new(PlaceholderExtendedQueryHandler)
-    }
+    // TODO: fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {}
 
-    fn startup_handler(&self) -> Arc<Self::StartupHandler> {
+    fn startup_handler(&self) -> Arc<impl StartupHandler> {
         self.handler.clone()
     }
-
-    fn copy_handler(&self) -> Arc<Self::CopyHandler> {
-        Arc::new(NoopCopyHandler)
-    }
-
-    fn error_handler(&self) -> Arc<Self::ErrorHandler> {
-        Arc::new(NoopErrorHandler)
-    }
 }
 
-pub async fn start_pg<T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + 'static>(
+pub async fn start_pg<T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Clone + 'static>(
     shutdown: Arc<Notify>,
-    ctx: Arc<T>,
+    ctx: T,
     tcp: TcpListener,
 ) {
     let factory = Arc::new(PgSpacetimeDBFactory::new(ctx));