kalzu-value-store/storage/storage.go

package storage

import (
	"fmt"
	"time"

	badger "github.com/dgraph-io/badger/v4"
	"github.com/sirupsen/logrus"

	"kvs/types"
)

// StorageService handles all BadgerDB operations and data management
type StorageService struct {
	db             *badger.DB
	config         *types.Config
	compressionSvc *CompressionService
	logger         *logrus.Logger
}

// NewStorageService creates a new storage service
func NewStorageService(db *badger.DB, config *types.Config, logger *logrus.Logger) (*StorageService, error) {
	var compressionSvc *CompressionService
	var err error

	// Initialize compression if enabled
	if config.CompressionEnabled {
		compressionSvc, err = NewCompressionService()
		if err != nil {
			return nil, fmt.Errorf("failed to initialize compression: %v", err)
		}
	}

	return &StorageService{
		db:             db,
		config:         config,
		compressionSvc: compressionSvc,
		logger:         logger,
	}, nil
}

// Close closes the storage service and its resources
func (s *StorageService) Close() {
	if s.compressionSvc != nil {
		s.compressionSvc.Close()
	}
}

// StoreWithTTL stores data with optional TTL and compression
func (s *StorageService) StoreWithTTL(txn *badger.Txn, key []byte, data []byte, ttl time.Duration) error {
	var finalData []byte
	var err error

	// Compress data if compression is enabled
	if s.config.CompressionEnabled && s.compressionSvc != nil {
		finalData, err = s.compressionSvc.CompressData(data)
		if err != nil {
			return fmt.Errorf("failed to compress data: %v", err)
		}
	} else {
		finalData = data
	}

	entry := badger.NewEntry(key, finalData)

	// Apply TTL if specified
	if ttl > 0 {
		entry = entry.WithTTL(ttl)
	}

	return txn.SetEntry(entry)
}

// RetrieveWithDecompression retrieves and decompresses data from BadgerDB
func (s *StorageService) RetrieveWithDecompression(txn *badger.Txn, key []byte) ([]byte, error) {
	item, err := txn.Get(key)
	if err != nil {
		return nil, err
	}

	var compressedData []byte
	err = item.Value(func(val []byte) error {
		compressedData = append(compressedData, val...)
		return nil
	})
	if err != nil {
		return nil, err
	}

	// Decompress data if compression is enabled
	if s.config.CompressionEnabled && s.compressionSvc != nil {
		return s.compressionSvc.DecompressData(compressedData)
	}

	return compressedData, nil
}

// CompressData compresses data using the compression service
func (s *StorageService) CompressData(data []byte) ([]byte, error) {
	if !s.config.CompressionEnabled || s.compressionSvc == nil {
		return data, nil
	}
	return s.compressionSvc.CompressData(data)
}

// DecompressData decompresses data using the compression service
func (s *StorageService) DecompressData(compressedData []byte) ([]byte, error) {
	if !s.config.CompressionEnabled || s.compressionSvc == nil {
		return compressedData, nil
	}
	return s.compressionSvc.DecompressData(compressedData)
}