memory.go

package dstore

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/url"
	"path"
	"sort"
	"strings"
	"sync"
	"time"
)

// MemoryStore is a store that keeps all data in memory. Useful for unit testing where a store is required.
// Work in progress, not all methods are implemented
type MemoryStore struct {
	*commonStore

	baseURL *url.URL

	data     map[string][]byte
	metadata map[string]map[string]string
	modified map[string]time.Time

	lock sync.RWMutex
}

func (m *MemoryStore) OpenObject(ctx context.Context, name string) (out io.ReadCloser, err error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	if _, ok := m.data[name]; !ok {
		return nil, ErrNotFound
	}

	reader := io.NopCloser(bytes.NewReader(m.data[name]))
	out, err = m.uncompressedReader(ctx, reader)
	return
}

func (m *MemoryStore) WriteObject(ctx context.Context, base string, f io.Reader, metadataKeyValues ...string) (err error) {
	m.lock.Lock()
	defer m.lock.Unlock()

	// Parse metadataKeyValues array
	if len(metadataKeyValues)%2 != 0 {
		return fmt.Errorf("metadataKeyValues must have an even number of strings (key-value pairs), got %d", len(metadataKeyValues))
	}

	if _, exists := m.data[base]; !m.overwrite && exists {
		return nil
	}

	w := bytes.NewBuffer(nil)
	if err := m.compressedCopy(ctx, w, f); err != nil {
		return err
	}

	m.data[base] = w.Bytes()
	m.modified[base] = time.Now()

	// Set metadata if provided
	if len(metadataKeyValues) > 0 {
		if m.metadata == nil {
			m.metadata = make(map[string]map[string]string)
		}
		metadata := make(map[string]string)
		for i := 0; i < len(metadataKeyValues); i += 2 {
			key := metadataKeyValues[i]
			value := metadataKeyValues[i+1]
			metadata[key] = value
		}
		m.metadata[base] = metadata
	}

	return nil
}

func (m *MemoryStore) FileExists(_ context.Context, base string) (bool, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	_, exists := m.data[base]
	return exists, nil
}

func (m *MemoryStore) ObjectPath(name string) string {
	return path.Join(strings.TrimLeft(m.baseURL.Path, "/"), m.pathWithExt(name))
}

func (m *MemoryStore) ObjectURL(name string) string {
	return objectURL(m.baseURL, m.pathWithExt(name))
}

func (m *MemoryStore) ObjectAttributes(_ context.Context, base string) (*ObjectAttributes, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	if !m.modified[base].IsZero() {
		return &ObjectAttributes{
			LastModified: m.modified[base],
			Size:         int64(len(m.data[base])),
			Metadata:     m.metadata[base],
		}, nil
	}

	return nil, ErrNotFound
}

func (m *MemoryStore) SetMetadata(ctx context.Context, base string, metadata map[string]string) error {
	m.lock.Lock()
	defer m.lock.Unlock()

	if _, exists := m.data[base]; !exists {
		return ErrNotFound
	}

	if m.metadata == nil {
		m.metadata = make(map[string]map[string]string)
	}

	// Create a copy of the metadata to avoid external modifications
	metadataCopy := make(map[string]string)
	for k, v := range metadata {
		metadataCopy[k] = v
	}

	m.metadata[base] = metadataCopy
	return nil
}

func (m *MemoryStore) PushLocalFile(ctx context.Context, localFile, toBaseName string) (err error) {
	remove, err := pushLocalFile(ctx, m, localFile, toBaseName)
	if err != nil {
		return err
	}
	return remove()
}

func (m *MemoryStore) CopyObject(_ context.Context, src, dest string) error {
	m.lock.Lock()
	defer m.lock.Unlock()

	if _, ok := m.data[src]; !ok {
		return ErrNotFound
	}

	m.data[dest] = m.data[src]
	return nil
}

func (m *MemoryStore) WalkFrom(_ context.Context, prefix, startingPoint string, f func(filename string) (err error)) error {
	panic("not yet supported for this store type")
}

func (m *MemoryStore) WalkFromTo(_ context.Context, prefix, startingPoint, exclusiveEndPoint string, f func(filename string) (err error)) error {
	panic("not yet supported for this store type")
}

func (m *MemoryStore) Walk(ctx context.Context, prefix string, f func(filename string) (err error)) error {
	panic("not yet supported for this store type")
}

func (m *MemoryStore) ListFiles(ctx context.Context, prefix string, max int) ([]string, error) {
	panic("not yet supported for this store type")
}

