github.com/khulnasoft-lab/tunnel-db@v0.0.0-20231117205118-74e1113bd007/pkg/db/data_source.go (about)

     1  package db
     2  
     3  import (
     4  	"encoding/json"
     5  
     6  	bolt "go.etcd.io/bbolt"
     7  	"golang.org/x/xerrors"
     8  
     9  	"github.com/khulnasoft-lab/tunnel-db/pkg/types"
    10  )
    11  
    12  const (
    13  	dataSourceBucket = "data-source"
    14  )
    15  
    16  func (dbc Config) PutDataSource(tx *bolt.Tx, bktName string, source types.DataSource) error {
    17  	bucket, err := tx.CreateBucketIfNotExists([]byte(dataSourceBucket))
    18  	if err != nil {
    19  		return xerrors.Errorf("failed to create %s bucket: %w", dataSourceBucket, err)
    20  	}
    21  	b, err := json.Marshal(source)
    22  	if err != nil {
    23  		return xerrors.Errorf("JSON marshal error: %w", err)
    24  	}
    25  
    26  	return bucket.Put([]byte(bktName), b)
    27  }
    28  
    29  func (dbc Config) getDataSource(tx *bolt.Tx, bktName string) (types.DataSource, error) {
    30  	bucket := tx.Bucket([]byte(dataSourceBucket))
    31  	if bucket == nil {
    32  		return types.DataSource{}, nil
    33  	}
    34  
    35  	b := bucket.Get([]byte(bktName))
    36  	if b == nil {
    37  		return types.DataSource{}, nil
    38  	}
    39  
    40  	var source types.DataSource
    41  	if err := json.Unmarshal(b, &source); err != nil {
    42  		return types.DataSource{}, xerrors.Errorf("JSON unmarshal error: %w", err)
    43  	}
    44  
    45  	return source, nil
    46  }