watchman/server/storage.go

package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"time"

	minio "github.com/minio/minio-go"
)

var nilLocation = os.Getenv("STORE_REGION")

const keyBucket = "key"

var bucketNames = []string{}

type Storage struct {
	client *minio.Client
}

func assertEnv(key string) (string, error) {
	value := os.Getenv(key)
	if value == "" {
		return "", errors.New(key + " not set")
	}
	return value, nil
}

func newStorage() (*Storage, error) {
	var id, loc, secret string
	var err error
	if id, err = assertEnv("STORE_ID"); err != nil {
		return nil, err
	}
	if secret, err = assertEnv("STORE_SECRET"); err != nil {
		return nil, err
	}
	if loc, err = assertEnv("STORE_LOC"); err != nil {
		return nil, err
	}
	if _, err = assertEnv("STORE_REGION"); err != nil {
		return nil, err
	}
	client, err := minio.New(loc, id, secret, false)
	if err != nil {
		return nil, err
	}
	if err := makeInitialBuckets(client); err != nil {
		return nil, err
	}
	return &Storage{
		client: client,
	}, err
}

func makeInitialBuckets(client *minio.Client) error {
	buckets, err := client.ListBuckets()
	if err != nil {
		return err
	}
	keysFound := false
	for _, bucketInfo := range buckets {
		if bucketInfo.Name == keyBucket {
			keysFound = true
		}
		bucketNames = append(bucketNames, bucketInfo.Name)
	}
	if !keysFound {
		bucketNames = append(bucketNames, keyBucket)
	}
	for _, bucketName := range bucketNames {
		if err := makeMinioBucket(client, bucketName); err != nil {
			return err
		}
	}
	return nil
}

func makeMinioBucket(client *minio.Client, bucketName string) error {
	log.Print("Making bucket", bucketName, "...")
	if err := client.MakeBucket(bucketName, nilLocation); err == nil {
		log.Print("...Made", bucketName)
		return nil
	}
	exists, err := client.BucketExists(bucketName)
	if err == nil && exists {
		log.Print("...Exists", bucketName)
		return nil
	}
	log.Print("...Can't make bucket", bucketName, err)
	return err
}

func (s *Storage) makeBucket(bucket string) error {
	if bucket == keyBucket {
		return errors.New("reserved bucket name")
	}
	err := makeMinioBucket(s.client, bucket)
	if err == nil {
		bucketNames = append(bucketNames, bucket)
	}
	return err
}

func (s *Storage) toGetKey(fname, id string) string {
	return fmt.Sprintf(
		"%v_%v",
		id,
		strings.Split(fname, ".")[0],
	)
}

func (s *Storage) toSetKey(fname, id string) string {
	return fmt.Sprintf(
		"%v_%v.%v",
		s.toGetKey(fname, id),
		time.Now().UnixNano(),
		strings.Split(fname, ".")[1],
	)
}

func (s *Storage) set(bucket, key, value string) error {
	key = encodeKey(key)
	log.Print("server setting in minio")
	if !validBucket(bucket) {
		log.Print("making bucket", bucket)
		if err := s.makeBucket(bucket); err != nil {
			log.Print("cantmake bucket", bucket)
			return err
		}
	}
	log.Print("putting", key)
	_, err := s.client.PutObject(bucket, key, bytes.NewBuffer([]byte(value)), int64(len(value)), minio.PutObjectOptions{})
	log.Print("put err:", err)
	return err
}

func (s *Storage) get(bucket, key string) (string, error) {
	key = encodeKey(key)
	if !validBucket(bucket) {
		return "", errors.New("unknown bucket")
	}
	done := make(chan struct{})
	defer close(done)
	lastKey := ""
	for obj := range s.client.ListObjectsV2(bucket, key, false, done) {
		if lastKey == "" {
			lastKey = obj.Key
		}
		if obj.Key > lastKey {
			lastKey = obj.Key
		}
	}
	if lastKey == "" {
		return "", errors.New("object not found")
	}
	obj, err := s.client.GetObject(bucket, lastKey, minio.GetObjectOptions{})
	if err != nil {
		return "", err
	}
	body, err := ioutil.ReadAll(obj)
	return string(body), err
}

func validBucket(bucket string) bool {
	for i := range bucketNames {
		if bucketNames[i] == bucket {
			return true
		}
	}
	return false
}

func encodeKey(key string) string {
	return strings.Replace(key, "/", "`", -1)
}

func decodeKey(key string) string {
	return strings.Replace(key, "`", "/", -1)
}