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/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;