package main
import (
"encoding/json"
"io"
"io/ioutil"
"gopkg.in/src-d/go-git.v4/config"
"gopkg.in/src-d/go-git.v4/core"
"github.com/aerospike/aerospike-client-go"
)
const (
urlField = "url"
referencesSet = "reference"
remotesSet = "remote"
)
type AerospikeStorage struct {
client *aerospike.Client
ns string
url string
os *AerospikeObjectStorage
rs *AerospikeReferenceStorage
cs *AerospikeConfigStorage
}
func NewAerospikeStorage(client *aerospike.Client, ns, url string) (*AerospikeStorage, error) {
if err := createIndexes(client, ns); err != nil {
return nil, err
}
return &AerospikeStorage{client: client, ns: ns, url: url}, nil
}
func (s *AerospikeStorage) ObjectStorage() core.ObjectStorage {
if s.os == nil {
s.os = &AerospikeObjectStorage{s.client, s.ns, s.url}
}
return s.os
}
func (s *AerospikeStorage) ReferenceStorage() core.ReferenceStorage {
if s.rs == nil {
s.rs = &AerospikeReferenceStorage{s.client, s.ns, s.url}
}
return s.rs
}
func (s *AerospikeStorage) ConfigStorage() config.ConfigStorage {
if s.cs == nil {
s.cs = &AerospikeConfigStorage{s.client, s.ns, s.url}
}
return s.cs
}
type AerospikeObjectStorage struct {
client *aerospike.Client
ns string
url string
}
func (s *AerospikeObjectStorage) NewObject() core.Object {
return &core.MemoryObject{}
}
func (s *AerospikeObjectStorage) Set(obj core.Object) (core.Hash, error) {
key, err := aerospike.NewKey(s.ns, obj.Type().String(), obj.Hash().String())
if err != nil {
return obj.Hash(), err
}
r, err := obj.Reader()
if err != nil {
return obj.Hash(), err
}
c, err := ioutil.ReadAll(r)
if err != nil {
return obj.Hash(), err
}
bins := aerospike.BinMap{
urlField: s.url,
"hash": obj.Hash().String(),
"type": obj.Type().String(),
"blob": c,
}
err = s.client.Put(nil, key, bins)
return obj.Hash(), err
}
func (s *AerospikeObjectStorage) Get(t core.ObjectType, h core.Hash) (core.Object, error) {
key, err := s.keyFromObject(h, t)
if err != nil {
return nil, err
}
rec, err := s.client.Get(nil, key)
if err != nil {
return nil, err
}
if rec == nil {
return nil, core.ErrObjectNotFound
}
return objectFromRecord(rec, t)
}
func (s *AerospikeObjectStorage) Iter(t core.ObjectType) (core.ObjectIter, error) {
stmnt := aerospike.NewStatement(s.ns, t.String())
err := stmnt.Addfilter(aerospike.NewEqualFilter(urlField, s.url))
rs, err := s.client.Query(nil, stmnt)
if err != nil {
return nil, err
}
return &AerospikeObjectIter{t, rs.Records}, nil
}
func (s *AerospikeObjectStorage) keyFromObject(h core.Hash, t core.ObjectType,
) (*aerospike.Key, error) {
return aerospike.NewKey(s.ns, t.String(), h.String())
}
type AerospikeObjectIter struct {
t core.ObjectType
ch chan *aerospike.Record
}
func (i *AerospikeObjectIter) Next() (core.Object, error) {
r := <-i.ch
if r == nil {
return nil, io.EOF
}
return objectFromRecord(r, i.t)
}
func (i *AerospikeObjectIter) ForEach(cb func(obj core.Object) error) error {
for {
obj, err := i.Next()
if err != nil {
if err == io.EOF {
return nil
}
return err
}
if err := cb(obj); err != nil {
if err == core.ErrStop {
return nil
}
return err
}
}
}
func (i *AerospikeObjectIter) Close() {}
func objectFromRecord(r *aerospike.Record, t core.ObjectType) (core.Object, error) {
content := r.Bins["blob"].([]byte)
o := &core.MemoryObject{}
o.SetType(t)
o.SetSize(int64(len(content)))
_, err := o.Write(content)
if err != nil {
return nil, err
}
return o, nil
}
type AerospikeReferenceStorage struct {
client *aerospike.Client
ns string
url string
}
// Set stores a reference.
func (s *AerospikeReferenceStorage) Set(ref *core.Reference) error {
key, err := aerospike.NewKey(s.ns, referencesSet, ref.Name().String())
if err != nil {
return err
}
raw := ref.Strings()
bins := aerospike.BinMap{
urlField: s.url,
"name": raw[0],
"target": raw[1],
}
return s.client.Put(nil, key, bins)
}
// Get returns a stored reference with the given name
func (s *AerospikeReferenceStorage) Get(n core.ReferenceName) (*core.Reference, error) {
key, err := aerospike.NewKey(s.ns, referencesSet, n.String())
if err != nil {
return nil, err
}
rec, err := s.client.Get(nil, key)
if err != nil {
return nil, err
}
return core.NewReferenceFromStrings(
rec.Bins["name"].(string),
rec.Bins["target"].(string),
), nil
}
// Iter returns a core.ReferenceIter
func (s *AerospikeReferenceStorage) Iter() (core.ReferenceIter, error) {
stmnt := aerospike.NewStatement(s.ns, referencesSet)
err := stmnt.Addfilter(aerospike.NewEqualFilter(urlField, s.url))
if err != nil {
return nil, err
}
rs, err := s.client.Query(nil, stmnt)
if err != nil {
return nil, err
}
var refs []*core.Reference
for r := range rs.Records {
refs = append(refs, core.NewReferenceFromStrings(
r.Bins["name"].(string),
r.Bins["target"].(string),
))
}
return core.NewReferenceSliceIter(refs), nil
}
type AerospikeConfigStorage struct {
client *aerospike.Client
ns string
url string
}
func (s *AerospikeConfigStorage) Remote(name string) (*config.RemoteConfig, error) {
key, err := aerospike.NewKey(s.ns, remotesSet, name)
if err != nil {
return nil, err
}
rec, err := s.client.Get(nil, key)
if err != nil {
return nil, err
}
return remoteFromRecord(rec)
}
func remoteFromRecord(r *aerospike.Record) (*config.RemoteConfig, error) {
content := r.Bins["blob"].([]byte)
c := &config.RemoteConfig{}
return c, json.Unmarshal(content, c)
}
func (s *AerospikeConfigStorage) Remotes() ([]*config.RemoteConfig, error) {
stmnt := aerospike.NewStatement(s.ns, remotesSet)
err := stmnt.Addfilter(aerospike.NewEqualFilter(urlField, s.url))
if err != nil {
return nil, err
}
rs, err := s.client.Query(nil, stmnt)
if err != nil {
return nil, err
return nil, err
}
var remotes []*config.RemoteConfig
for r := range rs.Records {
remote, err := remoteFromRecord(r)
if err != nil {
return nil, err
}
remotes = append(remotes, remote)
}
return remotes, nil
}
func (s *AerospikeConfigStorage) SetRemote(r *config.RemoteConfig) error {
key, err := aerospike.NewKey(s.ns, remotesSet, r.Name)
if err != nil {
return err
}
json, err := json.Marshal(r)
if err != nil {
return err
}
bins := aerospike.BinMap{
urlField: s.url,
"name": r.Name,
"blob": json,
}
return s.client.Put(nil, key, bins)
}
func (s *AerospikeConfigStorage) DeleteRemote(name string) error {
key, err := aerospike.NewKey(s.ns, remotesSet, name)
if err != nil {
return err
}
_, err = s.client.Delete(nil, key)
return err
}
func createIndexes(c *aerospike.Client, ns string) error {
for _, set := range [...]string{
referencesSet,
remotesSet,
core.BlobObject.String(),
core.TagObject.String(),
core.TreeObject.String(),
core.CommitObject.String(),
} {
if err := createIndex(c, ns, set); err != nil {
return err
}
}
return nil
}
func createIndex(c *aerospike.Client, ns, set string) error {
task, err := c.CreateIndex(nil, ns, set, set, urlField, aerospike.STRING)
if err != nil {
if err.Error() == "Index already exists" {
return nil
}
return err
}
return <-task.OnComplete()
}
