featurebase/bench/client.go

package bench

import (
	"context"
	"errors"
	"fmt"
	"github.com/pilosa/pilosa"
)

func firstHostClient(hosts []string) (*pilosa.Client, error) {
	client, err := pilosa.NewClient(hosts[0])
	if err != nil {
		return nil, err
	}
	return client, nil
}

func roundRobinClient(hosts []string, agentNum int) (*pilosa.Client, error) {
	clientNum := agentNum % len(hosts)
	return firstHostClient(hosts[clientNum:])
}

// HasClient provides a reusable component for Benchmark implementations which
// provides the Init method, a ClientType argument and a cli internal variable.
type HasClient struct {
	client      *pilosa.Client
	ClientType  string `json:"client-type"`
	ContentType string `json:"content-type"`
}

// Init for HasClient looks at the ClientType field and creates a pilosa client
// either using the first host in the list of hosts or based on the agent
// number mod len(hosts)
func (h *HasClient) Init(hosts []string, agentNum int) error {
	var err error
	switch h.ClientType {
	case "single":
		h.client, err = firstHostClient(hosts)
	case "round_robin":
		h.client, err = roundRobinClient(hosts, agentNum)
	default:
		err = fmt.Errorf("Unsupported ClientType: %v", h.ClientType)
	}
	if err != nil {
		return err
	}

	switch h.ContentType {
	case "protobuf":
		return nil
	case "pql":
		return nil
	default:
		return fmt.Errorf("Unsupported ContentType: %v", h.ContentType)

	}
}

func (h *HasClient) ExecuteQuery(contentType, db, query string, ctx context.Context) (interface{}, error) {
	if contentType == "protobuf" {
		return h.client.ExecuteQuery(ctx, db, query, true)
	} else if contentType == "pql" {
		return h.client.ExecutePQL(ctx, db, query)
	} else {
		return nil, errors.New("unsupport content type")
	}
}