<commit_before>// A http.RoundTripper that retries common errors, with convenience constructors.
//
// NOTE: This meant for TEMPORARY, TRANSIENT ERRORS.
// Do not use for waiting on operations or polling of resource state,
// especially if the expected state (operation done, resource ready, etc)
// takes longer to reach than the default client Timeout.
// In those cases, retryTimeDuration(...)/resource.Retry with appropriate timeout
// and error predicates/handling should be used as a wrapper around the request
// instead.
//
// Example Usage:
// For handwritten/Go clients, the retry transport should be provided via
// the main client or a shallow copy of the HTTP resources, depending on the
// API-specific retry predicates.
// Example Usage in Terraform Config:
//	client := oauth2.NewClient(ctx, tokenSource)
//	// Create with default retry predicates
//	client.Transport := NewTransportWithDefaultRetries(client.Transport, defaultTimeout)
//
//	// If API uses just default retry predicates:
//	c.clientCompute, err = compute.NewService(ctx, option.WithHTTPClient(client))
//	...
//	// If API needs custom additional retry predicates:
//	sqlAdminHttpClient := ClientWithAdditionalRetries(client, retryTransport,
//			isTemporarySqlError1,
//			isTemporarySqlError2)
//	c.clientSqlAdmin, err = compute.NewService(ctx, option.WithHTTPClient(sqlAdminHttpClient))
// ...

package google

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"google.golang.org/api/googleapi"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"time"
)

const defaultRetryTransportTimeoutSec = 30

// NewTransportWithDefaultRetries constructs a default retryTransport that will retry common temporary errors
func NewTransportWithDefaultRetries(t http.RoundTripper) *retryTransport {
	return &retryTransport{
		retryPredicates: defaultErrorRetryPredicates,
		internal:        t,
	}
}

// Helper method to create a shallow copy of an HTTP client with a shallow-copied retryTransport
// s.t. the base HTTP transport is the same (i.e. client connection pools are shared, retryPredicates are different)
func ClientWithAdditionalRetries(baseClient *http.Client, baseRetryTransport *retryTransport, predicates ...RetryErrorPredicateFunc) *http.Client {
	copied := *baseClient
	if baseRetryTransport == nil {
		baseRetryTransport = NewTransportWithDefaultRetries(baseClient.Transport)
	}
	copied.Transport = baseRetryTransport.WithAddedPredicates(predicates...)
	return &copied
}

// Returns a shallow copy of the retry transport with additional retry
// predicates but same wrapped http.RoundTripper
func (t *retryTransport) WithAddedPredicates(predicates ...RetryErrorPredicateFunc) *retryTransport {
	copyT := *t
	copyT.retryPredicates = append(t.retryPredicates, predicates...)
	return &copyT
}

type retryTransport struct {
	retryPredicates []RetryErrorPredicateFunc
	internal        http.RoundTripper
}

// RoundTrip implements the RoundTripper interface method.
// It retries the given HTTP request based on the retry predicates
// registered under the retryTransport.
func (t *retryTransport) RoundTrip(req *http.Request) (resp *http.Response, respErr error) {
	// Set timeout to default value.
	ctx := req.Context()
	var ccancel context.CancelFunc
	if _, ok := ctx.Deadline(); !ok {
		ctx, ccancel = context.WithTimeout(ctx, defaultRetryTransportTimeoutSec*time.Second)
		defer func() {
			if ctx.Err() == nil {
				// Cleanup child context created for retry loop if ctx not done.
				ccancel()
			}
		}()
	}

	attempts := 0
	backoff := time.Millisecond * 500
	nextBackoff := time.Millisecond * 500

	log.Printf("[DEBUG] Retry Transport: starting RoundTrip retry loop")
Retry:
	for {
		log.Printf("[DEBUG] Retry Transport: request attempt %d", attempts)

		// Copy the request - we dont want to use the original request as
		// RoundTrip contract says request body can/will be consumed
		newRequest, copyErr := copyHttpRequest(req)
		if copyErr != nil {
			respErr = errwrap.Wrapf("unable to copy invalid http.Request for retry: {{err}}", copyErr)
			break Retry
		}

		// Do the wrapped Roundtrip. This is one request in the retry loop.
		resp, respErr = t.internal.RoundTrip(newRequest)
		attempts++

		retryErr := t.checkForRetryableError(resp, respErr)
		if retryErr == nil {
			log.Printf("[DEBUG] Retry Transport: Stopping retries, last request was successful")
			break Retry
		}
		if !retryErr.Retryable {
			log.Printf("[DEBUG] Retry Transport: Stopping retries, last request failed with non-retryable error: %s", retryErr.Err)
			break Retry
		}

		log.Printf("[DEBUG] Retry Transport: Waiting %s before trying request again", backoff)
		select {
		case <-ctx.Done():
			log.Printf("[DEBUG] Retry Transport: Stopping retries, context done: %v", ctx.Err())
			break Retry
		case <-time.After(backoff):
			log.Printf("[DEBUG] Retry Transport: Finished waiting %s before next retry", backoff)

			// Fibonnaci backoff - 0.5, 1, 1.5, 2.5, 4, 6.5, 10.5, ...
			lastBackoff := backoff
			backoff = backoff + nextBackoff
			nextBackoff = lastBackoff
			continue
		}
	}

	log.Printf("[DEBUG] Retry Transport: Returning after %d attempts", attempts)
	return resp, respErr
}

// copyHttpRequest provides an copy of the given HTTP request for one RoundTrip.
// If the request has a non-empty body (io.ReadCloser), the body is deep copied
// so it can be consumed.
func copyHttpRequest(req *http.Request) (*http.Request, error) {
	newRequest := *req

	if req.Body == nil || req.Body == http.NoBody {
		return &newRequest, nil
	}

	// Helpers like http.NewRequest add a GetBody for copying.
	// If not given, we should reject the request.
	if req.GetBody == nil {
		return nil, errors.New("invalid HTTP request for transport, expected request.GetBody for non-empty Body")
	}

	bd, err := req.GetBody()
	if err != nil {
		return nil, err
	}

	newRequest.Body = bd
	return &newRequest, nil
}

// checkForRetryableError uses the googleapi.CheckResponse util to check for
// errors in the response, and determines whether there is a retryable error.
// in response/response error.
func (t *retryTransport) checkForRetryableError(resp *http.Response, respErr error) *resource.RetryError {
	var errToCheck error

	if respErr != nil {
		errToCheck = respErr
	} else {
		respToCheck := *resp
		// The RoundTrip contract states that the HTTP response/response error
		// returned cannot be edited. We need to consume the Body to check for
		// errors, so we need to create a copy if the Response has a body.
		if resp.Body != nil && resp.Body != http.NoBody {
			// Use httputil.DumpResponse since the only important info is
			// error code and messages in the response body.
			dumpBytes, err := httputil.DumpResponse(resp, true)
			if err != nil {
				return resource.NonRetryableError(fmt.Errorf("unable to check response for error: %v", err))
			}
			respToCheck.Body = ioutil.NopCloser(bytes.NewReader(dumpBytes))
		}
		errToCheck = googleapi.CheckResponse(&respToCheck)
	}

	if isRetryableError(errToCheck, t.retryPredicates...) {
		return resource.RetryableError(errToCheck)
	}
	return resource.NonRetryableError(errToCheck)
}
<commit_msg>Consume original request body from retry transport (#3276)<commit_after>// A http.RoundTripper that retries common errors, with convenience constructors.
//
// NOTE: This meant for TEMPORARY, TRANSIENT ERRORS.
// Do not use for waiting on operations or polling of resource state,
// especially if the expected state (operation done, resource ready, etc)
// takes longer to reach than the default client Timeout.
// In those cases, retryTimeDuration(...)/resource.Retry with appropriate timeout
// and error predicates/handling should be used as a wrapper around the request
// instead.
//
// Example Usage:
// For handwritten/Go clients, the retry transport should be provided via
// the main client or a shallow copy of the HTTP resources, depending on the
// API-specific retry predicates.
// Example Usage in Terraform Config:
//	client := oauth2.NewClient(ctx, tokenSource)
//	// Create with default retry predicates
//	client.Transport := NewTransportWithDefaultRetries(client.Transport, defaultTimeout)
//
//	// If API uses just default retry predicates:
//	c.clientCompute, err = compute.NewService(ctx, option.WithHTTPClient(client))
//	...
//	// If API needs custom additional retry predicates:
//	sqlAdminHttpClient := ClientWithAdditionalRetries(client, retryTransport,
//			isTemporarySqlError1,
//			isTemporarySqlError2)
//	c.clientSqlAdmin, err = compute.NewService(ctx, option.WithHTTPClient(sqlAdminHttpClient))
// ...

package google

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"google.golang.org/api/googleapi"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"time"
)

const defaultRetryTransportTimeoutSec = 30

// NewTransportWithDefaultRetries constructs a default retryTransport that will retry common temporary errors
func NewTransportWithDefaultRetries(t http.RoundTripper) *retryTransport {
	return &retryTransport{
		retryPredicates: defaultErrorRetryPredicates,
		internal:        t,
	}
}

// Helper method to create a shallow copy of an HTTP client with a shallow-copied retryTransport
// s.t. the base HTTP transport is the same (i.e. client connection pools are shared, retryPredicates are different)
func ClientWithAdditionalRetries(baseClient *http.Client, baseRetryTransport *retryTransport, predicates ...RetryErrorPredicateFunc) *http.Client {
	copied := *baseClient
	if baseRetryTransport == nil {
		baseRetryTransport = NewTransportWithDefaultRetries(baseClient.Transport)
	}
	copied.Transport = baseRetryTransport.WithAddedPredicates(predicates...)
	return &copied
}

// Returns a shallow copy of the retry transport with additional retry
// predicates but same wrapped http.RoundTripper
func (t *retryTransport) WithAddedPredicates(predicates ...RetryErrorPredicateFunc) *retryTransport {
	copyT := *t
	copyT.retryPredicates = append(t.retryPredicates, predicates...)
	return &copyT
}

type retryTransport struct {
	retryPredicates []RetryErrorPredicateFunc
	internal        http.RoundTripper
}

// RoundTrip implements the RoundTripper interface method.
// It retries the given HTTP request based on the retry predicates
// registered under the retryTransport.
func (t *retryTransport) RoundTrip(req *http.Request) (resp *http.Response, respErr error) {
	// Set timeout to default value.
	ctx := req.Context()
	var ccancel context.CancelFunc
	if _, ok := ctx.Deadline(); !ok {
		ctx, ccancel = context.WithTimeout(ctx, defaultRetryTransportTimeoutSec*time.Second)
		defer func() {
			if ctx.Err() == nil {
				// Cleanup child context created for retry loop if ctx not done.
				ccancel()
			}
		}()
	}

	attempts := 0
	backoff := time.Millisecond * 500
	nextBackoff := time.Millisecond * 500

	log.Printf("[DEBUG] Retry Transport: starting RoundTrip retry loop")
Retry:
	for {
		log.Printf("[DEBUG] Retry Transport: request attempt %d", attempts)

		// Copy the request - we dont want to use the original request as
		// RoundTrip contract says request body can/will be consumed
		newRequest, copyErr := copyHttpRequest(req)
		if copyErr != nil {
			respErr = errwrap.Wrapf("unable to copy invalid http.Request for retry: {{err}}", copyErr)
			break Retry
		}

		// Do the wrapped Roundtrip. This is one request in the retry loop.
		resp, respErr = t.internal.RoundTrip(newRequest)
		attempts++

		retryErr := t.checkForRetryableError(resp, respErr)
		if retryErr == nil {
			log.Printf("[DEBUG] Retry Transport: Stopping retries, last request was successful")
			break Retry
		}
		if !retryErr.Retryable {
			log.Printf("[DEBUG] Retry Transport: Stopping retries, last request failed with non-retryable error: %s", retryErr.Err)
			break Retry
		}

		log.Printf("[DEBUG] Retry Transport: Waiting %s before trying request again", backoff)
		select {
		case <-ctx.Done():
			log.Printf("[DEBUG] Retry Transport: Stopping retries, context done: %v", ctx.Err())
			break Retry
		case <-time.After(backoff):
			log.Printf("[DEBUG] Retry Transport: Finished waiting %s before next retry", backoff)

			// Fibonnaci backoff - 0.5, 1, 1.5, 2.5, 4, 6.5, 10.5, ...
			lastBackoff := backoff
			backoff = backoff + nextBackoff
			nextBackoff = lastBackoff
			continue
		}
	}

	// VCR depends on the original request body being consumed, so consume it here
	_, err := httputil.DumpRequestOut(req, true)
	if err != nil {
		log.Printf("[DEBUG] Retry Transport: Reading request failed: %v", err)
	}

	log.Printf("[DEBUG] Retry Transport: Returning after %d attempts", attempts)
	return resp, respErr
}

