Add test for max_backend_message_size

Author: petuhovskiy

tokio-postgres/src/config.rs

@@ -599,6 +599,14 @@ impl Config {
                     self.replication_mode(mode);
                 }
             }
+            "max_backend_message_size" => {
+                let limit = value.parse::<usize>().map_err(|_| {
+                    Error::config_parse(Box::new(InvalidValue("max_backend_message_size")))
+                })?;
+                if limit > 0 {
+                    self.max_backend_message_size(limit);
+                }
+            }
             key => {
                 return Err(Error::config_parse(Box::new(UnknownOption(
                     key.to_string(),

tokio-postgres/src/connect_raw.rs

@@ -90,9 +90,12 @@ where
     let stream = connect_tls(stream, config.ssl_mode, tls).await?;
 
     let mut stream = StartupStream {
-        inner: Framed::new(stream, PostgresCodec {
-            max_message_size: config.max_backend_message_size,
-        }),
+        inner: Framed::new(
+            stream,
+            PostgresCodec {
+                max_message_size: config.max_backend_message_size,
+            },
+        ),
         buf: BackendMessages::empty(),
         delayed: VecDeque::new(),
     };

tokio-postgres/tests/test/main.rs

@@ -280,6 +280,30 @@ async fn query_raw_txt() {
     assert!(rows[0].body_len() > 0);
 }
 
+#[tokio::test]
+async fn limit_max_backend_message_size() {
+    let client = connect("user=postgres max_backend_message_size=10000").await;
+    let small: Vec<tokio_postgres::Row> = client
+        .query_raw_txt("SELECT REPEAT('a', 20)", [])
+        .await
+        .unwrap()
+        .try_collect()
+        .await
+        .unwrap();
+
+    assert_eq!(small.len(), 1);
+    assert_eq!(small[0].as_text(0).unwrap().unwrap().len(), 20);
+
+    let large: Result<Vec<tokio_postgres::Row>, Error> = client
+        .query_raw_txt("SELECT REPEAT('a', 2000000)", [])
+        .await
+        .unwrap()
+        .try_collect()
+        .await;
+
+    assert!(large.is_err());
+}
+
 #[tokio::test]
 async fn command_tag() {
     let client = connect("user=postgres").await;