open-vault/vendor/github.com/couchbase/gocb/v2/cluster_query.go

315 lines
8.0 KiB
Go

package gocb
import (
"encoding/json"
"time"
gocbcore "github.com/couchbase/gocbcore/v9"
)
type jsonQueryMetrics struct {
ElapsedTime string `json:"elapsedTime"`
ExecutionTime string `json:"executionTime"`
ResultCount uint64 `json:"resultCount"`
ResultSize uint64 `json:"resultSize"`
MutationCount uint64 `json:"mutationCount,omitempty"`
SortCount uint64 `json:"sortCount,omitempty"`
ErrorCount uint64 `json:"errorCount,omitempty"`
WarningCount uint64 `json:"warningCount,omitempty"`
}
type jsonQueryWarning struct {
Code uint32 `json:"code"`
Message string `json:"msg"`
}
type jsonQueryResponse struct {
RequestID string `json:"requestID"`
ClientContextID string `json:"clientContextID"`
Status QueryStatus `json:"status"`
Warnings []jsonQueryWarning `json:"warnings"`
Metrics jsonQueryMetrics `json:"metrics"`
Profile interface{} `json:"profile"`
Signature interface{} `json:"signature"`
Prepared string `json:"prepared"`
}
// QueryMetrics encapsulates various metrics gathered during a queries execution.
type QueryMetrics struct {
ElapsedTime time.Duration
ExecutionTime time.Duration
ResultCount uint64
ResultSize uint64
MutationCount uint64
SortCount uint64
ErrorCount uint64
WarningCount uint64
}
func (metrics *QueryMetrics) fromData(data jsonQueryMetrics) error {
elapsedTime, err := time.ParseDuration(data.ElapsedTime)
if err != nil {
logDebugf("Failed to parse query metrics elapsed time: %s", err)
}
executionTime, err := time.ParseDuration(data.ExecutionTime)
if err != nil {
logDebugf("Failed to parse query metrics execution time: %s", err)
}
metrics.ElapsedTime = elapsedTime
metrics.ExecutionTime = executionTime
metrics.ResultCount = data.ResultCount
metrics.ResultSize = data.ResultSize
metrics.MutationCount = data.MutationCount
metrics.SortCount = data.SortCount
metrics.ErrorCount = data.ErrorCount
metrics.WarningCount = data.WarningCount
return nil
}
// QueryWarning encapsulates any warnings returned by a query.
type QueryWarning struct {
Code uint32
Message string
}
func (warning *QueryWarning) fromData(data jsonQueryWarning) error {
warning.Code = data.Code
warning.Message = data.Message
return nil
}
// QueryMetaData provides access to the meta-data properties of a query result.
type QueryMetaData struct {
RequestID string
ClientContextID string
Status QueryStatus
Metrics QueryMetrics
Signature interface{}
Warnings []QueryWarning
Profile interface{}
preparedName string
}
func (meta *QueryMetaData) fromData(data jsonQueryResponse) error {
metrics := QueryMetrics{}
if err := metrics.fromData(data.Metrics); err != nil {
return err
}
warnings := make([]QueryWarning, len(data.Warnings))
for wIdx, jsonWarning := range data.Warnings {
err := warnings[wIdx].fromData(jsonWarning)
if err != nil {
return err
}
}
meta.RequestID = data.RequestID
meta.ClientContextID = data.ClientContextID
meta.Status = data.Status
meta.Metrics = metrics
meta.Signature = data.Signature
meta.Warnings = warnings
meta.Profile = data.Profile
meta.preparedName = data.Prepared
return nil
}
// QueryResult allows access to the results of a query.
type QueryResult struct {
reader queryRowReader
rowBytes []byte
}
func newQueryResult(reader queryRowReader) *QueryResult {
return &QueryResult{
reader: reader,
}
}
// Next assigns the next result from the results into the value pointer, returning whether the read was successful.
func (r *QueryResult) Next() bool {
rowBytes := r.reader.NextRow()
if rowBytes == nil {
return false
}
r.rowBytes = rowBytes
return true
}
// Row returns the contents of the current row
func (r *QueryResult) Row(valuePtr interface{}) error {
if r.rowBytes == nil {
return ErrNoResult
}
if bytesPtr, ok := valuePtr.(*json.RawMessage); ok {
*bytesPtr = r.rowBytes
return nil
}
return json.Unmarshal(r.rowBytes, valuePtr)
}
// Err returns any errors that have occurred on the stream
func (r *QueryResult) Err() error {
return r.reader.Err()
}
// Close marks the results as closed, returning any errors that occurred during reading the results.
func (r *QueryResult) Close() error {
return r.reader.Close()
}
// One assigns the first value from the results into the value pointer.
// It will close the results but not before iterating through all remaining
// results, as such this should only be used for very small resultsets - ideally
// of, at most, length 1.
func (r *QueryResult) One(valuePtr interface{}) error {
// Read the bytes from the first row
valueBytes := r.reader.NextRow()
if valueBytes == nil {
return ErrNoResult
}
// Skip through the remaining rows
for r.reader.NextRow() != nil {
// do nothing with the row
}
return json.Unmarshal(valueBytes, valuePtr)
}
// MetaData returns any meta-data that was available from this query. Note that
// the meta-data will only be available once the object has been closed (either
// implicitly or explicitly).
func (r *QueryResult) MetaData() (*QueryMetaData, error) {
metaDataBytes, err := r.reader.MetaData()
if err != nil {
return nil, err
}
var jsonResp jsonQueryResponse
err = json.Unmarshal(metaDataBytes, &jsonResp)
if err != nil {
return nil, err
}
var metaData QueryMetaData
err = metaData.fromData(jsonResp)
if err != nil {
return nil, err
}
return &metaData, nil
}
type queryRowReader interface {
NextRow() []byte
Err() error
MetaData() ([]byte, error)
Close() error
PreparedName() (string, error)
}
// Query executes the query statement on the server.
func (c *Cluster) Query(statement string, opts *QueryOptions) (*QueryResult, error) {
if opts == nil {
opts = &QueryOptions{}
}
span := c.tracer.StartSpan("Query", opts.parentSpan).
SetTag("couchbase.service", "query")
defer span.Finish()
timeout := opts.Timeout
if timeout == 0 {
timeout = c.timeoutsConfig.QueryTimeout
}
deadline := time.Now().Add(timeout)
retryStrategy := c.retryStrategyWrapper
if opts.RetryStrategy != nil {
retryStrategy = newRetryStrategyWrapper(opts.RetryStrategy)
}
queryOpts, err := opts.toMap()
if err != nil {
return nil, QueryError{
InnerError: wrapError(err, "failed to generate query options"),
Statement: statement,
ClientContextID: opts.ClientContextID,
}
}
queryOpts["statement"] = statement
return c.execN1qlQuery(span, queryOpts, deadline, retryStrategy, opts.Adhoc)
}
func maybeGetQueryOption(options map[string]interface{}, name string) string {
if value, ok := options[name].(string); ok {
return value
}
return ""
}
func (c *Cluster) execN1qlQuery(
span requestSpan,
options map[string]interface{},
deadline time.Time,
retryStrategy *retryStrategyWrapper,
adHoc bool,
) (*QueryResult, error) {
provider, err := c.getQueryProvider()
if err != nil {
return nil, QueryError{
InnerError: wrapError(err, "failed to get query provider"),
Statement: maybeGetQueryOption(options, "statement"),
ClientContextID: maybeGetQueryOption(options, "client_context_id"),
}
}
eSpan := c.tracer.StartSpan("request_encoding", span.Context())
reqBytes, err := json.Marshal(options)
eSpan.Finish()
if err != nil {
return nil, QueryError{
InnerError: wrapError(err, "failed to marshall query body"),
Statement: maybeGetQueryOption(options, "statement"),
ClientContextID: maybeGetQueryOption(options, "client_context_id"),
}
}
var res queryRowReader
var qErr error
if adHoc {
res, qErr = provider.N1QLQuery(gocbcore.N1QLQueryOptions{
Payload: reqBytes,
RetryStrategy: retryStrategy,
Deadline: deadline,
TraceContext: span.Context(),
})
} else {
res, qErr = provider.PreparedN1QLQuery(gocbcore.N1QLQueryOptions{
Payload: reqBytes,
RetryStrategy: retryStrategy,
Deadline: deadline,
TraceContext: span.Context(),
})
}
if qErr != nil {
return nil, maybeEnhanceQueryError(qErr)
}
return newQueryResult(res), nil
}