// copyHttpRequest provides an copy of the given HTTP request for one RoundTrip.
// If the request has a non-empty body (io.ReadCloser), the body is deep copied
// so it can be consumed.
func copyHttpRequest(req *http.Request) (*http.Request, error) {
	newRequest := *req

	if req.Body == nil || req.Body == http.NoBody {
		return &newRequest, nil
	}

	// Helpers like http.NewRequest add a GetBody for copying.
	// If not given, we should reject the request.
	if req.GetBody == nil {
		return nil, errors.New("invalid HTTP request for transport, expected request.GetBody for non-empty Body")
	}

	bd, err := req.GetBody()
	if err != nil {
		return nil, err
	}

	newRequest.Body = bd
	return &newRequest, nil
}

// checkForRetryableError uses the googleapi.CheckResponse util to check for
// errors in the response, and determines whether there is a retryable error.
// in response/response error.
func (t *retryTransport) checkForRetryableError(resp *http.Response, respErr error) *resource.RetryError {
	var errToCheck error

	if respErr != nil {
		errToCheck = respErr
	} else {
		respToCheck := *resp
		// The RoundTrip contract states that the HTTP response/response error
		// returned cannot be edited. We need to consume the Body to check for
		// errors, so we need to create a copy if the Response has a body.
		if resp.Body != nil && resp.Body != http.NoBody {
			// Use httputil.DumpResponse since the only important info is
			// error code and messages in the response body.
			dumpBytes, err := httputil.DumpResponse(resp, true)
			if err != nil {
				return resource.NonRetryableError(fmt.Errorf("unable to check response for error: %v", err))
			}
			respToCheck.Body = ioutil.NopCloser(bytes.NewReader(dumpBytes))
		}
		errToCheck = googleapi.CheckResponse(&respToCheck)
	}

	if isRetryableError(errToCheck, t.retryPredicates...) {
		return resource.RetryableError(errToCheck)
	}
	return resource.NonRetryableError(errToCheck)
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package serialization

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"hash"
	"io"

	"v.io/v23/security"
	"v.io/v23/verror"
	"v.io/v23/vom"
)

var (
	errCantBeNilSigner = verror.Register(pkgPath+".errCantBeNilSigner", verror.NoRetry, "{1:}{2:} data:{3} signature:{4} signer:{5} cannot be nil{:_}")
	errCantSign        = verror.Register(pkgPath+".errCantSign", verror.NoRetry, "{1:}{2:} signing failed{:_}")
)

const defaultChunkSizeBytes = 1 << 20

// signingWriter implements io.WriteCloser.
type signingWriter struct {
	data      io.WriteCloser
	signature io.WriteCloser
	signer    Signer

	chunkSizeBytes int64
	curChunk       bytes.Buffer
	signatureHash  hash.Hash
	sigEnc         *vom.Encoder
}

func (w *signingWriter) Write(p []byte) (int, error) {
	bytesWritten := 0
	for len(p) > 0 {
		pLimited := p
		curChunkFreeBytes := w.chunkSizeBytes - int64(w.curChunk.Len())
		if int64(len(pLimited)) > curChunkFreeBytes {
			pLimited = pLimited[:curChunkFreeBytes]
		}

		n, err := w.curChunk.Write(pLimited)
		bytesWritten = bytesWritten + n
		if err != nil {
			return bytesWritten, err
		}
		p = p[n:]

		if err := w.commitChunk(false); err != nil {
			return bytesWritten, err
		}
	}
	return bytesWritten, nil
}

func (w *signingWriter) Close() error {
	if w.curChunk.Len() > 0 {
		if err := w.commitChunk(true); err != nil {
			defer w.close()
			return err
		}
	}
	if err := w.commitSignature(); err != nil {
		defer w.close()
		return err
	}
	return w.close()
}

// Options specifies parameters to tune a SigningWriteCloser.
type Options struct {
	// ChunkSizeBytes controls the maximum amount of memory devoted to buffering
	// data provided to Write calls. See NewSigningWriteCloser.
	ChunkSizeBytes int64
}

// Signer is the interface for digital signature operations used by NewSigningWriteCloser.
type Signer interface {
	Sign(message []byte) (security.Signature, error)
	PublicKey() security.PublicKey
}

// NewSigningWriteCloser returns an io.WriteCloser that writes data along
// with an appropriate signature that establishes the integrity and
// authenticity of the data. It behaves as follows:
// * A Write call writes chunks (of size provided by the Options or 1MB by default)
//   of data to the provided data WriteCloser and a hash of the chunks to the provided
//   signature WriteCloser.
// * A Close call writes a signature (computed using the provided signer) of
//   all the hashes written, and then closes the data and signature WriteClosers.
func NewSigningWriteCloser(data, signature io.WriteCloser, s Signer, opts *Options) (io.WriteCloser, error) {
	if (data == nil) || (signature == nil) || (s == nil) {
		return nil, verror.New(errCantBeNilSigner, nil, data, signature, s)
	}
	enc := vom.NewEncoder(signature)
	w := &signingWriter{data: data, signature: signature, signer: s, signatureHash: sha256.New(), chunkSizeBytes: defaultChunkSizeBytes, sigEnc: enc}

	if opts != nil {
		w.chunkSizeBytes = opts.ChunkSizeBytes
	}

	if err := w.commitHeader(); err != nil {
		return nil, err
	}
	return w, nil
}

func (w *signingWriter) commitHeader() error {
	if err := binary.Write(w.signatureHash, binary.LittleEndian, int64(w.chunkSizeBytes)); err != nil {
		return err
	}
	if err := w.sigEnc.Encode(SignedHeader{w.chunkSizeBytes}); err != nil {
		return err
	}
	return nil
}

func (w *signingWriter) commitChunk(force bool) error {
	if !force && int64(w.curChunk.Len()) < w.chunkSizeBytes {
		return nil
	}

	hashBytes := sha256.Sum256(w.curChunk.Bytes())
	if _, err := io.CopyN(w.data, &w.curChunk, int64(w.curChunk.Len())); err != nil {
		return err
	}
	if _, err := w.signatureHash.Write(hashBytes[:]); err != nil {
		return err
	}
	return w.sigEnc.Encode(SignedDataHash{hashBytes})
}

func (w *signingWriter) commitSignature() error {
	sig, err := w.signer.Sign(w.signatureHash.Sum(nil))
	if err != nil {
		return verror.New(errCantSign, nil, err)
	}

	return w.sigEnc.Encode(SignedDataSignature{sig})
}

func (w *signingWriter) close() error {
	var closeErr error
	if err := w.data.Close(); err != nil && closeErr == nil {
		closeErr = err
	}
	if err := w.signature.Close(); err != nil && closeErr == nil {
		closeErr = err
	}
	return closeErr
}
<commit_msg>lib/security/serialization: doc fix<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package serialization

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"hash"
	"io"

	"v.io/v23/security"
	"v.io/v23/verror"
	"v.io/v23/vom"
)

var (
	errCantBeNilSigner = verror.Register(pkgPath+".errCantBeNilSigner", verror.NoRetry, "{1:}{2:} data:{3} signature:{4} signer:{5} cannot be nil{:_}")
	errCantSign        = verror.Register(pkgPath+".errCantSign", verror.NoRetry, "{1:}{2:} signing failed{:_}")
)

const defaultChunkSizeBytes = 1 << 20

// signingWriter implements io.WriteCloser.
type signingWriter struct {
	data      io.WriteCloser
	signature io.WriteCloser
	signer    Signer

	chunkSizeBytes int64
	curChunk       bytes.Buffer
	signatureHash  hash.Hash
	sigEnc         *vom.Encoder
}

func (w *signingWriter) Write(p []byte) (int, error) {
	bytesWritten := 0
	for len(p) > 0 {
		pLimited := p
		curChunkFreeBytes := w.chunkSizeBytes - int64(w.curChunk.Len())
		if int64(len(pLimited)) > curChunkFreeBytes {
			pLimited = pLimited[:curChunkFreeBytes]
		}

		n, err := w.curChunk.Write(pLimited)
		bytesWritten = bytesWritten + n
		if err != nil {
			return bytesWritten, err
		}
		p = p[n:]

		if err := w.commitChunk(false); err != nil {
			return bytesWritten, err
		}
	}
	return bytesWritten, nil
}

func (w *signingWriter) Close() error {
	if w.curChunk.Len() > 0 {
		if err := w.commitChunk(true); err != nil {
			defer w.close()
			return err
		}
	}
	if err := w.commitSignature(); err != nil {
		defer w.close()
		return err
	}
	return w.close()
}

// Options specifies parameters to tune a SigningWriteCloser.
type Options struct {
	// ChunkSizeBytes controls the maximum amount of memory devoted to buffering
	// data provided to Write calls. See NewSigningWriteCloser.
	ChunkSizeBytes int64
}

// Signer is the interface for digital signature operations used by NewSigningWriteCloser.
type Signer interface {
	Sign(message []byte) (security.Signature, error)
	PublicKey() security.PublicKey
}

// NewSigningWriteCloser returns an io.WriteCloser that writes data along
// with an appropriate signature that establishes the integrity and
// authenticity of the data. It behaves as follows:
//     * A Write call writes chunks (of size provided by the Options or
//       1MB by default) of data to the provided data WriteCloser and a
//       hash of the chunks to the provided signature WriteCloser.
//     * A Close call writes a signature (computed using the provided
//       signer) of all the hashes written, and then closes the data and
//       signature WriteClosers.
func NewSigningWriteCloser(data, signature io.WriteCloser, s Signer, opts *Options) (io.WriteCloser, error) {
	if (data == nil) || (signature == nil) || (s == nil) {
		return nil, verror.New(errCantBeNilSigner, nil, data, signature, s)
	}
	enc := vom.NewEncoder(signature)
	w := &signingWriter{data: data, signature: signature, signer: s, signatureHash: sha256.New(), chunkSizeBytes: defaultChunkSizeBytes, sigEnc: enc}

	if opts != nil {
		w.chunkSizeBytes = opts.ChunkSizeBytes
	}

	if err := w.commitHeader(); err != nil {
		return nil, err
	}
	return w, nil
}

func (w *signingWriter) commitHeader() error {
	if err := binary.Write(w.signatureHash, binary.LittleEndian, int64(w.chunkSizeBytes)); err != nil {
		return err
	}
	if err := w.sigEnc.Encode(SignedHeader{w.chunkSizeBytes}); err != nil {
		return err
	}
	return nil
}

func (w *signingWriter) commitChunk(force bool) error {
	if !force && int64(w.curChunk.Len()) < w.chunkSizeBytes {
		return nil
	}

	hashBytes := sha256.Sum256(w.curChunk.Bytes())
	if _, err := io.CopyN(w.data, &w.curChunk, int64(w.curChunk.Len())); err != nil {
		return err
	}
	if _, err := w.signatureHash.Write(hashBytes[:]); err != nil {
		return err
	}
	return w.sigEnc.Encode(SignedDataHash{hashBytes})
}

func (w *signingWriter) commitSignature() error {
	sig, err := w.signer.Sign(w.signatureHash.Sum(nil))
	if err != nil {
		return verror.New(errCantSign, nil, err)
	}

	return w.sigEnc.Encode(SignedDataSignature{sig})
}

func (w *signingWriter) close() error {
	var closeErr error
	if err := w.data.Close(); err != nil && closeErr == nil {
		closeErr = err
	}
	if err := w.signature.Close(); err != nil && closeErr == nil {
		closeErr = err
	}
	return closeErr
}
<|endoftext|>


<commit_before>// Copyright 2012 Tetsuo Kiso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pegasos

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
)

// reaeder
//
// implements file reader to read files in libsvm format.

func readLines(r io.Reader) ([]Example, int) {
	rd := bufio.NewReader(r)
	lineNum := 1
	var data []Example
	maxId := 0
	for {
		line, err := rd.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		x, id, err := Tokenize(strings.TrimRight(line, "\n"))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal line at %d\n", lineNum)
			log.Fatal(err)
		}
		data = append(data, x)

		if id > maxId {
			maxId = id
		}

		lineNum++
	}
	return data, maxId
}

func ReadTrainingData(filename string) ([]Example, int) {
	file, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	return readLines(file)
}
<commit_msg>Skip illegal lines.<commit_after>// Copyright 2012 Tetsuo Kiso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pegasos

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
)

// reaeder
//
// implements file reader to read files in libsvm format.

func readLines(r io.Reader) ([]Example, int) {
	rd := bufio.NewReader(r)
	lineNum := 1
	var data []Example
	maxId := 0
	for {
		line, err := rd.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		x, id, err := Tokenize(strings.TrimRight(line, "\n"))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal line at %d\n", lineNum)
			continue
		}
		data = append(data, x)

		if id > maxId {
			maxId = id
		}

		lineNum++
	}
	return data, maxId
}

func ReadTrainingData(filename string) ([]Example, int) {
	file, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	return readLines(file)
}
<|endoftext|>


<commit_before>// Copyright 2016,2017 Yaacov Zamir <kobi.zamir@gmail.com>
// and other contributors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package storage
package memory

import (
	"log"
	"net/url"
	"regexp"
	"time"

	"github.com/MohawkTSDB/mohawk/src/storage"
)

