github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/unstructured/viper_wrap.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package unstructured
    21  
    22  import (
    23  	"bytes"
    24  	"os"
    25  	"path/filepath"
    26  	"strings"
    27  
    28  	"github.com/spf13/cast"
    29  	oviper "github.com/spf13/viper"
    30  
    31  	appsv1alpha1 "github.com/1aal/kubeblocks/apis/apps/v1alpha1"
    32  )
    33  
    34  type viperWrap struct {
    35  	*oviper.Viper
    36  
    37  	name   string
    38  	format appsv1alpha1.CfgFileFormat
    39  }
    40  
    41  func init() {
    42  	CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.Ini, createViper(appsv1alpha1.Ini))
    43  	// CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.YAML, createViper(appsv1alpha1.YAML))
    44  	CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.JSON, createViper(appsv1alpha1.JSON))
    45  	CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.Dotenv, createViper(appsv1alpha1.Dotenv))
    46  	CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.HCL, createViper(appsv1alpha1.HCL))
    47  	CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.TOML, createViper(appsv1alpha1.TOML))
    48  	CfgObjectRegistry().RegisterConfigCreator(appsv1alpha1.Properties, createViper(appsv1alpha1.Properties))
    49  }
    50  
    51  func (v *viperWrap) GetString(key string) (string, error) {
    52  	return cast.ToStringE(v.Get(key))
    53  }
    54  
    55  func (v *viperWrap) GetAllParameters() map[string]interface{} {
    56  	return v.AllSettings()
    57  }
    58  
    59  func (v *viperWrap) RemoveKey(key string) error {
    60  	// TODO viper does not support remove key
    61  	return nil
    62  }
    63  
    64  func (v *viperWrap) SubConfig(key string) ConfigObject {
    65  	return &viperWrap{
    66  		Viper:  v.Sub(key),
    67  		format: v.format,
    68  	}
    69  }
    70  
    71  func (v *viperWrap) Update(key string, value any) error {
    72  	v.Set(key, value)
    73  	return nil
    74  }
    75  
    76  func (v *viperWrap) Marshal() (string, error) {
    77  	const tmpFileName = "_config_tmp"
    78  
    79  	tmpDir, err := os.MkdirTemp(os.TempDir(), "configuration-")
    80  	if err != nil {
    81  		return "", err
    82  	}
    83  	defer os.RemoveAll(tmpDir)
    84  
    85  	cfgName := v.name
    86  	if cfgName == "" {
    87  		cfgName = tmpFileName
    88  	}
    89  	tmpFile := filepath.Join(tmpDir, strings.ReplaceAll(cfgName, ".", "_"))
    90  	return dumpCfgContent(v.Viper, tmpFile)
    91  }
    92  
    93  func (v viperWrap) Unmarshal(str string) error {
    94  	return v.ReadConfig(bytes.NewReader([]byte(str)))
    95  }
    96  
    97  func newCfgViper(cfgType appsv1alpha1.CfgFileFormat) *oviper.Viper {
    98  	defaultKeySep := DelimiterDot
    99  	if cfgType == appsv1alpha1.Properties || cfgType == appsv1alpha1.Dotenv {
   100  		defaultKeySep = CfgDelimiterPlaceholder
   101  	}
   102  	v := oviper.NewWithOptions(oviper.KeyDelimiter(defaultKeySep))
   103  	v.SetConfigType(strings.ToLower(string(cfgType)))
   104  	return v
   105  }
   106  
   107  func createViper(format appsv1alpha1.CfgFileFormat) ConfigObjectCreator {
   108  	return func(name string) ConfigObject {
   109  		return &viperWrap{
   110  			name:   name,
   111  			format: format,
   112  			Viper:  newCfgViper(format),
   113  		}
   114  	}
   115  }