From 0c88629a4fe87e5289c6380d222340723b0ddce2 Mon Sep 17 00:00:00 2001
From: Rudro-25 <rudro@appscode.com>
Date: Fri, 9 Aug 2024 16:50:10 +0600
Subject: [PATCH] update zk client

Signed-off-by: Rudro-25 <rudro@appscode.com>
---
 zookeeper/kubedb_client_builder.go | 36 +++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/zookeeper/kubedb_client_builder.go b/zookeeper/kubedb_client_builder.go
index 1b529aae..f42c51bd 100644
--- a/zookeeper/kubedb_client_builder.go
+++ b/zookeeper/kubedb_client_builder.go
@@ -33,10 +33,13 @@ const (
 )
 
 type KubeDBClientBuilder struct {
-	kc      client.Client
-	db      *dbapi.ZooKeeper
-	podName string
-	url     string
+	kc                client.Client
+	db                *dbapi.ZooKeeper
+	ctx               context.Context
+	podName           string
+	url               string
+	enableHTTPClient  bool
+	disableAMQPClient bool
 }
 
 func NewKubeDBClientBuilder(kc client.Client, db *dbapi.ZooKeeper) *KubeDBClientBuilder {
@@ -46,6 +49,14 @@ func NewKubeDBClientBuilder(kc client.Client, db *dbapi.ZooKeeper) *KubeDBClient
 	}
 }
 
+// NewKubeDBClientBuilderForHTTP returns a KubeDB client builder only for http client
+func NewKubeDBClientBuilderForHTTP(kc client.Client, db *dbapi.ZooKeeper) *KubeDBClientBuilder {
+	return NewKubeDBClientBuilder(kc, db).
+		WithContext(context.TODO()).
+		WithAMQPClientDisabled().
+		WithHTTPClientEnabled()
+}
+
 func (o *KubeDBClientBuilder) WithPod(podName string) *KubeDBClientBuilder {
 	o.podName = podName
 	return o
@@ -56,7 +67,22 @@ func (o *KubeDBClientBuilder) WithURL(url string) *KubeDBClientBuilder {
 	return o
 }
 
-func (o *KubeDBClientBuilder) GetZooKeeperClient(ctx context.Context) (*Client, error) {
+func (o *KubeDBClientBuilder) WithContext(ctx context.Context) *KubeDBClientBuilder {
+	o.ctx = ctx
+	return o
+}
+
+func (o *KubeDBClientBuilder) WithHTTPClientEnabled() *KubeDBClientBuilder {
+	o.enableHTTPClient = true
+	return o
+}
+
+func (o *KubeDBClientBuilder) WithAMQPClientDisabled() *KubeDBClientBuilder {
+	o.disableAMQPClient = true
+	return o
+}
+
+func (o *KubeDBClientBuilder) GetZooKeeperClient() (*Client, error) {
 	var err error
 	if o.podName != "" {
 		o.url = o.getPodURL()