type TimeValuePair struct {
	timeStamp int64
	value     float64
}

type TimeSeries struct {
	tags      map[string]string
	data      []TimeValuePair
	lastValue TimeValuePair
}

type Tenant struct {
	ts map[string]*TimeSeries
}

type Storage struct {
	timeGranularitySec int64
	timeRetentionSec   int64
	timeLastSec        int64

	tenant map[string]*Tenant
}

// Storage functions
// Required by storage interface

func (r Storage) Name() string {
	return "Storage-Memory"
}

func (r *Storage) Open(options url.Values) {
	// set last entry time
	r.timeLastSec = 0
	// set time granularity to 30 sec
	r.timeGranularitySec = 30
	// set time retention to 7 days
	r.timeRetentionSec = 7 * 24 * 60 * 60

	// open db connection
	r.tenant = make(map[string]*Tenant, 0)

	// start a maintenance worker that will clean the db periodically
	go r.maintenance()
}

func (r Storage) GetTenants() []storage.Tenant {
	res := make([]storage.Tenant, 0, len(r.tenant))

	// return a list of tenants
	for key := range r.tenant {
		res = append(res, storage.Tenant{ID: key})
	}

	return res
}

func (r Storage) GetItemList(tenant string, tags map[string]string) []storage.Item {
	res := make([]storage.Item, 0)
	t, ok := r.tenant[tenant]

	// check tenant
	if !ok {
		return res
	}

	for key, ts := range t.ts {
		if hasMatchingTag(tags, ts.tags) {
			lastValue := storage.DataItem{
				Timestamp: ts.lastValue.timeStamp,
				Value:     ts.lastValue.value,
			}

			res = append(res, storage.Item{
				ID:         key,
				Type:       "gauge",
				Tags:       ts.tags,
				LastValues: []storage.DataItem{lastValue},
			})
		}
	}

	return res
}

