From ef95b88e3777f0f454f98050cf9f42c32579ca1e Mon Sep 17 00:00:00 2001
From: raihankhan <raihan@appscode.com>
Date: Tue, 27 Aug 2024 18:52:52 +0600
Subject: [PATCH] Add virtual host setup

Signed-off-by: raihankhan <raihan@appscode.com>
---
 rabbitmq/kubedb_client_builder.go | 48 ++++++++++++++++++-------------
 1 file changed, 28 insertions(+), 20 deletions(-)

diff --git a/rabbitmq/kubedb_client_builder.go b/rabbitmq/kubedb_client_builder.go
index a3f2768f..d4a8c91a 100644
--- a/rabbitmq/kubedb_client_builder.go
+++ b/rabbitmq/kubedb_client_builder.go
@@ -19,12 +19,12 @@ package rabbitmq
 import (
 	"context"
 	"errors"
+	amqp "github.com/rabbitmq/amqp091-go"
 
 	"fmt"
 	"strings"
 
 	rmqhttp "github.com/michaelklishin/rabbit-hole/v2"
-	amqp "github.com/rabbitmq/amqp091-go"
 	core "k8s.io/api/core/v1"
 	kerr "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/types"
@@ -131,14 +131,34 @@ func (o *KubeDBClientBuilder) GetRabbitMQClient() (*Client, error) {
 	}
 
 	rmqClient := &Client{}
+	defaultVhost := "/"
 
-	if !o.disableAMQPClient {
-		if o.amqpURL == "" {
-			o.amqpURL = o.GetAMQPconnURL(username, password)
+	if o.enableHTTPClient {
+		if o.httpURL == "" {
+			o.httpURL = o.GetHTTPconnURL()
+		}
+		httpClient, err := rmqhttp.NewClient(o.httpURL, username, password)
+		if err != nil {
+			klog.Error(err, "Failed to get http client for rabbitmq")
+			return nil, err
 		}
 
-		if o.vhost == "" {
-			o.vhost = o.GetVirtualHostFromURL(o.amqpURL)
+		vhosts, err := httpClient.ListVhosts()
+		for _, vhost := range vhosts {
+			if vhost.Description == "Default virtual host" {
+				defaultVhost = vhost.Name
+				break
+			}
+		}
+		rmqClient.HTTPClient = HTTPClient{httpClient}
+	}
+
+	if !o.disableAMQPClient {
+		if o.amqpURL == "" {
+			if o.vhost == "" {
+				o.vhost = defaultVhost
+			}
+			o.amqpURL = o.GetAMQPconnURL(username, password, o.vhost)
 		}
 
 		rabbitConnection, err := amqp.DialConfig(o.amqpURL, amqp.Config{
@@ -153,23 +173,11 @@ func (o *KubeDBClientBuilder) GetRabbitMQClient() (*Client, error) {
 		rmqClient.AMQPClient = AMQPClient{rabbitConnection}
 	}
 
-	if o.enableHTTPClient {
-		if o.httpURL == "" {
-			o.httpURL = o.GetHTTPconnURL()
-		}
-		httpClient, err := rmqhttp.NewClient(o.httpURL, username, password)
-		if err != nil {
-			klog.Error(err, "Failed to get http client for rabbitmq")
-			return nil, err
-		}
-		rmqClient.HTTPClient = HTTPClient{httpClient}
-	}
-
 	return rmqClient, nil
 }
 
-func (o *KubeDBClientBuilder) GetAMQPconnURL(username string, password string) string {
-	return fmt.Sprintf("amqp://%s:%s@%s.%s.svc.cluster.local:%d/", username, password, o.db.ServiceName(), o.db.Namespace, kubedb.RabbitMQAMQPPort)
+func (o *KubeDBClientBuilder) GetAMQPconnURL(username string, password string, vhost string) string {
+	return fmt.Sprintf("amqp://%s:%s@%s.%s.svc.cluster.local:%d%s", username, password, o.db.ServiceName(), o.db.Namespace, kubedb.RabbitMQAMQPPort, vhost)
 }
 
 func (o *KubeDBClientBuilder) GetHTTPconnURL() string {