test: add integration tests for message size limit

tests-integration/src/test_util.rs

@@ -39,7 +39,7 @@ use object_store::test_util::TempFolder;
 use object_store::ObjectStore;
 use secrecy::ExposeSecret;
 use servers::grpc::greptime_handler::GreptimeRequestHandler;
-use servers::grpc::GrpcServer;
+use servers::grpc::{GrpcServer, GrpcServerConfig};
 use servers::http::{HttpOptions, HttpServerBuilder};
 use servers::metrics_handler::MetricsHandler;
 use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
@@ -423,13 +423,22 @@ pub async fn setup_grpc_server(
     store_type: StorageType,
     name: &str,
 ) -> (String, TestGuard, Arc<GrpcServer>) {
-    setup_grpc_server_with_user_provider(store_type, name, None).await
+    setup_grpc_server_with(store_type, name, None, None).await
 }
 
 pub async fn setup_grpc_server_with_user_provider(
     store_type: StorageType,
     name: &str,
     user_provider: Option<UserProviderRef>,
+) -> (String, TestGuard, Arc<GrpcServer>) {
+    setup_grpc_server_with(store_type, name, user_provider, None).await
+}
+
+pub async fn setup_grpc_server_with(
+    store_type: StorageType,
+    name: &str,
+    user_provider: Option<UserProviderRef>,
+    grpc_config: Option<GrpcServerConfig>,
 ) -> (String, TestGuard, Arc<GrpcServer>) {
     let instance = setup_standalone_instance(name, store_type).await;
 
@@ -447,8 +456,9 @@ pub async fn setup_grpc_server_with_user_provider(
         user_provider.clone(),
         runtime.clone(),
     ));
+
     let fe_grpc_server = Arc::new(GrpcServer::new(
-        None,
+        grpc_config,
         Some(ServerGrpcQueryHandlerAdaptor::arc(fe_instance_ref.clone())),
         Some(fe_instance_ref.clone()),
         Some(flight_handler),

tests-integration/tests/grpc.rs

@@ -24,10 +24,11 @@ use client::{Client, Database, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_catalog::consts::{MIN_USER_TABLE_ID, MITO_ENGINE};
 use common_query::Output;
 use common_recordbatch::RecordBatches;
+use servers::grpc::GrpcServerConfig;
 use servers::http::prometheus::{PromData, PromSeries, PrometheusJsonResponse, PrometheusResponse};
 use servers::server::Server;
 use tests_integration::test_util::{
-    setup_grpc_server, setup_grpc_server_with_user_provider, StorageType,
+    setup_grpc_server, setup_grpc_server_with, setup_grpc_server_with_user_provider, StorageType,
 };
 
 #[macro_export]
@@ -64,6 +65,9 @@ macro_rules! grpc_tests {
                 test_auto_create_table,
                 test_insert_and_select,
                 test_dbname,
+                test_grpc_message_size_ok,
+                test_grpc_message_size_limit_recv,
+                test_grpc_message_size_limit_send,
                 test_grpc_auth,
                 test_health_check,
                 test_prom_gateway_query,
@@ -115,6 +119,62 @@ pub async fn test_dbname(store_type: StorageType) {
     guard.remove_all().await;
 }
 
+pub async fn test_grpc_message_size_ok(store_type: StorageType) {
+    let config = GrpcServerConfig {
+        max_recv_message_size: 1024,
+        max_send_message_size: 1024,
+    };
+    let (addr, mut guard, fe_grpc_server) =
+        setup_grpc_server_with(store_type, "auto_create_table", None, Some(config)).await;
+
+    let grpc_client = Client::with_urls(vec![addr]);
+    let db = Database::new_with_dbname(
+        format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
+        grpc_client,
+    );
+    db.sql("show tables;").await.unwrap();
+    let _ = fe_grpc_server.shutdown().await;
+    guard.remove_all().await;
+}
+
+pub async fn test_grpc_message_size_limit_send(store_type: StorageType) {
+    let config = GrpcServerConfig {
+        max_recv_message_size: 1024,
+        max_send_message_size: 50,
+    };
+    let (addr, mut guard, fe_grpc_server) =
+        setup_grpc_server_with(store_type, "auto_create_table", None, Some(config)).await;
+
+    let grpc_client = Client::with_urls(vec![addr]);
+    let db = Database::new_with_dbname(
+        format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
+        grpc_client,
+    );
+    let err_msg = db.sql("show tables;").await.unwrap_err().to_string();
+    assert!(err_msg.contains("message length too large"), "{}", err_msg);
+    let _ = fe_grpc_server.shutdown().await;
+    guard.remove_all().await;
+}
+
+pub async fn test_grpc_message_size_limit_recv(store_type: StorageType) {
+    let config = GrpcServerConfig {
+        max_recv_message_size: 10,
+        max_send_message_size: 1024,
+    };
+    let (addr, mut guard, fe_grpc_server) =
+        setup_grpc_server_with(store_type, "auto_create_table", None, Some(config)).await;
+
+    let grpc_client = Client::with_urls(vec![addr]);
+    let db = Database::new_with_dbname(
+        format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
+        grpc_client,
+    );
+    let err_msg = db.sql("show tables;").await.unwrap_err().to_string();
+    assert!(err_msg.contains("message length too large"), "{}", err_msg);
+    let _ = fe_grpc_server.shutdown().await;
+    guard.remove_all().await;
+}
+
 pub async fn test_grpc_auth(store_type: StorageType) {
     let user_provider = user_provider_from_option(
         &"static_user_provider:cmd:greptime_user=greptime_pwd".to_string(),

tests-integration/tests/http.rs

@@ -609,6 +609,8 @@ node_id = 0
 require_lease_before_startup = true
 rpc_addr = "127.0.0.1:3001"
 rpc_runtime_size = 8
+rpc_max_recv_message_size = 536870912
+rpc_max_send_message_size = 536870912
 enable_telemetry = true
 
 [heartbeat]