Client-side support to storage class and tablespace settings

open-gpdb · Aug 24, 2024 · 546d37f · 546d37f
1 parent 2837dd9
commit 546d37f
Show file tree

Hide file tree

Showing 10 changed files with 61 additions and 30 deletions.
diff --git a/cmd/client/main.go b/cmd/client/main.go
@@ -12,19 +12,28 @@ import (
 	"github.com/yezzey-gp/yproxy/pkg/message"
 	"github.com/yezzey-gp/yproxy/pkg/object"
 	"github.com/yezzey-gp/yproxy/pkg/proc"
+	"github.com/yezzey-gp/yproxy/pkg/tablespace"
 	"github.com/yezzey-gp/yproxy/pkg/ylogger"
 )
 
-var cfgPath string
-var oldCfgPath string
-var logLevel string
-var decrypt bool
-var encrypt bool
-var offset uint64
-var segmentPort int
-var segmentNum int
-var confirm bool
-var garbage bool
+var (
+	cfgPath    string
+	oldCfgPath string
+	logLevel   string
+
+	decrypt bool
+	/* Put command flags */
+	encrypt      bool
+	storageClass string
+	tableSpace   string
+
+	offset uint64
+
+	segmentPort int
+	segmentNum  int
+	confirm     bool
+	garbage     bool
+)
 
 // TODOV
 func Runner(f func(net.Conn, *config.Instance, []string) error) func(*cobra.Command, []string) error {
@@ -96,7 +105,16 @@ func putFunc(con net.Conn, instanceCnf *config.Instance, args []string) error {
 	ycl := client.NewYClient(con)
 	r := proc.NewProtoReader(ycl)
 
-	msg := message.NewPutMessage(args[0], encrypt).Encode()
+	msg := message.NewPutMessageV2(args[0], encrypt, []message.PutSettings{
+		{
+			Name:  message.StorageClassSetting,
+			Value: storageClass,
+		},
+		{
+			Name:  message.TableSpaceSetting,
+			Value: tableSpace,
+		},
+	}).Encode()
 	_, err := con.Write(msg)
 	if err != nil {
 		return err
@@ -341,6 +359,8 @@ func init() {
 	rootCmd.AddCommand(copyCmd)
 
 	putCmd.PersistentFlags().BoolVarP(&encrypt, "encrypt", "e", false, "encrypt external object before put")
+	putCmd.PersistentFlags().StringVarP(&storageClass, "storage-class", "s", "STANDARD", "storage class for message upload")
+	putCmd.PersistentFlags().StringVarP(&tableSpace, "tablespace", "t", tablespace.DefaultTableSpace, "storage class for message upload")
 	rootCmd.AddCommand(putCmd)
 
 	rootCmd.AddCommand(listCmd)

diff --git a/pkg/message/message_test.go b/pkg/message/message_test.go
@@ -81,15 +81,15 @@ func TestPutV2Msg(t *testing.T) {
 		name     string
 		encrypt  bool
 		err      error
-		settings []message.PutSetting
+		settings []message.PutSettings
 	}
 
 	for _, tt := range []tcase{
 		{
 			"nam1",
 			true,
 			nil,
-			[]message.PutSetting{
+			[]message.PutSettings{
 				{
 					Name:  "a",
 					Value: "b",

diff --git a/pkg/message/put_message_v2.go b/pkg/message/put_message_v2.go
@@ -5,7 +5,10 @@ import (
 	"encoding/binary"
 )
 
-type PutSetting struct {
+const StorageClassSetting = "StorageClass"
+const TableSpaceSetting = "TableSpace"
+
+type PutSettings struct {
 	Name  string
 	Value string
 }
@@ -14,12 +17,12 @@ type PutMessageV2 struct {
 	Encrypt bool
 	Name    string
 
-	Settings []PutSetting
+	Settings []PutSettings
 }
 
 var _ ProtoMessage = &PutMessageV2{}
 
-func NewPutMessageV2(name string, encrypt bool, settings []PutSetting) *PutMessageV2 {
+func NewPutMessageV2(name string, encrypt bool, settings []PutSettings) *PutMessageV2 {
 	return &PutMessageV2{
 		Name:     name,
 		Encrypt:  encrypt,
@@ -90,7 +93,7 @@ func (c *PutMessageV2) Decode(body []byte) {
 
 	totalOff := 4 + off + 8
 
-	c.Settings = make([]PutSetting, settLen)
+	c.Settings = make([]PutSettings, settLen)
 
 	for i := 0; i < int(settLen); i++ {
 

diff --git a/pkg/mock/storage.go b/pkg/mock/storage.go
diff --git a/pkg/proc/interaction.go b/pkg/proc/interaction.go
@@ -21,7 +21,7 @@ func ProcessPutExtended(
 	s storage.StorageInteractor,
 	pr *ProtoReader,
 	name string,
-	encrypt bool, settings []message.PutSetting, cr crypt.Crypter, ycl client.YproxyClient) error {
+	encrypt bool, settings []message.PutSettings, cr crypt.Crypter, ycl client.YproxyClient) error {
 
 	ycl.SetExternalFilePath(name)
 
@@ -51,7 +51,11 @@ func ProcessPutExtended(
 				return
 			}
 		} else {
-			ylogger.Zero.Debug().Str("path", name).Msg("omit encryption for chunk")
+			ylogger.Zero.Debug().Str("path", name).Msg("omit encryption for upload chunks")
+		}
+
+		for _, set := range settings {
+			ylogger.Zero.Debug().Str("setting name", set.Name).Str("value", set.Value).Msg("setting for chunk")
 		}
 
 		defer w.Close()

diff --git a/pkg/storage/filestorage.go b/pkg/storage/filestorage.go
@@ -51,7 +51,7 @@ func (s *FileStorageInteractor) ListPath(prefix string) ([]*object.ObjectInfo, e
 	return data, err
 }
 
-func (s *FileStorageInteractor) PutFileToDest(name string, r io.Reader, _ []message.PutSetting) error {
+func (s *FileStorageInteractor) PutFileToDest(name string, r io.Reader, _ []message.PutSettings) error {
 	fPath := path.Join(s.cnf.StoragePrefix, name)
 	fDir := path.Dir(fPath)
 	os.MkdirAll(fDir, 0700)

diff --git a/pkg/storage/s3storage.go b/pkg/storage/s3storage.go
@@ -49,7 +49,7 @@ func (s *S3StorageInteractor) CatFileFromStorage(name string, offset int64) (io.
 	return object.Body, err
 }
 
-func (s *S3StorageInteractor) PutFileToDest(name string, r io.Reader, settings []message.PutSetting) error {
+func (s *S3StorageInteractor) PutFileToDest(name string, r io.Reader, settings []message.PutSettings) error {
 	sess, err := s.pool.GetSession(context.TODO())
 	if err != nil {
 		ylogger.Zero.Err(err).Msg("failed to acquire s3 session")
@@ -63,8 +63,8 @@ func (s *S3StorageInteractor) PutFileToDest(name string, r io.Reader, settings [
 		uploader.Concurrency = 1
 	})
 
-	storageClass := ResolveStorageSetting(settings, "StorageClass", "STANDARD")
-	tableSpace := ResolveStorageSetting(settings, "Tablespace", tablespace.DefaultTableSpace)
+	storageClass := ResolveStorageSetting(settings, message.StorageClassSetting, "STANDARD")
+	tableSpace := ResolveStorageSetting(settings, message.TableSpaceSetting, tablespace.DefaultTableSpace)
 
 	bucket, ok := s.bucketMap[tableSpace]
 	if !ok {

diff --git a/pkg/storage/storage.go b/pkg/storage/storage.go
@@ -15,7 +15,7 @@ type StorageReader interface {
 }
 
 type StorageWriter interface {
-	PutFileToDest(name string, r io.Reader, settings []message.PutSetting) error
+	PutFileToDest(name string, r io.Reader, settings []message.PutSettings) error
 	PatchFile(name string, r io.ReadSeeker, startOffset int64) error
 }
 
@@ -55,6 +55,10 @@ func NewStorage(cnf *config.Storage) (StorageInteractor, error) {
 
 func buildBucketMapFromCnf(cnf *config.Storage) map[string]string {
 	mp := cnf.TablespaceMap
+	if mp == nil {
+		/* fallback for backward-compatibilty if to TableSpace map configured */
+		mp = map[string]string{}
+	}
 	if _, ok := mp[tablespace.DefaultTableSpace]; !ok {
 		mp[tablespace.DefaultTableSpace] = cnf.StorageBucket
 	}

diff --git a/pkg/storage/utils.go b/pkg/storage/utils.go
@@ -2,7 +2,7 @@ package storage
 
 import "github.com/yezzey-gp/yproxy/pkg/message"
 
-func ResolveStorageSetting(settings []message.PutSetting, name, defaultVal string) string {
+func ResolveStorageSetting(settings []message.PutSettings, name, defaultVal string) string {
 
 	for _, s := range settings {
 		if s.Name == name {

diff --git a/pkg/storage/utils_test.go b/pkg/storage/utils_test.go
@@ -17,7 +17,7 @@ func TestResolveSettings(t *testing.T) {
 		name     string
 		defaultV string
 		exp      string
-		settings []message.PutSetting
+		settings []message.PutSettings
 	}
 
 	for _, tt := range []tcase{
@@ -31,7 +31,7 @@ func TestResolveSettings(t *testing.T) {
 			"ababa",
 			"aboba",
 			"aboba",
-			[]message.PutSetting{
+			[]message.PutSettings{
 				{
 					Name:  "djewikdeowp",
 					Value: "jdoiwejoidew",
@@ -43,7 +43,7 @@ func TestResolveSettings(t *testing.T) {
 			"ababa",
 			"aboba",
 			"valval",
-			[]message.PutSetting{
+			[]message.PutSettings{
 				{
 					Name:  "ababa",
 					Value: "valval",