enhance: Add flag to control skip bucket exist check

The bucket existence check need extra permission, this could block online check. This PR enable to skip bucket existence check. Signed-off-by: Congqi Xia <[email protected]>
milvus-io · Sep 6, 2024 · d681b79 · d681b79
1 parent dc550d4
commit d681b79
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 19 deletions.
diff --git a/oss/minio.go b/oss/minio.go
@@ -31,6 +31,23 @@ type MinioClientParam struct {
 
 	BucketName string
 	RootPath   string
+
+	skipCheckBucket bool
+}
+
+// MinioConnectParam is the function type to override client params
+type MinioConnectParam func(p *MinioClientParam)
+
+func WithSkipCheckBucket(v bool) MinioConnectParam {
+	return func(p *MinioClientParam) {
+		p.skipCheckBucket = v
+	}
+}
+
+func WithMinioAddr(addr string) MinioConnectParam {
+	return func(p *MinioClientParam) {
+		p.Addr = addr
+	}
 }
 
 // MinioClient wraps minio client, bucket info within
@@ -80,13 +97,17 @@ func NewMinioClient(ctx context.Context, p MinioClientParam) (*MinioClient, erro
 
 	fmt.Println("Connection successful!")
 
-	ok, err := client.BucketExists(ctx, p.BucketName)
-	if err != nil {
-		fmt.Printf("check bucket %s exists failed: %s\n", p.BucketName, err.Error())
-		return nil, err
-	}
-	if !ok {
-		return nil, errors.Newf("Bucket %s not exists", p.BucketName)
+	if p.skipCheckBucket {
+		fmt.Println("Skip bucket existence check...")
+	} else {
+		ok, err := client.BucketExists(ctx, p.BucketName)
+		if err != nil {
+			fmt.Printf("check bucket %s exists failed: %s\n", p.BucketName, err.Error())
+			return nil, err
+		}
+		if !ok {
+			return nil, errors.Newf("Bucket %s not exists", p.BucketName)
+		}
 	}
 
 	return &MinioClient{

diff --git a/states/binlog.go b/states/binlog.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/oss"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
 	"github.com/milvus-io/birdwatcher/storage"
@@ -26,7 +27,12 @@ type ScanBinlogsParam struct {
 }
 
 func (s *InstanceState) TestScanBinlogsCommand(ctx context.Context, p *ScanBinlogsParam) error {
-	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	params := []oss.MinioConnectParam{}
+	if p.MinioAddress != "" {
+		params = append(params, oss.WithMinioAddr(p.MinioAddress))
+	}
+
+	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, params...)
 	if err != nil {
 		return err
 	}

diff --git a/states/check_partition_key.go b/states/check_partition_key.go
@@ -16,6 +16,7 @@ import (
 
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/oss"
 	"github.com/milvus-io/birdwatcher/proto/v2.0/schemapb"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
@@ -46,7 +47,12 @@ func (s *InstanceState) CheckPartitionKeyCommand(ctx context.Context, p *CheckPa
 	var minioClient *minio.Client
 	var bucketName, rootPath string
 
-	minioClient, bucketName, rootPath, err = s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	params := []oss.MinioConnectParam{}
+	if p.MinioAddress != "" {
+		params = append(params, oss.WithMinioAddr(p.MinioAddress))
+	}
+
+	minioClient, bucketName, rootPath, err = s.GetMinioClientFromCfg(ctx, params...)
 	if err != nil {
 		return err
 	}

diff --git a/states/download_pk.go b/states/download_pk.go
@@ -18,6 +18,7 @@ import (
 
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/oss"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
 )
@@ -46,7 +47,12 @@ func (s *InstanceState) DownloadPKCommand(ctx context.Context, p *DownloadPKPara
 		return err
 	}
 
-	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	params := []oss.MinioConnectParam{}
+	if p.MinioAddress != "" {
+		params = append(params, oss.WithMinioAddr(p.MinioAddress))
+	}
+
+	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, params...)
 	if err != nil {
 		return err
 	}

diff --git a/states/download_segment.go b/states/download_segment.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/oss"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
 )
@@ -33,7 +34,12 @@ func (s *InstanceState) DownloadSegmentCommand(ctx context.Context, p *DownloadS
 		return err
 	}
 
-	minioClient, bucketName, _, err := s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	params := []oss.MinioConnectParam{}
+	if p.MinioAddress != "" {
+		params = append(params, oss.WithMinioAddr(p.MinioAddress))
+	}
+
+	minioClient, bucketName, _, err := s.GetMinioClientFromCfg(ctx, params...)
 	if err != nil {
 		return err
 	}

diff --git a/states/inspect_primary_key.go b/states/inspect_primary_key.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/oss"
 	"github.com/milvus-io/birdwatcher/proto/v2.0/schemapb"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
@@ -117,7 +118,12 @@ func (s *InstanceState) inspectRemote(ctx context.Context, p *InspectPKParam) (m
 		return nil, nil, err
 	}
 
-	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	params := []oss.MinioConnectParam{}
+	if p.MinioAddress != "" {
+		params = append(params, oss.WithMinioAddr(p.MinioAddress))
+	}
+
+	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, params...)
 	if err != nil {
 		fmt.Println("Failed to create folder,", err.Error())
 	}

diff --git a/states/minio.go b/states/minio.go
@@ -32,11 +32,11 @@ type MinioConnectParam struct {
 }
 
 func (s *InstanceState) TestMinioCfgCommand(ctx context.Context, p *TestMinioCfgParam) error {
-	_, _, _, err := s.GetMinioClientFromCfg(ctx, "")
+	_, _, _, err := s.GetMinioClientFromCfg(ctx)
 	return err
 }
 
-func (s *InstanceState) GetMinioClientFromCfg(ctx context.Context, minioAddr string) (client *minio.Client, bucketName, rootPath string, err error) {
+func (s *InstanceState) GetMinioClientFromCfg(ctx context.Context, params ...oss.MinioConnectParam) (client *minio.Client, bucketName, rootPath string, err error) {
 	sessions, err := common.ListSessions(s.client, s.basePath)
 	if err != nil {
 		return nil, "", "", err
@@ -98,9 +98,6 @@ func (s *InstanceState) GetMinioClientFromCfg(ctx context.Context, minioAddr str
 			useSSL = config.GetValue()
 		}
 	}
-	if minioAddr != "" {
-		addr = minioAddr
-	}
 
 	mp := oss.MinioClientParam{
 		CloudProvider: cloudProvider,
@@ -120,6 +117,10 @@ func (s *InstanceState) GetMinioClientFromCfg(ctx context.Context, minioAddr str
 		mp.UseSSL = true
 	}
 
+	for _, param := range params {
+		param(&mp)
+	}
+
 	mClient, err := oss.NewMinioClient(ctx, mp)
 	if err != nil {
 		return nil, "", "", err

diff --git a/states/scan_binlog.go b/states/scan_binlog.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/oss"
 	"github.com/milvus-io/birdwatcher/proto/v2.0/schemapb"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
@@ -25,6 +26,7 @@ type ScanBinlogParams struct {
 	Fields              []string `name:"fields"`
 	Expr                string   `name:"expr"`
 	MinioAddress        string   `name:"minioAddr"`
+	SkipBucketCheck     bool     `name:"skipBucketCheck" default:"false" desc:"skip bucket exist check due to permission issue"`
 	Action              string   `name:"action" default:"count"`
 }
 
@@ -61,9 +63,14 @@ func (s *InstanceState) ScanBinlogCommand(ctx context.Context, p *ScanBinlogPara
 		return err
 	}
 
-	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	params := []oss.MinioConnectParam{oss.WithSkipCheckBucket(p.SkipBucketCheck)}
+	if p.MinioAddress != "" {
+		params = append(params, oss.WithMinioAddr(p.MinioAddress))
+	}
+
+	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, params...)
 	if err != nil {
-		fmt.Println("Failed to create folder,", err.Error())
+		fmt.Println("Failed to create client,", err.Error())
 		return err
 	}