func (m *MemoryStore) DeleteObject(ctx context.Context, base string) error {
	m.lock.Lock()
	defer m.lock.Unlock()

	delete(m.data, base)
	delete(m.modified, base)
	return nil
}

func (m *MemoryStore) BaseURL() *url.URL {
	return &url.URL{}
}

func (m *MemoryStore) SubStore(subFolder string) (Store, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	newFiles := map[string][]byte{}
	newModified := map[string]time.Time{}

	for k, v := range m.data {
		if !strings.HasPrefix(k, subFolder) {
			continue
		}

		newFiles[strings.TrimPrefix(k, subFolder)] = v
		newModified[strings.TrimPrefix(k, subFolder)] = m.modified[k]
	}

	return &MemoryStore{
		commonStore: m.commonStore,
		baseURL:     m.baseURL,
		data:        newFiles,
		modified:    newModified,
	}, nil
}

func (m *MemoryStore) Clone(ctx context.Context, opts ...Option) (Store, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	ms, err := newMemoryStoreContext(ctx, m.baseURL, m.extension, m.compressionType, m.overwrite, opts...)
	if err != nil {
		return nil, err
	}

	ms.data = m.data
	ms.modified = m.modified

	return ms, nil
}

func NewMemoryStore(baseURL *url.URL, extension, compressionType string, overwrite bool, opts ...Option) (*MemoryStore, error) {
	return newMemoryStoreContext(context.Background(), baseURL, extension, compressionType, overwrite, opts...)
}

func newMemoryStoreContext(_ context.Context, baseURL *url.URL, extension, compressionType string, overwrite bool, opts ...Option) (*MemoryStore, error) {
	common, err := newCommonStore(baseURL, extension, compressionType, overwrite, opts...)
	if err != nil {
		return nil, err
	}

	return &MemoryStore{
		commonStore: common,
		baseURL:     baseURL,
		data:        map[string][]byte{},
		modified:    map[string]time.Time{},
	}, nil
}

// WalkAttributes walks the objects held in memory. It reads the map directly rather than going
// through Walk, which this store does not implement yet.
func (m *MemoryStore) WalkAttributes(ctx context.Context, prefix string, f func(entry ObjectEntry) error) error {
	m.lock.RLock()
	entries := make([]ObjectEntry, 0, len(m.data))
	for _, name := range m.sortedFiles() {
		if !strings.HasPrefix(name, prefix) {
			continue
		}
		entries = append(entries, ObjectEntry{
			Name:         name,
			Size:         int64(len(m.data[name])),
			LastModified: m.modified[name],
		})
	}
	m.lock.RUnlock()

	for _, entry := range entries {
		if err := f(entry); err != nil {
			if errors.Is(err, StopIteration) {
				return nil
			}
			return err
		}
	}

	return nil
}

// sortedFiles returns the object names in lexicographic order, the order a listing would use.
// The caller must hold the lock.
func (m *MemoryStore) sortedFiles() []string {
	names := make([]string, 0, len(m.data))
	for name := range m.data {
		names = append(names, name)
	}
	sort.Strings(names)

	return names
}

// ListFolders returns the immediate sub-folders of prefix, reading them off the keys already
// held in memory.
func (m *MemoryStore) ListFolders(ctx context.Context, prefix string, max int) ([]string, error) {
	return m.ListFoldersFromTo(ctx, prefix, "", "", max)
}

// ListFoldersFromTo filters the keys held in memory.
func (m *MemoryStore) ListFoldersFromTo(ctx context.Context, prefix, inclusiveFrom, exclusiveTo string, max int) ([]string, error) {
	prefix = asFolderPrefix(prefix)
	if err := checkFolderRange(prefix, inclusiveFrom, exclusiveTo); err != nil {
		return nil, err
	}

	folders := newLimitedFolders(max)
	if folders.full() {
		return folders.folders, nil
	}

	m.lock.RLock()
	defer m.lock.RUnlock()

	seen := map[string]bool{}
	for _, name := range m.sortedFiles() {
		if !strings.HasPrefix(name, prefix) {
			continue
		}

		folder, _, isFolder := strings.Cut(strings.TrimPrefix(name, prefix), "/")
		if !isFolder || folder == "" || seen[folder] {
			continue
		}
		seen[folder] = true

		if !folderInRange(prefix+folder+"/", inclusiveFrom, exclusiveTo) {
			continue
		}

		folders.add(prefix + folder + "/")
		if folders.full() {
			break
		}
	}

	return folders.folders, nil
}