ent/dialect/gremlin/request.go

// Copyright 2019-present Facebook Inc. All rights reserved.
// This source code is licensed under the Apache 2.0 license found
// in the LICENSE file in the root directory of this source tree.

package gremlin

import (
	"bytes"
	"encoding/base64"
	"time"

	"github.com/google/uuid"
	"github.com/pkg/errors"
)

type (
	// A Request models a request message sent to the server.
	Request struct {
		RequestID string                 `json:"requestId" graphson:"g:UUID"`
		Operation string                 `json:"op"`
		Processor string                 `json:"processor"`
		Arguments map[string]interface{} `json:"args"`
	}

	// RequestOption enables request customization.
	RequestOption func(*Request)

	// Credentials holds request plain auth credentials.
	Credentials struct{ Username, Password string }
)

// NewEvalRequest returns a new evaluation request request.
func NewEvalRequest(query string, opts ...RequestOption) *Request {
	r := &Request{
		RequestID: uuid.New().String(),
		Operation: OpsEval,
		Arguments: map[string]interface{}{
			ArgsGremlin:  query,
			ArgsLanguage: "gremlin-groovy",
		},
	}
	for i := range opts {
		opts[i](r)
	}
	return r
}

// NewAuthRequest returns a new auth request.
func NewAuthRequest(requestID, username, password string) *Request {
	return &Request{
		RequestID: requestID,
		Operation: OpsAuthentication,
		Arguments: map[string]interface{}{
			ArgsSasl: Credentials{
				Username: username,
				Password: password,
			},
			ArgsSaslMechanism: "PLAIN",
		},
	}
}

// WithBindings sets request bindings.
func WithBindings(bindings map[string]interface{}) RequestOption {
	return func(r *Request) {
		r.Arguments[ArgsBindings] = bindings
	}
}

// WithEvalTimeout sets script evaluation timeout.
func WithEvalTimeout(timeout time.Duration) RequestOption {
	return func(r *Request) {
		r.Arguments[ArgsEvalTimeout] = int64(timeout / time.Millisecond)
	}
}

// MarshalText implements encoding.TextMarshaler interface.
func (c Credentials) MarshalText() ([]byte, error) {
	var buf bytes.Buffer
	buf.WriteByte(0)
	buf.WriteString(c.Username)
	buf.WriteByte(0)
	buf.WriteString(c.Password)

	enc := base64.StdEncoding
	text := make([]byte, enc.EncodedLen(buf.Len()))
	enc.Encode(text, buf.Bytes())
	return text, nil
}

// UnmarshalText implements encoding.TextUnmarshaler interface.
func (c *Credentials) UnmarshalText(text []byte) error {
	enc := base64.StdEncoding
	data := make([]byte, enc.DecodedLen(len(text)))

	n, err := enc.Decode(data, text)
	if err != nil {
		return err
	}
	data = data[:n]

	parts := bytes.SplitN(data, []byte{0}, 3)
	if len(parts) != 3 {
		return errors.New("bad credentials data")
	}

	c.Username = string(parts[1])
	c.Password = string(parts[2])
	return nil
}