update zk client

Signed-off-by: Rudro-25 <[email protected]>
kubedb · Sep 3, 2024 · f22d8b5 · f22d8b5
1 parent b8c278d
commit f22d8b5
Showing 1 changed file with 31 additions and 5 deletions.
diff --git 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()