func (r *Storage) GetRawData(tenant string, id string, end int64, start int64, limit int64, order string) []storage.DataItem {
	res := make([]storage.DataItem, 0)

	arraySize := r.timeRetentionSec / r.timeGranularitySec
	pStart := r.getPosForTimestamp(start)
	pEnd := r.getPosForTimestamp(end)

	// make sure start and end times is in the retention time
	start, end = r.checkTimespan(start, end)

	// sanity check pEnd
	if pEnd <= pStart {
		pEnd += arraySize
	}

	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// fill data out array
	count := int64(0)
	for i := pEnd; count < limit && i >= pStart; i-- {
		d := r.tenant[tenant].ts[id].data[i%arraySize]

		// if this is a valid point
		if d.timeStamp < end && d.timeStamp >= start {
			count++
			res = append(res, storage.DataItem{
				Timestamp: d.timeStamp,
				Value:     d.value,
			})
		}
	}

	// order
	if order == "ASC" {
		for i := 0; i < len(res)/2; i++ {
			j := len(res) - i - 1
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r Storage) GetStatData(tenant string, id string, end int64, start int64, limit int64, order string, bucketDuration int64) []storage.StatItem {
	res := make([]storage.StatItem, 0)

	// make sure start and end times is in the retention time
	start, end = r.checkTimespan(start, end)

	// bucketDuration can't be smaller then granularity
	if bucketDuration < r.timeGranularitySec {
		bucketDuration = r.timeGranularitySec
	}
	bucketDuration = r.timeGranularitySec * (bucketDuration / r.timeGranularitySec)

	// start and tend must be integer multiplections of bucketDuration
	bucketDurationMilli := bucketDuration * 1000
	start = bucketDurationMilli * (start / bucketDurationMilli)
	end = bucketDurationMilli * (1 + end/bucketDurationMilli)

	arraySize := r.timeRetentionSec / r.timeGranularitySec
	pStep := bucketDuration / r.timeGranularitySec
	pStart := r.getPosForTimestamp(start)
	pEnd := r.getPosForTimestamp(end)

	// sanity check pEnd
	if pEnd <= pStart {
		pEnd += arraySize
	}

	// sanity check step
	if pStep < 1 {
		pStep = 1
	}
	if pStep > (pEnd - pStart) {
		pStep = pEnd - pStart
	}

	startTimestamp := end
	stepMillisec := pStep * r.timeGranularitySec * 1000

	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// fill data out array
	count := int64(0)
	for b := pEnd; count < limit && b > pStart && startTimestamp >= stepMillisec; b -= pStep {
		samples := int64(0)
		sum := float64(0)
		first := float64(0)
		last := float64(0)
		min := float64(0)
		max := float64(0)

		// loop on all points in bucket
		for i := (b - pStep); i < b; i++ {
			d := r.tenant[tenant].ts[id].data[i%arraySize]
			if d.timeStamp <= end && d.timeStamp > start {
				samples++

				// calculate bucket stat values
				if samples == 1 {
					// first sample
					first = d.value
					min = first
					max = first
				} else {
					// all samples except first sample
					if min > d.value {
						min = d.value
					}
					if max < d.value {
						max = d.value
					}
				}

				last = d.value
				sum += d.value
			}
		}

		// all points are valid
		startTimestamp -= stepMillisec
		count++

		// all points are valid
		if samples > 0 {
			res = append(res, storage.StatItem{
				Start:   startTimestamp,
				End:     startTimestamp + stepMillisec,
				Empty:   false,
				Samples: samples,
				First:   first,
				Last:    last,
				Min:     min,
				Max:     max,
				Avg:     sum / float64(samples),
				Sum:     sum,
			})
		} else {
			count++
			res = append(res, storage.StatItem{
				Start: startTimestamp,
				End:   startTimestamp + stepMillisec,
				Empty: true,
			})
		}
	}

	// order
	if order == "ASC" {
		for i := 0; i < len(res)/2; i++ {
			j := len(res) - i - 1
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r *Storage) PostRawData(tenant string, id string, t int64, v float64) bool {
	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// update time value pair to the time serias
	p := r.getPosForTimestamp(t)
	r.tenant[tenant].ts[id].data[p] = TimeValuePair{timeStamp: t, value: v}

	// update last value
	if r.tenant[tenant].ts[id].lastValue.timeStamp < t {
		r.tenant[tenant].ts[id].lastValue.timeStamp = t
		r.tenant[tenant].ts[id].lastValue.value = v
	}

	// update last
	tSec := t / 1000
	if tSec > r.timeLastSec {
		r.timeLastSec = tSec
	}

	return true
}

func (r *Storage) PutTags(tenant string, id string, tags map[string]string) bool {
	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// update time serias tags
	if len(tags) > 0 {
		for key, value := range tags {
			r.tenant[tenant].ts[id].tags[key] = value
		}
	}

	return true
}

func (r *Storage) DeleteData(tenant string, id string, end int64, start int64) bool {
	return true
}

func (r *Storage) DeleteTags(tenant string, id string, tags []string) bool {
	return true
}

// Helper functions
// Not required by storage interface

func (r *Storage) checkTimespan(start int64, end int64) (int64, int64) {
	memFirstTime := (r.timeLastSec - r.timeRetentionSec) * 1000
	memLastTime := r.timeLastSec * 1000

	// make sure start and end times is in the retention time
	if start < memFirstTime {
		start = memFirstTime
	}
	if end > memLastTime {
		end = memLastTime + 1
	}

	return start, end
}

func (r *Storage) getPosForTimestamp(timestamp int64) int64 {
	arraySize := r.timeRetentionSec / r.timeGranularitySec
	arrayPos := timestamp / 1000 / r.timeGranularitySec

	return arrayPos % arraySize
}

func (r *Storage) checkID(tenant string, id string) {
	var ok bool

	// check for tenant
	if _, ok = r.tenant[tenant]; !ok {
		r.tenant[tenant] = &Tenant{ts: make(map[string]*TimeSeries)}
	}

	// check for TimeSeries
	if _, ok = r.tenant[tenant].ts[id]; !ok {
		r.tenant[tenant].ts[id] = &TimeSeries{
			tags: make(map[string]string),
			data: make([]TimeValuePair, r.timeRetentionSec/r.timeGranularitySec),
		}
	}
}

func hasMatchingTag(tags map[string]string, itemTags map[string]string) bool {
	out := true

	// if no tags, all items match
	if len(tags) == 0 {
		return true
	}

	// if item has no tags, item is invalid
	if len(itemTags) == 0 {
		return false
	}

	// loop on all the tags, we need _all_ query tags to match tags on item
	for key, value := range tags {
		r, _ := regexp.Compile("^" + value + "$")
		out = out && r.MatchString(itemTags[key])
	}

	return out
}

func (r *Storage) maintenance() {
	// clean data every 120 minutes
	c := time.Tick(120 * time.Minute)

	// once a tick clean data
	for range c {
		log.Printf("maintenance: start\n")
		r.cleanData()
	}
}

func (r *Storage) cleanData() {
	var lastTimeStampSec int64
	validTimeStamp := time.Now().Unix() - r.timeRetentionSec

	// loop on all tenants
	for _, t := range r.tenant {
		// loop on all time series in this tenant
		for key, ts := range t.ts {
			lastTimeStampSec = ts.lastValue.timeStamp / 1000

			// if last value is more then time span old, remove data
			if lastTimeStampSec <= validTimeStamp {
				log.Printf("maintenance: delete item %s\n", key)
				delete(t.ts, key)
			}
		}

		// TODO: delete tenant if no time seriess
	}
}
<commit_msg>Refactor memory stats<commit_after>// Copyright 2016,2017 Yaacov Zamir <kobi.zamir@gmail.com>
// and other contributors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package storage
package memory

import (
	"log"
	"net/url"
	"regexp"
	"time"

	"github.com/MohawkTSDB/mohawk/src/storage"
)

type TimeValuePair struct {
	timeStamp int64
	value     float64
}

type TimeSeries struct {
	tags      map[string]string
	data      []TimeValuePair
	lastValue TimeValuePair
}

type Tenant struct {
	ts map[string]*TimeSeries
}

type Storage struct {
	timeGranularitySec int64
	timeRetentionSec   int64
	timeLastSec        int64

	tenant map[string]*Tenant
}

// Storage functions
// Required by storage interface

func (r Storage) Name() string {
	return "Storage-Memory"
}

func (r *Storage) Open(options url.Values) {
	// set last entry time
	r.timeLastSec = 0
	// set time granularity to 30 sec
	r.timeGranularitySec = 30
	// set time retention to 7 days
	r.timeRetentionSec = 7 * 24 * 60 * 60

	// open db connection
	r.tenant = make(map[string]*Tenant, 0)

	// start a maintenance worker that will clean the db periodically
	go r.maintenance()
}

func (r Storage) GetTenants() []storage.Tenant {
	res := make([]storage.Tenant, 0, len(r.tenant))

	// return a list of tenants
	for key := range r.tenant {
		res = append(res, storage.Tenant{ID: key})
	}

	return res
}

func (r Storage) GetItemList(tenant string, tags map[string]string) []storage.Item {
	res := make([]storage.Item, 0)
	t, ok := r.tenant[tenant]

	// check tenant
	if !ok {
		return res
	}

	for key, ts := range t.ts {
		if hasMatchingTag(tags, ts.tags) {
			lastValue := storage.DataItem{
				Timestamp: ts.lastValue.timeStamp,
				Value:     ts.lastValue.value,
			}

			res = append(res, storage.Item{
				ID:         key,
				Type:       "gauge",
				Tags:       ts.tags,
				LastValues: []storage.DataItem{lastValue},
			})
		}
	}

	return res
}

func (r *Storage) GetRawData(tenant string, id string, end int64, start int64, limit int64, order string) []storage.DataItem {
	res := make([]storage.DataItem, 0)

	arraySize := r.timeRetentionSec / r.timeGranularitySec
	pStart := r.getPosForTimestamp(start)
	pEnd := r.getPosForTimestamp(end)

	// make sure start and end times is in the retention time
	start, end = r.checkTimespan(start, end)

	// sanity check pEnd
	if pEnd <= pStart {
		pEnd += arraySize
	}

	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// fill data out array
	count := int64(0)
	for i := pEnd; count < limit && i >= pStart; i-- {
		d := r.tenant[tenant].ts[id].data[i%arraySize]

		// if this is a valid point
		if d.timeStamp < end && d.timeStamp >= start {
			count++
			res = append(res, storage.DataItem{
				Timestamp: d.timeStamp,
				Value:     d.value,
			})
		}
	}

	// order
	if order == "ASC" {
		for i := 0; i < len(res)/2; i++ {
			j := len(res) - i - 1
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r Storage) GetStatTimes(end int64, start int64, limit int64, bucketDuration int64) (int64, int64, int64, int64, int64, int64) {
	// make sure start and end times is in the retention time
	start, end = r.checkTimespan(start, end)

	// bucketDuration can't be smaller then granularity
	if bucketDuration < r.timeGranularitySec {
		bucketDuration = r.timeGranularitySec
	}
	bucketDuration = r.timeGranularitySec * (bucketDuration / r.timeGranularitySec)

	// start and tend must be integer multiplections of bucketDuration
	bucketDurationMilli := bucketDuration * 1000
	start = bucketDurationMilli * (start / bucketDurationMilli)
	end = bucketDurationMilli * (1 + end/bucketDurationMilli)

	arraySize := r.timeRetentionSec / r.timeGranularitySec
	pStep := bucketDuration / r.timeGranularitySec
	pStart := r.getPosForTimestamp(start)
	pEnd := r.getPosForTimestamp(end)

	// sanity check pEnd
	if pEnd <= pStart {
		pEnd += arraySize
	}

	// sanity check step
	if pStep < 1 {
		pStep = 1
	}
	if pStep > (pEnd - pStart) {
		pStep = pEnd - pStart
	}

	return end, start, pEnd, pStart, pStep, arraySize
}

func (r Storage) GetStatData(tenant string, id string, end int64, start int64, limit int64, order string, bucketDuration int64) []storage.StatItem {
	res := make([]storage.StatItem, 0)
	end, start, pEnd, pStart, pStep, arraySize := r.GetStatTimes(end, start, limit, bucketDuration)

	startTimestamp := end
	stepMillisec := pStep * r.timeGranularitySec * 1000

	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// fill data out array
	count := int64(0)
	for b := pEnd; count < limit && b > pStart && startTimestamp >= stepMillisec; b -= pStep {
		samples := int64(0)
		sum := float64(0)
		first := float64(0)
		last := float64(0)
		min := float64(0)
		max := float64(0)

		// loop on all points in bucket
		for i := (b - pStep); i < b; i++ {
			d := r.tenant[tenant].ts[id].data[i%arraySize]
			if d.timeStamp <= end && d.timeStamp > start {
				samples++

				// calculate bucket stat values
				if samples == 1 {
					// first sample
					first = d.value
					min = first
					max = first
				} else {
					// all samples except first sample
					if min > d.value {
						min = d.value
					}
					if max < d.value {
						max = d.value
					}
				}

				last = d.value
				sum += d.value
			}
		}

		// all points are valid
		startTimestamp -= stepMillisec
		count++

		// all points are valid
		if samples > 0 {
			res = append(res, storage.StatItem{
				Start:   startTimestamp,
				End:     startTimestamp + stepMillisec,
				Empty:   false,
				Samples: samples,
				First:   first,
				Last:    last,
				Min:     min,
				Max:     max,
				Avg:     sum / float64(samples),
				Sum:     sum,
			})
		} else {
			count++
			res = append(res, storage.StatItem{
				Start: startTimestamp,
				End:   startTimestamp + stepMillisec,
				Empty: true,
			})
		}
	}

	// order
	if order == "ASC" {
		for i := 0; i < len(res)/2; i++ {
			j := len(res) - i - 1
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r *Storage) PostRawData(tenant string, id string, t int64, v float64) bool {
	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// update time value pair to the time serias
	p := r.getPosForTimestamp(t)
	r.tenant[tenant].ts[id].data[p] = TimeValuePair{timeStamp: t, value: v}

	// update last value
	if r.tenant[tenant].ts[id].lastValue.timeStamp < t {
		r.tenant[tenant].ts[id].lastValue.timeStamp = t
		r.tenant[tenant].ts[id].lastValue.value = v
	}

	// update last
	tSec := t / 1000
	if tSec > r.timeLastSec {
		r.timeLastSec = tSec
	}

	return true
}

func (r *Storage) PutTags(tenant string, id string, tags map[string]string) bool {
	// check if tenant and id exists, create them if necessary
	r.checkID(tenant, id)

	// update time serias tags
	if len(tags) > 0 {
		for key, value := range tags {
			r.tenant[tenant].ts[id].tags[key] = value
		}
	}

	return true
}

func (r *Storage) DeleteData(tenant string, id string, end int64, start int64) bool {
	return true
}

func (r *Storage) DeleteTags(tenant string, id string, tags []string) bool {
	return true
}

// Helper functions
// Not required by storage interface

func (r *Storage) checkTimespan(start int64, end int64) (int64, int64) {
	memFirstTime := (r.timeLastSec - r.timeRetentionSec) * 1000
	memLastTime := r.timeLastSec * 1000

	// make sure start and end times is in the retention time
	if start < memFirstTime {
		start = memFirstTime
	}
	if end > memLastTime {
		end = memLastTime + 1
	}

	return start, end
}

func (r *Storage) getPosForTimestamp(timestamp int64) int64 {
	arraySize := r.timeRetentionSec / r.timeGranularitySec
	arrayPos := timestamp / 1000 / r.timeGranularitySec

	return arrayPos % arraySize
}

func (r *Storage) checkID(tenant string, id string) {
	var ok bool

	// check for tenant
	if _, ok = r.tenant[tenant]; !ok {
		r.tenant[tenant] = &Tenant{ts: make(map[string]*TimeSeries)}
	}

	// check for TimeSeries
	if _, ok = r.tenant[tenant].ts[id]; !ok {
		r.tenant[tenant].ts[id] = &TimeSeries{
			tags: make(map[string]string),
			data: make([]TimeValuePair, r.timeRetentionSec/r.timeGranularitySec),
		}
	}
}

func hasMatchingTag(tags map[string]string, itemTags map[string]string) bool {
	out := true

	// if no tags, all items match
	if len(tags) == 0 {
		return true
	}

	// if item has no tags, item is invalid
	if len(itemTags) == 0 {
		return false
	}

	// loop on all the tags, we need _all_ query tags to match tags on item
	for key, value := range tags {
		r, _ := regexp.Compile("^" + value + "$")
		out = out && r.MatchString(itemTags[key])
	}

	return out
}

func (r *Storage) maintenance() {
	// clean data every 120 minutes
	c := time.Tick(120 * time.Minute)

	// once a tick clean data
	for range c {
		log.Printf("maintenance: start\n")
		r.cleanData()
	}
}

func (r *Storage) cleanData() {
	var lastTimeStampSec int64
	validTimeStamp := time.Now().Unix() - r.timeRetentionSec

	// loop on all tenants
	for _, t := range r.tenant {
		// loop on all time series in this tenant
		for key, ts := range t.ts {
			lastTimeStampSec = ts.lastValue.timeStamp / 1000

			// if last value is more then time span old, remove data
			if lastTimeStampSec <= validTimeStamp {
				log.Printf("maintenance: delete item %s\n", key)
				delete(t.ts, key)
			}
		}

		// TODO: delete tenant if no time seriess
	}
}
<|endoftext|>


<commit_before>package backend_test

import (
	"fmt"
	"io"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"

	"github.com/restic/restic/backend"
)

// assert fails the test if the condition is false.
func assert(tb testing.TB, condition bool, msg string, v ...interface{}) {
	if !condition {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: "+msg+"\033[39m\n\n", append([]interface{}{filepath.Base(file), line}, v...)...)
		tb.FailNow()
	}
}

// ok fails the test if an err is not nil.
func ok(tb testing.TB, err error) {
	if err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: unexpected error: %s\033[39m\n\n", filepath.Base(file), line, err.Error())
		tb.FailNow()
	}
}

// equals fails the test if exp is not equal to act.
func equals(tb testing.TB, exp, act interface{}) {
	if !reflect.DeepEqual(exp, act) {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n", filepath.Base(file), line, exp, act)
		tb.FailNow()
	}
}

func str2id(s string) backend.ID {
	id, err := backend.ParseID(s)
	if err != nil {
		panic(err)
	}

	return id
}

type mockBackend struct {
	list       func(backend.Type, <-chan struct{}) <-chan string
	get        func(backend.Type, backend.ID) ([]byte, error)
	getReader  func(backend.Type, backend.ID) (io.ReadCloser, error)
	create     func(backend.Type, []byte) (backend.ID, error)
	createFrom func(backend.Type, io.Reader) (backend.ID, error)
	test       func(backend.Type, backend.ID) (bool, error)
	remove     func(backend.Type, backend.ID) error
	close      func() error
}

func (m mockBackend) List(t backend.Type, done <-chan struct{}) <-chan string {
	return m.list(t, done)
}

func (m mockBackend) Get(t backend.Type, id backend.ID) ([]byte, error) {
	return m.get(t, id)
}

func (m mockBackend) GetReader(t backend.Type, id backend.ID) (io.ReadCloser, error) {
	return m.getReader(t, id)
}

func (m mockBackend) Create(t backend.Type, data []byte) (backend.ID, error) {
	return m.create(t, data)
}

func (m mockBackend) CreateFrom(t backend.Type, r io.Reader) (backend.ID, error) {
	return m.createFrom(t, r)
}

func (m mockBackend) Test(t backend.Type, id backend.ID) (bool, error) {
	return m.test(t, id)
}

func (m mockBackend) Remove(t backend.Type, id backend.ID) error {
	return m.remove(t, id)
}

func (m mockBackend) Close() error {
	return m.close()
}

var samples = backend.IDs{
	str2id("20bdc1402a6fc9b633aaffffffffffffffffffffffffffffffffffffffffffff"),
	str2id("20bdc1402a6fc9b633ccd578c4a92d0f4ef1a457fa2e16c596bc73fb409d6cc0"),
	str2id("20bdc1402a6fc9b633ffffffffffffffffffffffffffffffffffffffffffffff"),
	str2id("20ff988befa5fc40350f00d531a767606efefe242c837aaccb80673f286be53d"),
	str2id("326cb59dfe802304f96ee9b5b9af93bdee73a30f53981e5ec579aedb6f1d0f07"),
	str2id("86b60b9594d1d429c4aa98fa9562082cabf53b98c7dc083abe5dae31074dd15a"),
	str2id("96c8dbe225079e624b5ce509f5bd817d1453cd0a85d30d536d01b64a8669aeae"),
	str2id("fa31d65b87affcd167b119e9d3d2a27b8236ca4836cb077ed3e96fcbe209b792"),
}

func TestPrefixLength(t *testing.T) {
	list := samples

	m := mockBackend{}
	m.list = func(t backend.Type, done <-chan struct{}) <-chan string {
		ch := make(chan string)
		go func() {
			defer close(ch)
			for _, id := range list {
				select {
				case ch <- id.String():
				case <-done:
					return
				}
			}
		}()
		return ch
	}

	l, err := backend.PrefixLength(m, backend.Snapshot)
	ok(t, err)
	equals(t, 19, l)

	list = samples[:3]
	l, err = backend.PrefixLength(m, backend.Snapshot)
	ok(t, err)
	equals(t, 19, l)

	list = samples[3:]
	l, err = backend.PrefixLength(m, backend.Snapshot)
	ok(t, err)
	equals(t, 8, l)
}
<commit_msg>Refactor mockBackend for test<commit_after>package backend_test

import (
	"fmt"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"

	"github.com/restic/restic/backend"
)

// assert fails the test if the condition is false.
func assert(tb testing.TB, condition bool, msg string, v ...interface{}) {
	if !condition {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: "+msg+"\033[39m\n\n", append([]interface{}{filepath.Base(file), line}, v...)...)
		tb.FailNow()
	}
}

// ok fails the test if an err is not nil.
func ok(tb testing.TB, err error) {
	if err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: unexpected error: %s\033[39m\n\n", filepath.Base(file), line, err.Error())
		tb.FailNow()
	}
}

// equals fails the test if exp is not equal to act.
func equals(tb testing.TB, exp, act interface{}) {
	if !reflect.DeepEqual(exp, act) {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n", filepath.Base(file), line, exp, act)
		tb.FailNow()
	}
}

func str2id(s string) backend.ID {
	id, err := backend.ParseID(s)
	if err != nil {
		panic(err)
	}

	return id
}

type mockBackend struct {
	list func(backend.Type, <-chan struct{}) <-chan string
}

func (m mockBackend) List(t backend.Type, done <-chan struct{}) <-chan string {
	return m.list(t, done)
}

var samples = backend.IDs{
	str2id("20bdc1402a6fc9b633aaffffffffffffffffffffffffffffffffffffffffffff"),
	str2id("20bdc1402a6fc9b633ccd578c4a92d0f4ef1a457fa2e16c596bc73fb409d6cc0"),
	str2id("20bdc1402a6fc9b633ffffffffffffffffffffffffffffffffffffffffffffff"),
	str2id("20ff988befa5fc40350f00d531a767606efefe242c837aaccb80673f286be53d"),
	str2id("326cb59dfe802304f96ee9b5b9af93bdee73a30f53981e5ec579aedb6f1d0f07"),
	str2id("86b60b9594d1d429c4aa98fa9562082cabf53b98c7dc083abe5dae31074dd15a"),
	str2id("96c8dbe225079e624b5ce509f5bd817d1453cd0a85d30d536d01b64a8669aeae"),
	str2id("fa31d65b87affcd167b119e9d3d2a27b8236ca4836cb077ed3e96fcbe209b792"),
}

func TestPrefixLength(t *testing.T) {
	list := samples

	m := mockBackend{}
	m.list = func(t backend.Type, done <-chan struct{}) <-chan string {
		ch := make(chan string)
		go func() {
			defer close(ch)
			for _, id := range list {
				select {
				case ch <- id.String():
				case <-done:
					return
				}
			}
		}()
		return ch
	}

	l, err := backend.PrefixLength(m, backend.Snapshot)
	ok(t, err)
	equals(t, 19, l)

	list = samples[:3]
	l, err = backend.PrefixLength(m, backend.Snapshot)
	ok(t, err)
	equals(t, 19, l)

	list = samples[3:]
	l, err = backend.PrefixLength(m, backend.Snapshot)
	ok(t, err)
	equals(t, 8, l)
}
<|endoftext|>


<commit_before>package main

import (
	"net/url"
	"time"

	"io"

	"errors"

	"github.com/PuerkitoBio/goquery"
	"github.com/catatsuy/private-isu/benchmarker/checker"
	"github.com/catatsuy/private-isu/benchmarker/util"
)

// 1ページに表示される画像にリクエストする
// TODO: 画像には並列リクエストするべきでは？
func loadImages(s *checker.Session, imageUrls []string) {
	for _, url := range imageUrls {
		imgReq := checker.NewAssetAction(url, &checker.Asset{})
		imgReq.Description = "投稿画像"
		imgReq.Play(s)
	}
}

func extractImages(body io.Reader) ([]string, error) {
	imageUrls := []string{}

	doc, err := goquery.NewDocumentFromReader(body)
	if err != nil {
		return nil, errors.New("ページが正しく読み込めませんでした")
	}

	doc.Find("img").Each(func(_ int, selection *goquery.Selection) {
		url, _ := selection.Attr("src")
		imageUrls = append(imageUrls, url)
	}).Length()

	return imageUrls, nil
}

// 普通のページに表示されるべき静的ファイルに一通りアクセス
func loadAssets(s *checker.Session) {
	a := checker.NewAssetAction("/favicon.ico", &checker.Asset{})
	a.ExpectedLocation = "/favicon.ico"
	a.Description = "favicon.ico"
	a.Play(s)

	a = checker.NewAssetAction("js/jquery-2.2.0.js", &checker.Asset{})
	a.ExpectedLocation = "js/jquery-2.2.0.js"
	a.Description = "js/jquery-2.2.0.js"
	a.Play(s)

	a = checker.NewAssetAction("js/jquery.timeago.js", &checker.Asset{})
	a.ExpectedLocation = "js/jquery.timeago.js"
	a.Description = "js/jquery.timeago.js"
	a.Play(s)

	a = checker.NewAssetAction("js/jquery.timeago.ja.js", &checker.Asset{})
	a.ExpectedLocation = "js/jquery.timeago.ja.js"
	a.Description = "js/jquery.timeago.ja.js"
	a.Play(s)

	a = checker.NewAssetAction("/js/main.js", &checker.Asset{})
	a.ExpectedLocation = "/js/main.js"
	a.Description = "js/main.js"
	a.Play(s)

	a = checker.NewAssetAction("/css/style.css", &checker.Asset{})
	a.ExpectedLocation = "/css/style.css"
	a.Description = "/css/style.css"
	a.Play(s)
}

// / にリクエストしてもっと見るを10ページ辿る
// 負荷をかけられればいいので
func indexMoreAndMoreScenario(s *checker.Session) {
	var imageUrls []string

	imagePerPageChecker := func(s *checker.Session, body io.Reader) error {
		imageUrls, err := extractImages(body)
		if err != nil {
			return err
		}
		if len(imageUrls) < PostsPerPage {
			return errors.New("1ページに表示される画像の数が足りません")
		}
		return nil
	}

	index := checker.NewAction("GET", "/")
	index.ExpectedLocation = "/"
	index.Description = "インデックスページ"
	index.CheckFunc = imagePerPageChecker
	index.Play(s)

	loadAssets(s)
	loadImages(s, imageUrls)

	offset := util.RandomNumber(10) // 10は適当。URLをバラけさせるため
	for i := 0; i < 10; i++ {       // 10ページ辿る
		maxCreatedAt := time.Date(2016, time.January, 2, 11, 46, 21-PostsPerPage*i+offset, 0, time.FixedZone("Asia/Tokyo", 9*60*60))

		imageUrls = []string{}
		posts := checker.NewAction("GET", "/posts?max_created_at="+url.QueryEscape(maxCreatedAt.Format(time.RFC3339)))
		posts.Description = "インデックスページの「もっと見る」"
		posts.CheckFunc = imagePerPageChecker
		posts.Play(s)

		loadImages(s, imageUrls)
	}
}

// インデックスページを5回表示するだけ（負荷かける用）
func loadIndexScenario(s *checker.Session) {
	var imageUrls []string

	imagePerPageChecker := func(s *checker.Session, body io.Reader) error {
		imageUrls, err := extractImages(body)
		if err != nil {
			return err
		}
		if len(imageUrls) < PostsPerPage {
			return errors.New("1ページに表示される画像の数が足りません")
		}
		return nil
	}

	index := checker.NewAction("GET", "/")
	index.ExpectedLocation = "/"
	index.Description = "インデックスページ"
	index.CheckFunc = imagePerPageChecker
	index.Play(s)

	loadAssets(s)
	loadImages(s, imageUrls)

	for i := 0; i < 4; i++ {
		// あとの4回はDOMをパースしない。トップページをキャッシュして超高速に返されたとき対策
		index := checker.NewAction("GET", "/")
		index.ExpectedLocation = "/"
		index.Description = "インデックスページ"

		loadAssets(s)
		loadImages(s, imageUrls) // 画像は初回と同じものにリクエスト投げる
	}
}
<commit_msg>15秒以上たったらページを辿るのをやめる<commit_after>package main

import (
	"net/url"
	"time"

	"io"

	"errors"

	"github.com/PuerkitoBio/goquery"
	"github.com/catatsuy/private-isu/benchmarker/checker"
	"github.com/catatsuy/private-isu/benchmarker/util"
)

// 1ページに表示される画像にリクエストする
// TODO: 画像には並列リクエストするべきでは？
func loadImages(s *checker.Session, imageUrls []string) {
	for _, url := range imageUrls {
		imgReq := checker.NewAssetAction(url, &checker.Asset{})
		imgReq.Description = "投稿画像"
		imgReq.Play(s)
	}
}

func extractImages(body io.Reader) ([]string, error) {
	imageUrls := []string{}

	doc, err := goquery.NewDocumentFromReader(body)
	if err != nil {
		return nil, errors.New("ページが正しく読み込めませんでした")
	}

	doc.Find("img").Each(func(_ int, selection *goquery.Selection) {
		url, _ := selection.Attr("src")
		imageUrls = append(imageUrls, url)
	}).Length()

	return imageUrls, nil
}

// 普通のページに表示されるべき静的ファイルに一通りアクセス
func loadAssets(s *checker.Session) {
	a := checker.NewAssetAction("/favicon.ico", &checker.Asset{})
	a.ExpectedLocation = "/favicon.ico"
	a.Description = "favicon.ico"
	a.Play(s)

	a = checker.NewAssetAction("js/jquery-2.2.0.js", &checker.Asset{})
	a.ExpectedLocation = "js/jquery-2.2.0.js"
	a.Description = "js/jquery-2.2.0.js"
	a.Play(s)

	a = checker.NewAssetAction("js/jquery.timeago.js", &checker.Asset{})
	a.ExpectedLocation = "js/jquery.timeago.js"
	a.Description = "js/jquery.timeago.js"
	a.Play(s)

	a = checker.NewAssetAction("js/jquery.timeago.ja.js", &checker.Asset{})
	a.ExpectedLocation = "js/jquery.timeago.ja.js"
	a.Description = "js/jquery.timeago.ja.js"
	a.Play(s)

	a = checker.NewAssetAction("/js/main.js", &checker.Asset{})
	a.ExpectedLocation = "/js/main.js"
	a.Description = "js/main.js"
	a.Play(s)

	a = checker.NewAssetAction("/css/style.css", &checker.Asset{})
	a.ExpectedLocation = "/css/style.css"
	a.Description = "/css/style.css"
	a.Play(s)
}

// インデックスにリクエストして「もっと見る」を最大10ページ辿る
// 15秒たったら問答無用で打ち切る
func indexMoreAndMoreScenario(s *checker.Session) {
	var imageUrls []string
	now := time.Now()

	imagePerPageChecker := func(s *checker.Session, body io.Reader) error {
		imageUrls, err := extractImages(body)
		if err != nil {
			return err
		}
		if len(imageUrls) < PostsPerPage {
			return errors.New("1ページに表示される画像の数が足りません")
		}
		return nil
	}

	index := checker.NewAction("GET", "/")
	index.ExpectedLocation = "/"
	index.Description = "インデックスページ"
	index.CheckFunc = imagePerPageChecker
	index.Play(s)

	loadAssets(s)
	loadImages(s, imageUrls)

	offset := util.RandomNumber(10) // 10は適当。URLをバラけさせるため
	for i := 0; i < 10; i++ {       // 10ページ辿る
		maxCreatedAt := time.Date(2016, time.January, 2, 11, 46, 21-PostsPerPage*i+offset, 0, time.FixedZone("Asia/Tokyo", 9*60*60))

		imageUrls = []string{}
		posts := checker.NewAction("GET", "/posts?max_created_at="+url.QueryEscape(maxCreatedAt.Format(time.RFC3339)))
		posts.Description = "インデックスページの「もっと見る」"
		posts.CheckFunc = imagePerPageChecker
		posts.Play(s)

		loadImages(s, imageUrls)

		if time.Now().Sub(now) > 15*time.Second {
			break
		}
	}
}

// インデックスページを5回表示するだけ（負荷かける用）
// 15秒たったら問答無用で打ち切る
func loadIndexScenario(s *checker.Session) {
	var imageUrls []string
	now := time.Now()

	imagePerPageChecker := func(s *checker.Session, body io.Reader) error {
		imageUrls, err := extractImages(body)
		if err != nil {
			return err
		}
		if len(imageUrls) < PostsPerPage {
			return errors.New("1ページに表示される画像の数が足りません")
		}
		return nil
	}

	index := checker.NewAction("GET", "/")
	index.ExpectedLocation = "/"
	index.Description = "インデックスページ"
	index.CheckFunc = imagePerPageChecker
	index.Play(s)

	loadAssets(s)
	loadImages(s, imageUrls)

	for i := 0; i < 4; i++ {
		// あとの4回はDOMをパースしない。トップページをキャッシュして超高速に返されたとき対策
		index := checker.NewAction("GET", "/")
		index.ExpectedLocation = "/"
		index.Description = "インデックスページ"

		loadAssets(s)
		loadImages(s, imageUrls) // 画像は初回と同じものにリクエスト投げる

		if time.Now().Sub(now) > 15*time.Second {
			break
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go-uuid/uuid"
	"code.google.com/p/gorest"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	_ "github.com/lib/pq"
	"log"
	"net/http"
	"os"
	"strings"
)

var g_db *sql.DB

var (
	errEntityDoesntExist       = errors.New("The entity doesn't exist.")
	errParentEntityDoesntExist = errors.New("The parent entity doesn't exist.")
)

func main() {
	conf := ReadConfig()

	var err error
	g_db, err = sql.Open("postgres", fmt.Sprintf("user=%v password=%v dbname=%v sslmode=disable", conf.DbUser, conf.DbPassword, conf.DbName))

	if err != nil {
		log.Fatal(err)
	}

	gorest.RegisterService(new(ServersService))
	gorest.RegisterService(new(InstancesService))
	gorest.RegisterService(new(PlayersService))
	gorest.RegisterService(new(EventsService))

	http.Handle("/", gorest.Handle())
	http.ListenAndServe(fmt.Sprintf(":%v", conf.HttpPort), nil)
}

type EventsService struct {
	gorest.RestService `root:"/events/{serverUrlId:string}/{instanceUrlId:string}" consumes:"application/json" produces:"application/json"`
	postEvents         gorest.EndPoint `method:"POST" path:"/" postdata:"[]Event"`
}

type ServersService struct {
	gorest.RestService `root:"/servers" consumes:"application/json" produces:"application/json"`
	putServer          gorest.EndPoint `method:"PUT" path:"/{urlId:string}" postdata:"Server"`
}

type InstancesService struct {
	gorest.RestService `root:"/instances/{serverUrlId:string}" consumes:"application/json" produces:"application/json"`
	putInstance        gorest.EndPoint `method:"PUT" path:"/{urlId:string}" postdata:"Instance"`
}

type PlayersService struct {
	gorest.RestService `root:"/players" consumes:"application/json" produces:"application/json"`
	getUser            gorest.EndPoint `method:"GET" path:"/{id:int64}" output:"Player"`
	putUser            gorest.EndPoint `method:"PUT" path:"/{id:int64}" postdata:"Player"`
}

type Event struct {
	Id       string
	Time     string
	Type     string
	Metadata string
}

type Instance struct {
	Id       string
	UrlId    string
	ServerId string
	Name     string
}

type Server struct {
	Id    string // The id is the "secret key" that the server admins put in the madrust-stats plugin config.
	UrlId string
	Name  string
}

type Player struct {
	Id   int64
	Name string
}

type Config struct {
	DbName     string
	DbUser     string
	DbPassword string
	HttpPort   int
}

func (serv EventsService) PostEvents(events []Event, serverUrlid string, instanceUrlId string) {
	for _, event := range events {
		fmt.Print(event.Metadata)
	}
}

func (serv InstancesService) PutInstance(instance Instance, serverUrlId string, urlId string) {
	if strings.TrimSpace(instance.Name) == "" {
		SendErrorResponse(serv.ResponseBuilder(), 400, "Name is required.")
		return
	}

	_, err := GetInstance(serverUrlId, urlId)

	switch {
	case err == errParentEntityDoesntExist:
		SendErrorResponse(serv.ResponseBuilder(), 404, fmt.Sprintf("No server exists with urlId %q.", serverUrlId))
		return
	case err == nil:
		SendErrorResponse(serv.ResponseBuilder(), 403, fmt.Sprintf("The urlId %q is already taken.", urlId))
		return
	case err != nil && err != errEntityDoesntExist:
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	}

	server, err := GetServer(serverUrlId)

	if err != nil {
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	}

	instance.Id = uuid.New()
	instance.ServerId = server.Id
	instance.UrlId = urlId

	_, insertErr := g_db.Exec("INSERT INTO instances (id, url_id, server_id, name) VALUES ($1, $2, $3, $4)", instance.Id, instance.UrlId, instance.ServerId, instance.Name)

	if insertErr != nil {
		SendInternalServerError(serv.ResponseBuilder(), insertErr)
		return
	}

	serv.ResponseBuilder().SetResponseCode(201).Overide(true)
}

func (serv ServersService) PutServer(server Server, urlId string) {
	_, err := GetServer(urlId)

	switch {
	case err != nil && err != errEntityDoesntExist:
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	case err == nil:
		SendErrorResponse(serv.ResponseBuilder(), 403, fmt.Sprintf("The urlId %q is already taken.", urlId))
		return
	}

	server.Id = uuid.New()
	server.UrlId = urlId

	_, insertErr := g_db.Exec("INSERT INTO servers (id, url_id, name) VALUES ($1, $2, $3)", server.Id, server.UrlId, server.Name)

	if insertErr != nil {
		SendInternalServerError(serv.ResponseBuilder(), insertErr)
	}

	serv.ResponseBuilder().SetResponseCode(201).Overide(true)
}

func (serv PlayersService) GetUser(id int64) Player {
	var steamid int64
	var username string

	err := g_db.QueryRow("SELECT id, name FROM players WHERE id = $1", id).Scan(&steamid, &username)

	switch {
	case err == sql.ErrNoRows:
		SendErrorResponse(serv.ResponseBuilder(), 404, "User not found.")
	case err != nil:
		SendInternalServerError(serv.ResponseBuilder(), err)
	}

	return Player{Id: steamid, Name: username}
}

func (serv PlayersService) PutUser(user Player, id int64) {

	if user.Id != id {
		SendErrorResponse(serv.ResponseBuilder(), 400, "Mismatched ids")
		return
	}

	var exists int
	err := g_db.QueryRow("SELECT count(*) FROM players WHERE id = $1", id).Scan(&exists)

	if err != nil {
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	}

	if exists > 0 {
		SendErrorResponse(serv.ResponseBuilder(), 400, fmt.Sprintf("User %v already exists", id))
		return
	}

	_, insertErr := g_db.Exec("INSERT INTO players (id, name) VALUES ($1, $2)", user.Id, user.Name)

	if insertErr != nil {
		SendInternalServerError(serv.ResponseBuilder(), insertErr)
		return
	}

	serv.ResponseBuilder().SetResponseCode(201).Overide(true)
}

func ReadConfig() (config Config) {
	file, err := os.Open("conf.json")
	defer file.Close()

	if err != nil {
		log.Fatal(err)
	}

	decoder := json.NewDecoder(file)
	config = Config{}
	err = decoder.Decode(&config)

	if err != nil {
		log.Fatal(err)
	}

	return config
}

func SendErrorResponse(responseBuilder *gorest.ResponseBuilder, responseCode int, message string) {
	responseBuilder.SetResponseCode(responseCode).WriteAndOveride([]byte(message))
}

func SendInternalServerError(responseBuilder *gorest.ResponseBuilder, err error) {
	SendErrorResponse(responseBuilder, 500, fmt.Sprintf("An unexpected error occurred: %v", err))
}

func GetServer(urlId string) (Server, error) {
	var id string
	var name string

	err := g_db.QueryRow("SELECT id, name FROM servers WHERE url_id = $1", urlId).Scan(&id, &name)

	switch {
	case err == sql.ErrNoRows:
		return Server{}, errEntityDoesntExist
	case err != nil:
		return Server{}, err
	}

	return Server{Id: id, UrlId: urlId, Name: name}, nil
}

func GetInstance(serverUrlId string, urlId string) (Instance, error) {
	server, err := GetServer(serverUrlId)

	switch {
	case err == errEntityDoesntExist:
		return Instance{}, errParentEntityDoesntExist
	case err != nil:
		return Instance{}, err
	}

	var id string
	var name string

	err = g_db.QueryRow("SELECT id, name FROM instances WHERE server_id = $1 AND url_id = $2", server.Id, urlId).Scan(&id, &name)

	switch {
	case err == sql.ErrNoRows:
		return Instance{}, errEntityDoesntExist
	case err != nil:
		return Instance{}, err
	}

	return Instance{Id: id, UrlId: urlId, ServerId: server.Id}, nil
}
<commit_msg>Events service should use non-url (uuid) ids.<commit_after>package main

import (
	"code.google.com/p/go-uuid/uuid"
	"code.google.com/p/gorest"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	_ "github.com/lib/pq"
	"log"
	"net/http"
	"os"
	"strings"
)

var g_db *sql.DB

var (
	errEntityDoesntExist       = errors.New("The entity doesn't exist.")
	errParentEntityDoesntExist = errors.New("The parent entity doesn't exist.")
)

func main() {
	conf := ReadConfig()

	var err error
	g_db, err = sql.Open("postgres", fmt.Sprintf("user=%v password=%v dbname=%v sslmode=disable", conf.DbUser, conf.DbPassword, conf.DbName))

	if err != nil {
		log.Fatal(err)
	}

	gorest.RegisterService(new(ServersService))
	gorest.RegisterService(new(InstancesService))
	gorest.RegisterService(new(PlayersService))
	gorest.RegisterService(new(EventsService))

	http.Handle("/", gorest.Handle())
	http.ListenAndServe(fmt.Sprintf(":%v", conf.HttpPort), nil)
}

type EventsService struct {
	gorest.RestService `root:"/events/{serverId:string}/{instanceId:string}" consumes:"application/json" produces:"application/json"`
	postEvents         gorest.EndPoint `method:"POST" path:"/" postdata:"[]Event"`
}

type ServersService struct {
	gorest.RestService `root:"/servers" consumes:"application/json" produces:"application/json"`
	putServer          gorest.EndPoint `method:"PUT" path:"/{urlId:string}" postdata:"Server"`
}

type InstancesService struct {
	gorest.RestService `root:"/instances/{serverUrlId:string}" consumes:"application/json" produces:"application/json"`
	putInstance        gorest.EndPoint `method:"PUT" path:"/{urlId:string}" postdata:"Instance"`
}

type PlayersService struct {
	gorest.RestService `root:"/players" consumes:"application/json" produces:"application/json"`
	getUser            gorest.EndPoint `method:"GET" path:"/{id:int64}" output:"Player"`
	putUser            gorest.EndPoint `method:"PUT" path:"/{id:int64}" postdata:"Player"`
}

type Event struct {
	Id       string
	Time     string
	Type     string
	Metadata string
}

type Instance struct {
	Id       string
	UrlId    string
	ServerId string
	Name     string
}

type Server struct {
	Id    string // The id is the "secret key" that the server admins put in the madrust-stats plugin config.
	UrlId string
	Name  string
}

type Player struct {
	Id   int64
	Name string
}

type Config struct {
	DbName     string
	DbUser     string
	DbPassword string
	HttpPort   int
}

func (serv EventsService) PostEvents(events []Event, serverId string, instanceId string) {
	for _, event := range events {
		fmt.Print(event.Metadata)
	}
}

func (serv InstancesService) PutInstance(instance Instance, serverUrlId string, urlId string) {
	if strings.TrimSpace(instance.Name) == "" {
		SendErrorResponse(serv.ResponseBuilder(), 400, "Name is required.")
		return
	}

	_, err := GetInstance(serverUrlId, urlId)

	switch {
	case err == errParentEntityDoesntExist:
		SendErrorResponse(serv.ResponseBuilder(), 404, fmt.Sprintf("No server exists with urlId %q.", serverUrlId))
		return
	case err == nil:
		SendErrorResponse(serv.ResponseBuilder(), 403, fmt.Sprintf("The urlId %q is already taken.", urlId))
		return
	case err != nil && err != errEntityDoesntExist:
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	}

	server, err := GetServer(serverUrlId)

	if err != nil {
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	}

	instance.Id = uuid.New()
	instance.ServerId = server.Id
	instance.UrlId = urlId

	_, insertErr := g_db.Exec("INSERT INTO instances (id, url_id, server_id, name) VALUES ($1, $2, $3, $4)", instance.Id, instance.UrlId, instance.ServerId, instance.Name)

	if insertErr != nil {
		SendInternalServerError(serv.ResponseBuilder(), insertErr)
		return
	}

	serv.ResponseBuilder().SetResponseCode(201).Overide(true)
}

func (serv ServersService) PutServer(server Server, urlId string) {
	_, err := GetServer(urlId)

	switch {
	case err != nil && err != errEntityDoesntExist:
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	case err == nil:
		SendErrorResponse(serv.ResponseBuilder(), 403, fmt.Sprintf("The urlId %q is already taken.", urlId))
		return
	}

	server.Id = uuid.New()
	server.UrlId = urlId

	_, insertErr := g_db.Exec("INSERT INTO servers (id, url_id, name) VALUES ($1, $2, $3)", server.Id, server.UrlId, server.Name)

	if insertErr != nil {
		SendInternalServerError(serv.ResponseBuilder(), insertErr)
	}

	serv.ResponseBuilder().SetResponseCode(201).Overide(true)
}

func (serv PlayersService) GetUser(id int64) Player {
	var steamid int64
	var username string

	err := g_db.QueryRow("SELECT id, name FROM players WHERE id = $1", id).Scan(&steamid, &username)

	switch {
	case err == sql.ErrNoRows:
		SendErrorResponse(serv.ResponseBuilder(), 404, "User not found.")
	case err != nil:
		SendInternalServerError(serv.ResponseBuilder(), err)
	}

	return Player{Id: steamid, Name: username}
}

func (serv PlayersService) PutUser(user Player, id int64) {

	if user.Id != id {
		SendErrorResponse(serv.ResponseBuilder(), 400, "Mismatched ids")
		return
	}

	var exists int
	err := g_db.QueryRow("SELECT count(*) FROM players WHERE id = $1", id).Scan(&exists)

	if err != nil {
		SendInternalServerError(serv.ResponseBuilder(), err)
		return
	}

	if exists > 0 {
		SendErrorResponse(serv.ResponseBuilder(), 400, fmt.Sprintf("User %v already exists", id))
		return
	}

	_, insertErr := g_db.Exec("INSERT INTO players (id, name) VALUES ($1, $2)", user.Id, user.Name)

	if insertErr != nil {
		SendInternalServerError(serv.ResponseBuilder(), insertErr)
		return
	}

	serv.ResponseBuilder().SetResponseCode(201).Overide(true)
}

func ReadConfig() (config Config) {
	file, err := os.Open("conf.json")
	defer file.Close()

	if err != nil {
		log.Fatal(err)
	}

	decoder := json.NewDecoder(file)
	config = Config{}
	err = decoder.Decode(&config)

	if err != nil {
		log.Fatal(err)
	}

	return config
}

func SendErrorResponse(responseBuilder *gorest.ResponseBuilder, responseCode int, message string) {
	responseBuilder.SetResponseCode(responseCode).WriteAndOveride([]byte(message))
}

func SendInternalServerError(responseBuilder *gorest.ResponseBuilder, err error) {
	SendErrorResponse(responseBuilder, 500, fmt.Sprintf("An unexpected error occurred: %v", err))
}

func GetServer(urlId string) (Server, error) {
	var id string
	var name string

	err := g_db.QueryRow("SELECT id, name FROM servers WHERE url_id = $1", urlId).Scan(&id, &name)

	switch {
	case err == sql.ErrNoRows:
		return Server{}, errEntityDoesntExist
	case err != nil:
		return Server{}, err
	}

	return Server{Id: id, UrlId: urlId, Name: name}, nil
}

func GetInstance(serverUrlId string, urlId string) (Instance, error) {
	server, err := GetServer(serverUrlId)

	switch {
	case err == errEntityDoesntExist:
		return Instance{}, errParentEntityDoesntExist
	case err != nil:
		return Instance{}, err
	}

	var id string
	var name string

	err = g_db.QueryRow("SELECT id, name FROM instances WHERE server_id = $1 AND url_id = $2", server.Id, urlId).Scan(&id, &name)

	switch {
	case err == sql.ErrNoRows:
		return Instance{}, errEntityDoesntExist
	case err != nil:
		return Instance{}, err
	}

	return Instance{Id: id, UrlId: urlId, ServerId: server.Id}, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//This file will be the future home for more policies
package gocql

//RetryableQuery is an interface that represents a query or batch statement that
//exposes the correct functions for the retry policy logic to evaluate correctly.
type RetryableQuery interface {
	Attempts() int
}

// RetryPolicy interace is used by gocql to determine if a query can be attempted
// again after a retryable error has been received. The interface allows gocql
// users to implement their own logic to determine if a query can be attempted
// again.
//
// See SimpleRetryPolicy as an example of implementing and using a RetryPolicy
// interface.
type RetryPolicy interface {
	Attempt(RetryableQuery) bool
}

/*
SimpleRetryPolicy has simple logic for attempting a query a fixed number of times.

See below for examples of usage:

	//Assign to the cluster
	cluster.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: 3}

	//Assign to a query
 	query.RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 1})
*/
type SimpleRetryPolicy struct {
	NumRetries int //Number of times to retry a query
}

// Attempt tells gocql to attempt the query again based on query.Attempts being less
// than the NumRetries defined in the policy.
func (s *SimpleRetryPolicy) Attempt(q RetryableQuery) bool {
	return q.Attempts() <= s.NumRetries
}
<commit_msg>Added missing expansion of RetryableQuery API.<commit_after>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//This file will be the future home for more policies
package gocql

//RetryableQuery is an interface that represents a query or batch statement that
//exposes the correct functions for the retry policy logic to evaluate correctly.
type RetryableQuery interface {
	Attempts() int
	GetConsistency() Consistency
}

// RetryPolicy interace is used by gocql to determine if a query can be attempted
// again after a retryable error has been received. The interface allows gocql
// users to implement their own logic to determine if a query can be attempted
// again.
//
// See SimpleRetryPolicy as an example of implementing and using a RetryPolicy
// interface.
type RetryPolicy interface {
	Attempt(RetryableQuery) bool
}

/*
SimpleRetryPolicy has simple logic for attempting a query a fixed number of times.

See below for examples of usage:

	//Assign to the cluster
	cluster.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: 3}

	//Assign to a query
 	query.RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 1})
*/
type SimpleRetryPolicy struct {
	NumRetries int //Number of times to retry a query
}

// Attempt tells gocql to attempt the query again based on query.Attempts being less
// than the NumRetries defined in the policy.
func (s *SimpleRetryPolicy) Attempt(q RetryableQuery) bool {
	return q.Attempts() <= s.NumRetries
}
<|endoftext|>


<commit_before>// Package cockroach implements the cockroach store
package cockroach

import (
	"database/sql"
	"fmt"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/lib/pq"
	"github.com/micro/go-micro/v2/store"
	"github.com/pkg/errors"
)

// DefaultDatabase is the namespace that the sql store
// will use if no namespace is provided.
var (
	DefaultDatabase = "micro"
	DefaultTable    = "micro"
)

var (
	re = regexp.MustCompile("[^a-zA-Z0-9]+")

	statements = map[string]string{
		"list":       "SELECT key, value, expiry FROM %s.%s;",
		"read":       "SELECT key, value, expiry FROM %s.%s WHERE key = $1;",
		"readMany":   "SELECT key, value, expiry FROM %s.%s WHERE key LIKE $1;",
		"readOffset": "SELECT key, value, expiry FROM %s.%s WHERE key LIKE $1 ORDER BY key DESC LIMIT $2 OFFSET $3;",
		"write":      "INSERT INTO %s.%s(key, value, expiry) VALUES ($1, $2::bytea, $3) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, expiry = EXCLUDED.expiry;",
		"delete":     "DELETE FROM %s.%s WHERE key = $1;",
	}
)

type sqlStore struct {
	options store.Options
	db      *sql.DB

	sync.RWMutex
	// known databases
	databases map[string]bool
}

func (s *sqlStore) getDB(database, table string) (string, string) {
	if len(database) == 0 {
		if len(s.options.Database) > 0 {
			database = s.options.Database
		} else {
			database = DefaultDatabase
		}
	}

	if len(table) == 0 {
		if len(s.options.Table) > 0 {
			table = s.options.Table
		} else {
			database = DefaultTable
		}
	}

	// store.namespace must only contain letters, numbers and underscores
	database = re.ReplaceAllString(database, "_")
	table = re.ReplaceAllString(table, "_")

	return database, table
}

func (s *sqlStore) createDB(database, table string) {
	database, table = s.getDB(database, table)

	s.Lock()
	_, ok := s.databases[database+":"+table]
	if !ok {
		s.initDB(database, table)
		s.databases[database+":"+table] = true
	}
	s.Unlock()
}

func (s *sqlStore) initDB(database, table string) error {
	// Create the namespace's database
	_, err := s.db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s;", database))
	if err != nil {
		return err
	}

	_, err = s.db.Exec(fmt.Sprintf("SET DATABASE = %s;", database))
	if err != nil {
		return errors.Wrap(err, "Couldn't set database")
	}

	// Create a table for the namespace's prefix
	_, err = s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s
	(
		key text NOT NULL,
		value bytea,
		expiry timestamp with time zone,
		CONSTRAINT %s_pkey PRIMARY KEY (key)
	);`, table, table))
	if err != nil {
		return errors.Wrap(err, "Couldn't create table")
	}

	// Create Index
	_, err = s.db.Exec(fmt.Sprintf(`CREATE INDEX IF NOT EXISTS "%s" ON %s.%s USING btree ("key");`, "key_index_"+table, database, table))
	if err != nil {
		return err
	}

	return nil
}

func (s *sqlStore) configure() error {
	if len(s.options.Nodes) == 0 {
		s.options.Nodes = []string{"postgresql://root@localhost:26257?sslmode=disable"}
	}

	source := s.options.Nodes[0]
	// check if it is a standard connection string eg: host=%s port=%d user=%s password=%s dbname=%s sslmode=disable
	// if err is nil which means it would be a URL like postgre://xxxx?yy=zz
	_, err := url.Parse(source)
	if err != nil {
		if !strings.Contains(source, " ") {
			source = fmt.Sprintf("host=%s", source)
		}
	}

	// create source from first node
	db, err := sql.Open("postgres", source)
	if err != nil {
		return err
	}

	if err := db.Ping(); err != nil {
		return err
	}

	if s.db != nil {
		s.db.Close()
	}

	// save the values
	s.db = db

	// get DB
	database, table := s.getDB(s.options.Database, s.options.Table)

	// initialise the database
	return s.initDB(database, table)
}

func (s *sqlStore) prepare(database, table, query string) (*sql.Stmt, error) {
	st, ok := statements[query]
	if !ok {
		return nil, errors.New("unsupported statement")
	}

	// get DB
	database, table = s.getDB(database, table)

	q := fmt.Sprintf(st, database, table)
	stmt, err := s.db.Prepare(q)
	if err != nil {
		return nil, err
	}
	return stmt, nil
}

func (s *sqlStore) Close() error {
	if s.db != nil {
		return s.db.Close()
	}
	return nil
}

func (s *sqlStore) Init(opts ...store.Option) error {
	for _, o := range opts {
		o(&s.options)
	}
	// reconfigure
	return s.configure()
}

// List all the known records
func (s *sqlStore) List(opts ...store.ListOption) ([]string, error) {
	var options store.ListOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	s.createDB(options.Database, options.Table)

	st, err := s.prepare(options.Database, options.Table, "list")
	if err != nil {
		return nil, err
	}
	defer st.Close()

	rows, err := st.Query()
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	defer rows.Close()

	var keys []string
	var timehelper pq.NullTime

	for rows.Next() {
		record := &store.Record{}
		if err := rows.Scan(&record.Key, &record.Value, &timehelper); err != nil {
			return keys, err
		}
		if timehelper.Valid {
			if timehelper.Time.Before(time.Now()) {
				// record has expired
				go s.Delete(record.Key)
			} else {
				record.Expiry = time.Until(timehelper.Time)
				keys = append(keys, record.Key)
			}
		} else {
			keys = append(keys, record.Key)
		}

	}
	rowErr := rows.Close()
	if rowErr != nil {
		// transaction rollback or something
		return keys, rowErr
	}
	if err := rows.Err(); err != nil {
		return keys, err
	}
	return keys, nil
}

// Read a single key
func (s *sqlStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var options store.ReadOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	s.createDB(options.Database, options.Table)

	if options.Prefix || options.Suffix {
		return s.read(key, options)
	}

	var records []*store.Record
	var timehelper pq.NullTime

	st, err := s.prepare(options.Database, options.Table, "read")
	if err != nil {
		return nil, err
	}
	defer st.Close()

	row := st.QueryRow(key)
	record := &store.Record{}
	if err := row.Scan(&record.Key, &record.Value, &timehelper); err != nil {
		if err == sql.ErrNoRows {
			return records, store.ErrNotFound
		}
		return records, err
	}
	if timehelper.Valid {
		if timehelper.Time.Before(time.Now()) {
			// record has expired
			go s.Delete(key)
			return records, store.ErrNotFound
		}
		record.Expiry = time.Until(timehelper.Time)
		records = append(records, record)
	} else {
		records = append(records, record)
	}

	return records, nil
}

// Read Many records
func (s *sqlStore) read(key string, options store.ReadOptions) ([]*store.Record, error) {
	pattern := "%"
	if options.Prefix {
		pattern = key + pattern
	}
	if options.Suffix {
		pattern = pattern + key
	}

	var rows *sql.Rows
	var err error

	if options.Limit != 0 {
		st, err := s.prepare(options.Database, options.Table, "readOffset")
		if err != nil {
			return nil, err
		}
		defer st.Close()

		rows, err = st.Query(pattern, options.Limit, options.Offset)
	} else {
		st, err := s.prepare(options.Database, options.Table, "readMany")
		if err != nil {
			return nil, err
		}
		defer st.Close()

		rows, err = st.Query(pattern)
	}
	if err != nil {
		if err == sql.ErrNoRows {
			return []*store.Record{}, nil
		}
		return []*store.Record{}, errors.Wrap(err, "sqlStore.read failed")
	}

	defer rows.Close()

	var records []*store.Record
	var timehelper pq.NullTime

	for rows.Next() {
		record := &store.Record{}
		if err := rows.Scan(&record.Key, &record.Value, &timehelper); err != nil {
			return records, err
		}
		if timehelper.Valid {
			if timehelper.Time.Before(time.Now()) {
				// record has expired
				go s.Delete(record.Key)
			} else {
				record.Expiry = time.Until(timehelper.Time)
				records = append(records, record)
			}
		} else {
			records = append(records, record)
		}
	}
	rowErr := rows.Close()
	if rowErr != nil {
		// transaction rollback or something
		return records, rowErr
	}
	if err := rows.Err(); err != nil {
		return records, err
	}

	return records, nil
}

// Write records
func (s *sqlStore) Write(r *store.Record, opts ...store.WriteOption) error {
	var options store.WriteOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	s.createDB(options.Database, options.Table)

	st, err := s.prepare(options.Database, options.Table, "write")
	if err != nil {
		return err
	}
	defer st.Close()

	if r.Expiry != 0 {
		_, err = st.Exec(r.Key, r.Value, time.Now().Add(r.Expiry))
	} else {
		_, err = st.Exec(r.Key, r.Value, nil)
	}

	if err != nil {
		return errors.Wrap(err, "Couldn't insert record "+r.Key)
	}

	return nil
}

// Delete records with keys
func (s *sqlStore) Delete(key string, opts ...store.DeleteOption) error {
	var options store.DeleteOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	s.createDB(options.Database, options.Table)

	st, err := s.prepare(options.Database, options.Table, "delete")
	if err != nil {
		return err
	}
	defer st.Close()

	result, err := st.Exec(key)
	if err != nil {
		return err
	}

	_, err = result.RowsAffected()
	if err != nil {
		return err
	}

	return nil
}

func (s *sqlStore) Options() store.Options {
	return s.options
}

func (s *sqlStore) String() string {
	return "cockroach"
}

// NewStore returns a new micro Store backed by sql
func NewStore(opts ...store.Option) store.Store {
	options := store.Options{
		Database: DefaultDatabase,
		Table:    DefaultTable,
	}

	for _, o := range opts {
		o(&options)
	}

	// new store
	s := new(sqlStore)
	// set the options
	s.options = options
	// mark known databases
	s.databases = make(map[string]bool)
	// best-effort configure the store
	s.configure()

	// return store
	return s
}
<commit_msg>Handle cockroach createDB error (#1603)<commit_after>// Package cockroach implements the cockroach store
package cockroach

import (
	"database/sql"
	"fmt"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/lib/pq"
	"github.com/micro/go-micro/v2/store"
	"github.com/pkg/errors"
)

// DefaultDatabase is the namespace that the sql store
// will use if no namespace is provided.
var (
	DefaultDatabase = "micro"
	DefaultTable    = "micro"
)

var (
	re = regexp.MustCompile("[^a-zA-Z0-9]+")

	statements = map[string]string{
		"list":       "SELECT key, value, expiry FROM %s.%s;",
		"read":       "SELECT key, value, expiry FROM %s.%s WHERE key = $1;",
		"readMany":   "SELECT key, value, expiry FROM %s.%s WHERE key LIKE $1;",
		"readOffset": "SELECT key, value, expiry FROM %s.%s WHERE key LIKE $1 ORDER BY key DESC LIMIT $2 OFFSET $3;",
		"write":      "INSERT INTO %s.%s(key, value, expiry) VALUES ($1, $2::bytea, $3) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, expiry = EXCLUDED.expiry;",
		"delete":     "DELETE FROM %s.%s WHERE key = $1;",
	}
)

type sqlStore struct {
	options store.Options
	db      *sql.DB

	sync.RWMutex
	// known databases
	databases map[string]bool
}

func (s *sqlStore) getDB(database, table string) (string, string) {
	if len(database) == 0 {
		if len(s.options.Database) > 0 {
			database = s.options.Database
		} else {
			database = DefaultDatabase
		}
	}

	if len(table) == 0 {
		if len(s.options.Table) > 0 {
			table = s.options.Table
		} else {
			database = DefaultTable
		}
	}

	// store.namespace must only contain letters, numbers and underscores
	database = re.ReplaceAllString(database, "_")
	table = re.ReplaceAllString(table, "_")

	return database, table
}

func (s *sqlStore) createDB(database, table string) error {
	database, table = s.getDB(database, table)

	s.Lock()
	defer s.Unlock()

	if _, ok := s.databases[database+":"+table]; ok {
		return nil
	}

	if err := s.initDB(database, table); err != nil {
		return err
	}

	s.databases[database+":"+table] = true
	return nil
}

func (s *sqlStore) initDB(database, table string) error {
	// Create the namespace's database
	_, err := s.db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s;", database))
	if err != nil {
		return err
	}

	_, err = s.db.Exec(fmt.Sprintf("SET DATABASE = %s;", database))
	if err != nil {
		return errors.Wrap(err, "Couldn't set database")
	}

	// Create a table for the namespace's prefix
	_, err = s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s
	(
		key text NOT NULL,
		value bytea,
		expiry timestamp with time zone,
		CONSTRAINT %s_pkey PRIMARY KEY (key)
	);`, table, table))
	if err != nil {
		return errors.Wrap(err, "Couldn't create table")
	}

	// Create Index
	_, err = s.db.Exec(fmt.Sprintf(`CREATE INDEX IF NOT EXISTS "%s" ON %s.%s USING btree ("key");`, "key_index_"+table, database, table))
	if err != nil {
		return err
	}

	return nil
}

func (s *sqlStore) configure() error {
	if len(s.options.Nodes) == 0 {
		s.options.Nodes = []string{"postgresql://root@localhost:26257?sslmode=disable"}
	}

	source := s.options.Nodes[0]
	// check if it is a standard connection string eg: host=%s port=%d user=%s password=%s dbname=%s sslmode=disable
	// if err is nil which means it would be a URL like postgre://xxxx?yy=zz
	_, err := url.Parse(source)
	if err != nil {
		if !strings.Contains(source, " ") {
			source = fmt.Sprintf("host=%s", source)
		}
	}

	// create source from first node
	db, err := sql.Open("postgres", source)
	if err != nil {
		return err
	}

	if err := db.Ping(); err != nil {
		return err
	}

	if s.db != nil {
		s.db.Close()
	}

	// save the values
	s.db = db

	// get DB
	database, table := s.getDB(s.options.Database, s.options.Table)

	// initialise the database
	return s.initDB(database, table)
}

func (s *sqlStore) prepare(database, table, query string) (*sql.Stmt, error) {
	st, ok := statements[query]
	if !ok {
		return nil, errors.New("unsupported statement")
	}

	// get DB
	database, table = s.getDB(database, table)

	q := fmt.Sprintf(st, database, table)
	stmt, err := s.db.Prepare(q)
	if err != nil {
		return nil, err
	}
	return stmt, nil
}

func (s *sqlStore) Close() error {
	if s.db != nil {
		return s.db.Close()
	}
	return nil
}

func (s *sqlStore) Init(opts ...store.Option) error {
	for _, o := range opts {
		o(&s.options)
	}
	// reconfigure
	return s.configure()
}

// List all the known records
func (s *sqlStore) List(opts ...store.ListOption) ([]string, error) {
	var options store.ListOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	if err := s.createDB(options.Database, options.Table); err != nil {
		return nil, err
	}

	st, err := s.prepare(options.Database, options.Table, "list")
	if err != nil {
		return nil, err
	}
	defer st.Close()

	rows, err := st.Query()
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	defer rows.Close()

	var keys []string
	var timehelper pq.NullTime

	for rows.Next() {
		record := &store.Record{}
		if err := rows.Scan(&record.Key, &record.Value, &timehelper); err != nil {
			return keys, err
		}
		if timehelper.Valid {
			if timehelper.Time.Before(time.Now()) {
				// record has expired
				go s.Delete(record.Key)
			} else {
				record.Expiry = time.Until(timehelper.Time)
				keys = append(keys, record.Key)
			}
		} else {
			keys = append(keys, record.Key)
		}

	}
	rowErr := rows.Close()
	if rowErr != nil {
		// transaction rollback or something
		return keys, rowErr
	}
	if err := rows.Err(); err != nil {
		return keys, err
	}
	return keys, nil
}

// Read a single key
func (s *sqlStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var options store.ReadOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	if err := s.createDB(options.Database, options.Table); err != nil {
		return nil, err
	}

	if options.Prefix || options.Suffix {
		return s.read(key, options)
	}

	var records []*store.Record
	var timehelper pq.NullTime

	st, err := s.prepare(options.Database, options.Table, "read")
	if err != nil {
		return nil, err
	}
	defer st.Close()

	row := st.QueryRow(key)
	record := &store.Record{}
	if err := row.Scan(&record.Key, &record.Value, &timehelper); err != nil {
		if err == sql.ErrNoRows {
			return records, store.ErrNotFound
		}
		return records, err
	}
	if timehelper.Valid {
		if timehelper.Time.Before(time.Now()) {
			// record has expired
			go s.Delete(key)
			return records, store.ErrNotFound
		}
		record.Expiry = time.Until(timehelper.Time)
		records = append(records, record)
	} else {
		records = append(records, record)
	}

	return records, nil
}

// Read Many records
func (s *sqlStore) read(key string, options store.ReadOptions) ([]*store.Record, error) {
	pattern := "%"
	if options.Prefix {
		pattern = key + pattern
	}
	if options.Suffix {
		pattern = pattern + key
	}

	var rows *sql.Rows
	var err error

	if options.Limit != 0 {
		st, err := s.prepare(options.Database, options.Table, "readOffset")
		if err != nil {
			return nil, err
		}
		defer st.Close()

		rows, err = st.Query(pattern, options.Limit, options.Offset)
	} else {
		st, err := s.prepare(options.Database, options.Table, "readMany")
		if err != nil {
			return nil, err
		}
		defer st.Close()

		rows, err = st.Query(pattern)
	}
	if err != nil {
		if err == sql.ErrNoRows {
			return []*store.Record{}, nil
		}
		return []*store.Record{}, errors.Wrap(err, "sqlStore.read failed")
	}

	defer rows.Close()

	var records []*store.Record
	var timehelper pq.NullTime

	for rows.Next() {
		record := &store.Record{}
		if err := rows.Scan(&record.Key, &record.Value, &timehelper); err != nil {
			return records, err
		}
		if timehelper.Valid {
			if timehelper.Time.Before(time.Now()) {
				// record has expired
				go s.Delete(record.Key)
			} else {
				record.Expiry = time.Until(timehelper.Time)
				records = append(records, record)
			}
		} else {
			records = append(records, record)
		}
	}
	rowErr := rows.Close()
	if rowErr != nil {
		// transaction rollback or something
		return records, rowErr
	}
	if err := rows.Err(); err != nil {
		return records, err
	}

	return records, nil
}

// Write records
func (s *sqlStore) Write(r *store.Record, opts ...store.WriteOption) error {
	var options store.WriteOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	if err := s.createDB(options.Database, options.Table); err != nil {
		return err
	}

	st, err := s.prepare(options.Database, options.Table, "write")
	if err != nil {
		return err
	}
	defer st.Close()

	if r.Expiry != 0 {
		_, err = st.Exec(r.Key, r.Value, time.Now().Add(r.Expiry))
	} else {
		_, err = st.Exec(r.Key, r.Value, nil)
	}

	if err != nil {
		return errors.Wrap(err, "Couldn't insert record "+r.Key)
	}

	return nil
}

// Delete records with keys
func (s *sqlStore) Delete(key string, opts ...store.DeleteOption) error {
	var options store.DeleteOptions
	for _, o := range opts {
		o(&options)
	}

	// create the db if not exists
	if err := s.createDB(options.Database, options.Table); err != nil {
		return err
	}

	st, err := s.prepare(options.Database, options.Table, "delete")
	if err != nil {
		return err
	}
	defer st.Close()

	result, err := st.Exec(key)
	if err != nil {
		return err
	}

	_, err = result.RowsAffected()
	if err != nil {
		return err
	}

	return nil
}

func (s *sqlStore) Options() store.Options {
	return s.options
}

func (s *sqlStore) String() string {
	return "cockroach"
}

// NewStore returns a new micro Store backed by sql
func NewStore(opts ...store.Option) store.Store {
	options := store.Options{
		Database: DefaultDatabase,
		Table:    DefaultTable,
	}

	for _, o := range opts {
		o(&options)
	}

	// new store
	s := new(sqlStore)
	// set the options
	s.options = options
	// mark known databases
	s.databases = make(map[string]bool)
	// best-effort configure the store
	s.configure()

	// return store
	return s
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"

	"github.com/mbrt/gmailfilter/pkg/config"
	"github.com/mbrt/gmailfilter/pkg/export"
)

func readConfig(path string) (config.Config, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return config.Config{}, errors.Wrap(err, fmt.Sprintf("cannot read %s", path))
	}

	var res config.Config
	err = yaml.Unmarshal(b, &res)
	return res, err

}

func errorf(format string, a ...interface{}) {
	/* #nosec */
	_, _ = fmt.Fprintf(os.Stderr, format, a...)
}

func fatal(format string, a ...interface{}) {
	errorf(format, a...)
	if !strings.HasSuffix(format, "\n") {
		errorf("") // Add newline
	}
	os.Exit(1)
}

func main() {
	flag.Parse()
	if flag.NArg() < 1 {
		fatal("usage: %s <config-file>", os.Args[0])
	}
	cfg, err := readConfig(flag.Arg(0))
	if err != nil {
		fatal("error in config parse: %s", err)
	}

	rules, err := export.GenerateRules(cfg)
	if err != nil {
		fatal("error generating rules: %s", err)
	}

	err = export.DefaultXMLExporter().MarshalEntries(cfg.Author, rules, os.Stdout)
	if err != nil {
		fatal("error exporting to XML: %s", err)
	}
}
<commit_msg>fix lint<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"

	"github.com/mbrt/gmailfilter/pkg/config"
	"github.com/mbrt/gmailfilter/pkg/export"
)

func readConfig(path string) (config.Config, error) {
	/* #nosec */
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return config.Config{}, errors.Wrap(err, fmt.Sprintf("cannot read %s", path))
	}

	var res config.Config
	err = yaml.Unmarshal(b, &res)
	return res, err

}

func errorf(format string, a ...interface{}) {
	/* #nosec */
	_, _ = fmt.Fprintf(os.Stderr, format, a...)
}

func fatal(format string, a ...interface{}) {
	errorf(format, a...)
	if !strings.HasSuffix(format, "\n") {
		errorf("") // Add newline
	}
	os.Exit(1)
}

func main() {
	flag.Parse()
	if flag.NArg() < 1 {
		fatal("usage: %s <config-file>", os.Args[0])
	}
	cfg, err := readConfig(flag.Arg(0))
	if err != nil {
		fatal("error in config parse: %s", err)
	}

	rules, err := export.GenerateRules(cfg)
	if err != nil {
		fatal("error generating rules: %s", err)
	}

	err = export.DefaultXMLExporter().MarshalEntries(cfg.Author, rules, os.Stdout)
	if err != nil {
		fatal("error exporting to XML: %s", err)
	}
}
<|endoftext|>