From 4febeb52845e1e9c66b8d3c07f0e173308ca588f Mon Sep 17 00:00:00 2001
From: Pengyuan Bian <pengyuan@clickhouse.com>
Date: Wed, 6 Mar 2024 14:16:09 -0800
Subject: [PATCH] Make it possible to configure max_recv_batch_size in Azure
 service bus.

---
 pubsub/azuresb/azuresb.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/pubsub/azuresb/azuresb.go b/pubsub/azuresb/azuresb.go
index 9876b0b7fe..b6fc751c43 100644
--- a/pubsub/azuresb/azuresb.go
+++ b/pubsub/azuresb/azuresb.go
@@ -61,6 +61,7 @@ import (
 	"net/url"
 	"os"
 	"path"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -233,6 +234,14 @@ func (o *URLOpener) OpenSubscriptionURL(ctx context.Context, u *url.URL) (*pubsu
 		}
 		opts.ListenerTimeout = d
 	}
+	if mrbss := q.Get("max_recv_batch_size"); mrbss != "" {
+		q.Del("max_recv_batch_size")
+		mrbs, err := strconv.Atoi(mrbss)
+		if err != nil {
+			return nil, fmt.Errorf("open subscription %v: invalid max_recv_batch_size %q: %v", u, mrbss, err)
+		}
+		opts.ReceiveBatcherOptions.MaxBatchSize = mrbs
+	}
 	for param := range q {
 		return nil, fmt.Errorf("open subscription %v: invalid query parameter %q", u, param)
 	}