<commit_before>package views

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"testing"
)

func TestUploadedFileHandler(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(UploadedFileHandler))
	defer ts.Close()
	os.Mkdir("./files", 0777)
	defer os.RemoveAll("./files")
	file, err := os.Create("./files/testfile")
	defer file.Close()
	defer os.Remove(file.Name())
	expectedContent := []byte("test content")
	file.Write(expectedContent)
	res, err := http.Get(ts.URL + "/files/testfile")
	if err != nil {
		t.Error("Error occured while getting response from test server:", err)
	}
	actualContent, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		t.Error("Error occured while reading content from response Body: ", err)
	}

	if !bytes.Equal(actualContent, expectedContent) {
		t.Errorf("Actual content (%s) did not match expected content (%s)", actualContent, expectedContent)
	}
}

//TestAddEmptyCategory tests that if the category field is empty it should do nothing
func TestAddEmptyCategory(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(AddCategoryFunc))
	defer ts.Close()
	req, err := http.NewRequest("POST", ts.URL, nil)
	req.Form, _ = url.ParseQuery("category=")
	if err != nil {
		t.Errorf("Error occured while constracting request:%s", err)
	}
	w := httptest.NewRecorder()
	AddCategoryFunc(w, req)
	body := w.Body.String()
	if len(body) != 0 {
		t.Error("Body should be empty. Instead contained data: ", body)
	}
}

func TestEditTaskWithWrongMethod(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(EditTaskFunc))
	defer ts.Close()
	req, err := http.NewRequest("OPTIONS", ts.URL, nil)
	if err != nil {
		t.Errorf("Error occured while constracting request:%s", err)
	}
	w := httptest.NewRecorder()
	EditTaskFunc(w, req)

	if w.Code != http.StatusFound && message != "Method not allowed" {
		t.Errorf("Message was: %s Return code was: %d. Should have been message: %s return code: %d", message, w.Code, "Method not allowed", http.StatusFound)
	}
}

func TestEditTaskWrongTaskName(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(EditTaskFunc))
	defer ts.Close()
	ts.URL = ts.URL + "/edit/invalidID"
	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Errorf("Error occured while constracting request:%s", err)
	}
	w := httptest.NewRecorder()
	EditTaskFunc(w, req)
	//TODO: Error should be returned as part of the string so that the message can tell
	//why there was a problem.
	if w.Code != http.StatusBadRequest {
		t.Errorf("Actual status: (%d); Expected status:(%d)", w.Code, http.StatusBadRequest)
	}
}
<commit_msg>Form is just a map[string][]string. It is easy to manipulate. But url.ParseQuery can also be used.<commit_after>package views

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
)

func TestUploadedFileHandler(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(UploadedFileHandler))
	defer ts.Close()
	os.Mkdir("./files", 0777)
	defer os.RemoveAll("./files")
	file, err := os.Create("./files/testfile")
	defer file.Close()
	defer os.Remove(file.Name())
	expectedContent := []byte("test content")
	file.Write(expectedContent)
	res, err := http.Get(ts.URL + "/files/testfile")
	if err != nil {
		t.Error("Error occured while getting response from test server:", err)
	}
	actualContent, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		t.Error("Error occured while reading content from response Body: ", err)
	}

	if !bytes.Equal(actualContent, expectedContent) {
		t.Errorf("Actual content (%s) did not match expected content (%s)", actualContent, expectedContent)
	}
}

//TestAddEmptyCategory tests that if the category field is empty it should do nothing
func TestAddEmptyCategory(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(AddCategoryFunc))
	defer ts.Close()
	req, err := http.NewRequest("POST", ts.URL, nil)
	req.Form = make(map[string][]string, 0)
	req.Form.Add("category", "")
	// req.Form, _ = url.ParseQuery("category=")
	if err != nil {
		t.Errorf("Error occured while constracting request:%s", err)
	}
	w := httptest.NewRecorder()
	AddCategoryFunc(w, req)
	body := w.Body.String()
	if len(body) != 0 {
		t.Error("Body should be empty. Instead contained data: ", body)
	}
}

func TestEditTaskWithWrongMethod(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(EditTaskFunc))
	defer ts.Close()
	req, err := http.NewRequest("OPTIONS", ts.URL, nil)
	if err != nil {
		t.Errorf("Error occured while constracting request:%s", err)
	}
	w := httptest.NewRecorder()
	EditTaskFunc(w, req)

	if w.Code != http.StatusFound && message != "Method not allowed" {
		t.Errorf("Message was: %s Return code was: %d. Should have been message: %s return code: %d", message, w.Code, "Method not allowed", http.StatusFound)
	}
}

func TestEditTaskWrongTaskName(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(EditTaskFunc))
	defer ts.Close()
	ts.URL = ts.URL + "/edit/invalidID"
	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Errorf("Error occured while constracting request:%s", err)
	}
	w := httptest.NewRecorder()
	EditTaskFunc(w, req)
	//TODO: Error should be returned as part of the string so that the message can tell
	//why there was a problem.
	if w.Code != http.StatusBadRequest {
		t.Errorf("Actual status: (%d); Expected status:(%d)", w.Code, http.StatusBadRequest)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 maps

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
	"googlemaps.github.io/maps/internal"
)

// Client may be used to make requests to the Google Maps WebService APIs
type Client struct {
	httpClient        *http.Client
	apiKey            string
	baseURL           string
	clientID          string
	signature         []byte
	requestsPerSecond int
	rateLimiter       chan int
}

// ClientOption is the type of constructor options for NewClient(...).
type ClientOption func(*Client) error

var defaultRequestsPerSecond = 10

// NewClient constructs a new Client which can make requests to the Google Maps WebService APIs.
func NewClient(options ...ClientOption) (*Client, error) {
	c := &Client{requestsPerSecond: defaultRequestsPerSecond}
	WithHTTPClient(&http.Client{})(c)
	for _, option := range options {
		err := option(c)
		if err != nil {
			return nil, err
		}
	}
	if c.apiKey == "" && (c.clientID == "" || len(c.signature) == 0) {
		return nil, errors.New("maps: API Key or Maps for Work credentials missing")
	}

	if c.requestsPerSecond > 0 {
		// Implement a bursty rate limiter.
		// Allow up to 1 second worth of requests to be made at once.
		c.rateLimiter = make(chan int, c.requestsPerSecond)
		// Prefill rateLimiter with 1 seconds worth of requests.
		for i := 0; i < c.requestsPerSecond; i++ {
			c.rateLimiter <- 1
		}
		go func() {
			// Wait a second for pre-filled quota to drain
			time.Sleep(time.Second)
			// Then, refill rateLimiter continuously
			for _ = range time.Tick(time.Second / time.Duration(c.requestsPerSecond)) {
				c.rateLimiter <- 1
			}
		}()
	}

	return c, nil
}

// WithHTTPClient configures a Maps API client with a http.Client to make requests over.
func WithHTTPClient(c *http.Client) ClientOption {
	return func(client *Client) error {
		if _, ok := c.Transport.(*transport); !ok {
			t := c.Transport
			if t != nil {
				c.Transport = &transport{Base: t}
			} else {
				c.Transport = &transport{Base: http.DefaultTransport}
			}
		}
		client.httpClient = c
		return nil
	}
}

// WithAPIKey configures a Maps API client with an API Key
func WithAPIKey(apiKey string) ClientOption {
	return func(c *Client) error {
		c.apiKey = apiKey
		return nil
	}
}

// WithClientIDAndSignature configures a Maps API client for a Maps for Work application
// The signature is assumed to be URL modified Base64 encoded
func WithClientIDAndSignature(clientID, signature string) ClientOption {
	return func(c *Client) error {
		c.clientID = clientID
		decoded, err := base64.URLEncoding.DecodeString(signature)
		if err != nil {
			return err
		}
		c.signature = decoded
		return nil
	}
}

// WithRateLimit configures the rate limit for back end requests. Default is to
// limit to 10 requests per second. A value of zero disables rate limiting.
func WithRateLimit(requestsPerSecond int) ClientOption {
	return func(c *Client) error {
		c.requestsPerSecond = requestsPerSecond
		return nil
	}
}

type apiConfig struct {
	host            string
	path            string
	acceptsClientID bool
}

type apiRequest interface {
	params() url.Values
}

func (c *Client) awaitRateLimiter(ctx context.Context) error {
	if c.rateLimiter == nil {
		return nil
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-c.rateLimiter:
		// Execute request.
		return nil
	}
}

func (c *Client) get(ctx context.Context, config *apiConfig, apiReq apiRequest) (*http.Response, error) {
	if err := c.awaitRateLimiter(ctx); err != nil {
		return nil, err
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}
	req, err := http.NewRequest("GET", host+config.path, nil)
	if err != nil {
		return nil, err
	}
	q, err := c.generateAuthQuery(config.path, apiReq.params(), config.acceptsClientID)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) post(ctx context.Context, config *apiConfig, apiReq interface{}) (*http.Response, error) {
	if err := c.awaitRateLimiter(ctx); err != nil {
		return nil, err
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}

	body, err := json.Marshal(apiReq)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", host+config.path, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	q, err := c.generateAuthQuery(config.path, url.Values{}, config.acceptsClientID)
	if err != nil {
		return nil, err
	}

	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) getJSON(ctx context.Context, config *apiConfig, apiReq apiRequest, resp interface{}) error {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

func (c *Client) postJSON(ctx context.Context, config *apiConfig, apiReq interface{}, resp interface{}) error {
	httpResp, err := c.post(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

type binaryResponse struct {
	statusCode  int
	contentType string
	data        io.ReadCloser
}

func (c *Client) getBinary(ctx context.Context, config *apiConfig, apiReq apiRequest) (binaryResponse, error) {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return binaryResponse{}, err
	}

	return binaryResponse{httpResp.StatusCode, httpResp.Header.Get("Content-Type"), httpResp.Body}, nil
}

func (c *Client) generateAuthQuery(path string, q url.Values, acceptClientID bool) (string, error) {
	if c.apiKey != "" {
		q.Set("key", c.apiKey)
		return q.Encode(), nil
	}
	if acceptClientID {
		return internal.SignURL(path, c.clientID, c.signature, q)
	}
	return "", errors.New("maps: API Key missing")
}

// commonResponse contains the common response fields to most API calls inside
// the Google Maps APIs. This is used internally.
type commonResponse struct {
	// Status contains the status of the request, and may contain debugging
	// information to help you track down why the call failed.
	Status string `json:"status"`

	// ErrorMessage is the explanatory field added when Status is an error.
	ErrorMessage string `json:"error_message"`
}

// StatusError returns an error iff this object has a non-OK Status.
func (c *commonResponse) StatusError() error {
	if c.Status != "OK" {
		return fmt.Errorf("maps: %s - %s", c.Status, c.ErrorMessage)
	}
	return nil
}
<commit_msg>Include Channel on Client<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 maps

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
	"googlemaps.github.io/maps/internal"
)

// Client may be used to make requests to the Google Maps WebService APIs
type Client struct {
	httpClient        *http.Client
	apiKey            string
	baseURL           string
	clientID          string
	signature         []byte
	requestsPerSecond int
	rateLimiter       chan int
	channel           string
}

// ClientOption is the type of constructor options for NewClient(...).
type ClientOption func(*Client) error

var defaultRequestsPerSecond = 10

// NewClient constructs a new Client which can make requests to the Google Maps WebService APIs.
func NewClient(options ...ClientOption) (*Client, error) {
	c := &Client{requestsPerSecond: defaultRequestsPerSecond}
	WithHTTPClient(&http.Client{})(c)
	for _, option := range options {
		err := option(c)
		if err != nil {
			return nil, err
		}
	}
	if c.apiKey == "" && (c.clientID == "" || len(c.signature) == 0) {
		return nil, errors.New("maps: API Key or Maps for Work credentials missing")
	}

	if c.requestsPerSecond > 0 {
		// Implement a bursty rate limiter.
		// Allow up to 1 second worth of requests to be made at once.
		c.rateLimiter = make(chan int, c.requestsPerSecond)
		// Prefill rateLimiter with 1 seconds worth of requests.
		for i := 0; i < c.requestsPerSecond; i++ {
			c.rateLimiter <- 1
		}
		go func() {
			// Wait a second for pre-filled quota to drain
			time.Sleep(time.Second)
			// Then, refill rateLimiter continuously
			for _ = range time.Tick(time.Second / time.Duration(c.requestsPerSecond)) {
				c.rateLimiter <- 1
			}
		}()
	}

	return c, nil
}

// WithHTTPClient configures a Maps API client with a http.Client to make requests over.
func WithHTTPClient(c *http.Client) ClientOption {
	return func(client *Client) error {
		if _, ok := c.Transport.(*transport); !ok {
			t := c.Transport
			if t != nil {
				c.Transport = &transport{Base: t}
			} else {
				c.Transport = &transport{Base: http.DefaultTransport}
			}
		}
		client.httpClient = c
		return nil
	}
}

// WithAPIKey configures a Maps API client with an API Key
func WithAPIKey(apiKey string) ClientOption {
	return func(c *Client) error {
		c.apiKey = apiKey
		return nil
	}
}

// WithChannel configures a Maps API client with a Channel
func WithChannel(channel string) ClientOption {
	return func(c *Client) error {
		c.channel = channel
		return nil
	}
}

// WithClientIDAndSignature configures a Maps API client for a Maps for Work application
// The signature is assumed to be URL modified Base64 encoded
func WithClientIDAndSignature(clientID, signature string) ClientOption {
	return func(c *Client) error {
		c.clientID = clientID
		decoded, err := base64.URLEncoding.DecodeString(signature)
		if err != nil {
			return err
		}
		c.signature = decoded
		return nil
	}
}

// WithRateLimit configures the rate limit for back end requests. Default is to
// limit to 10 requests per second. A value of zero disables rate limiting.
func WithRateLimit(requestsPerSecond int) ClientOption {
	return func(c *Client) error {
		c.requestsPerSecond = requestsPerSecond
		return nil
	}
}

type apiConfig struct {
	host            string
	path            string
	acceptsClientID bool
}

type apiRequest interface {
	params() url.Values
}

func (c *Client) awaitRateLimiter(ctx context.Context) error {
	if c.rateLimiter == nil {
		return nil
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-c.rateLimiter:
		// Execute request.
		return nil
	}
}

func (c *Client) get(ctx context.Context, config *apiConfig, apiReq apiRequest) (*http.Response, error) {
	if err := c.awaitRateLimiter(ctx); err != nil {
		return nil, err
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}
	req, err := http.NewRequest("GET", host+config.path, nil)
	if err != nil {
		return nil, err
	}
	q, err := c.generateAuthQuery(config.path, apiReq.params(), config.acceptsClientID)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) post(ctx context.Context, config *apiConfig, apiReq interface{}) (*http.Response, error) {
	if err := c.awaitRateLimiter(ctx); err != nil {
		return nil, err
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}

	body, err := json.Marshal(apiReq)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", host+config.path, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	q, err := c.generateAuthQuery(config.path, url.Values{}, config.acceptsClientID)
	if err != nil {
		return nil, err
	}

	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) getJSON(ctx context.Context, config *apiConfig, apiReq apiRequest, resp interface{}) error {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

func (c *Client) postJSON(ctx context.Context, config *apiConfig, apiReq interface{}, resp interface{}) error {
	httpResp, err := c.post(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

type binaryResponse struct {
	statusCode  int
	contentType string
	data        io.ReadCloser
}

func (c *Client) getBinary(ctx context.Context, config *apiConfig, apiReq apiRequest) (binaryResponse, error) {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return binaryResponse{}, err
	}

	return binaryResponse{httpResp.StatusCode, httpResp.Header.Get("Content-Type"), httpResp.Body}, nil
}

func (c *Client) generateAuthQuery(path string, q url.Values, acceptClientID bool) (string, error) {
	if c.channel != "" {
		q.Set("channel", c.channel)
	}
	if c.apiKey != "" {
		q.Set("key", c.apiKey)
		return q.Encode(), nil
	}
	if acceptClientID {
		return internal.SignURL(path, c.clientID, c.signature, q)
	}
	return "", errors.New("maps: API Key missing")
}

// commonResponse contains the common response fields to most API calls inside
// the Google Maps APIs. This is used internally.
type commonResponse struct {
	// Status contains the status of the request, and may contain debugging
	// information to help you track down why the call failed.
	Status string `json:"status"`

	// ErrorMessage is the explanatory field added when Status is an error.
	ErrorMessage string `json:"error_message"`
}

// StatusError returns an error iff this object has a non-OK Status.
func (c *commonResponse) StatusError() error {
	if c.Status != "OK" {
		return fmt.Errorf("maps: %s - %s", c.Status, c.ErrorMessage)
	}
	return nil
}
<|endoftext|>


<commit_before>package gowbem

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"

	"github.com/vmware/govmomi/vim25/progress"
)

type HasFault interface {
	Fault() error
}

type DecodeError struct {
	bytes []byte
	err   error
}

func (f *DecodeError) Error() string {
	return f.err.Error() + ", xml as follow:\r\n" + string(f.bytes)
}

type FaultError struct {
	bytes []byte
	err   error
}

func (f *FaultError) Error() string {
	return f.err.Error() + ", xml as follow:\r\n" + string(f.bytes)
}

type RoundTripper interface {
	RoundTrip(action string, reqBody interface{}, resBody HasFault, cached *bytes.Buffer) error
}

var cn uint64 // Client counter

type Client struct {
	rn uint64 // Request counter

	http.Client

	u        url.URL
	insecure bool

	cn_str string // Client counter
	cn     uint64 // Client counter
	cached *bytes.Buffer
}

func NewClient(u *url.URL, insecure bool) *Client {
	c := &Client{}
	c.init(u, insecure)
	return c
}

func (c *Client) init(u *url.URL, insecure bool) {
	c.u = *u
	c.insecure = insecure
	c.cn = atomic.AddUint64(&cn, 1)
	c.rn = 0
	c.cached = bytes.NewBuffer(make([]byte, 0, 8*1024*1024))

	c.cn_str = strconv.FormatUint(c.cn, 10)

	if c.u.Scheme == "https" {
		c.Client.Transport = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: c.insecure}}
	}

	c.Jar, _ = cookiejar.New(nil)
	//c.u.User = nil
}

func (c *Client) generateId() string {
	return c.cn_str + "-" + GenerateId()
}

func (c *Client) URL() url.URL {
	return c.u
}

type marshaledClient struct {
	Cookies  []*http.Cookie
	URL      *url.URL
	Insecure bool
}

func (c *Client) MarshalJSON() ([]byte, error) {
	m := marshaledClient{
		Cookies:  c.Jar.Cookies(&c.u),
		URL:      &c.u,
		Insecure: c.insecure,
	}

	return json.Marshal(m)
}

func (c *Client) UnmarshalJSON(b []byte) error {
	var m marshaledClient

	err := json.Unmarshal(b, &m)
	if err != nil {
		return err
	}

	*c = *NewClient(m.URL, m.Insecure)
	c.Jar.SetCookies(m.URL, m.Cookies)
	return nil
}

func (c *Client) RoundTrip(action string, headers map[string]string, reqBody interface{}, resBody HasFault) error {
	var httpreq *http.Request
	var httpres *http.Response
	var dumpWriter io.WriteCloser
	var err error

	num := atomic.AddUint64(&c.rn, 1)
	c.cached.Reset()

	c.cached.WriteString(xml.Header)

	if err = xml.NewEncoder(c.cached).Encode(reqBody); err != nil {
		panic(err)
	}
	rawreqbody := c.cached

	httpreq, err = http.NewRequest(action, c.u.String(), rawreqbody)
	if err != nil {
		panic(err)
	}
	//if pwd, ok := c.u.User.Password(); ok {
	//	httpreq.SetBasicAuth(c.u.User.Username(), pwd)
	//}

	httpreq.Header.Set(`Content-Type`, `text/xml; charset="utf-8"`)
	if 0 != len(headers) {
		for k, v := range headers {
			httpreq.Header.Set(k, v)
		}
	}
	//httpreq.Header.Set(`SOAPAction`, `urn:vim25/5.5`)

	if DebugEnabled() {
		b, _ := httputil.DumpRequest(httpreq, false)
		dumpWriter = DebugNewFile(fmt.Sprintf("%d-%04d.log", c.cn, num))
		defer dumpWriter.Close()
		dumpWriter.Write(b)
		dumpWriter.Write(c.cached.Bytes())
	}

	//tstart := time.Now()
	httpres, err = c.Client.Do(httpreq)
	//tstop := time.Now()

	// if DebugEnabled() {
	// 	now := time.Now().Format("2006-01-02T15-04-05.999999999")
	// 	ms := tstop.Sub(tstart) / time.Millisecond
	// 	fmt.Fprintf(c.log, "%s: %4d took %6dms\n", now, num, ms)
	// }

	if err != nil {
		return err
	}

	//var rawresbody io.Reader = httpres.Body
	defer httpres.Body.Close()

	if httpres.ContentLength <= 0 && httpres.StatusCode < http.StatusOK && httpres.StatusCode >= http.StatusMultipleChoices {
		cimError := httpres.Header.Get("CIMError")
		errorDetail := httpres.Header.Get("PGErrorDetail")
		if "" == cimError {
			if "" != errorDetail {
				return errors.New(errorDetail)
			}
			return errors.New(httpres.Status)
		} else {
			if "" != errorDetail {
				return errors.New(cimError + ": " + errorDetail)
			}
			return errors.New(cimError)
		}
	}

	c.cached.Reset()
	if _, err = io.Copy(c.cached, httpres.Body); nil != err {
		return err
	}

	if DebugEnabled() {
		b, _ := httputil.DumpResponse(httpres, false)
		dumpWriter.Write(b)
		dumpWriter.Write(c.cached.Bytes())
	}

	if 200 != httpres.StatusCode {
		if 0 == c.cached.Len() {
			return errors.New(httpres.Status)
		}
		return errors.New(httpres.Status + ":" + c.cached.String())
	}

	dec := xml.NewDecoder(bytes.NewReader(c.cached.Bytes()))
	err = dec.Decode(resBody)
	if err != nil {
		return &DecodeError{bytes: c.cached.Bytes(), err: err}
	}

	if fault := resBody.Fault(); fault != nil {
		return &FaultError{bytes: c.cached.Bytes(), err: fault}
	}

	return nil
}

// ParseURL wraps url.Parse to rewrite the URL.Host field
// In the case of VM guest uploads or NFC lease URLs, a Host
// field with a value of "*" is rewritten to the Client's URL.Host.
func (c *Client) ParseURL(urlStr string) (*url.URL, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	host := strings.Split(u.Host, ":")
	if host[0] == "*" {
		// Also use Client's port, to support port forwarding
		u.Host = c.URL().Host
	}

	return u, nil
}

type Upload struct {
	Type          string
	Method        string
	ContentLength int64
	Progress      progress.Sinker
}

var DefaultUpload = Upload{
	Type:   "application/octet-stream",
	Method: "PUT",
}

// Upload PUTs the local file to the given URL
func (c *Client) Upload(f io.Reader, u *url.URL, param *Upload) error {
	var err error

	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, f, param.ContentLength)
		f = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	req, err := http.NewRequest(param.Method, u.String(), f)
	if err != nil {
		return err
	}

	req.ContentLength = param.ContentLength
	req.Header.Set("Content-Type", param.Type)

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	switch res.StatusCode {
	case http.StatusOK:
	case http.StatusCreated:
	default:
		err = errors.New(res.Status)
	}

	return err
}

// UploadFile PUTs the local file to the given URL
func (c *Client) UploadFile(file string, u *url.URL, param *Upload) error {
	if param == nil {
		p := DefaultUpload // Copy since we set ContentLength
		param = &p
	}

	s, err := os.Stat(file)
	if err != nil {
		return err
	}

	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	param.ContentLength = s.Size()

	return c.Upload(f, u, param)
}

type Download struct {
	Method   string
	Progress progress.Sinker
}

var DefaultDownload = Download{
	Method: "GET",
}

// DownloadFile GETs the given URL to a local file
func (c *Client) DownloadFile(file string, u *url.URL, param *Download) error {
	var err error

	if param == nil {
		param = &DefaultDownload
	}

	fh, err := os.Create(file)
	if err != nil {
		return err
	}
	defer fh.Close()

	req, err := http.NewRequest(param.Method, u.String(), nil)
	if err != nil {
		return err
	}

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	defer res.Body.Close()

	switch res.StatusCode {
	case http.StatusOK:
	default:
		err = errors.New(res.Status)
	}

	if err != nil {
		return err
	}

	var r io.Reader = res.Body
	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, res.Body, res.ContentLength)
		r = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	_, err = io.Copy(fh, r)
	if err != nil {
		return err
	}

	// Assign error before returning so that it gets picked up by the deferred
	// function marking the progress reader as done.
	err = fh.Close()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>修复PG的错误响应中没有长度也不关闭 tcp 连接导致方法调用不返回的问题<commit_after>package gowbem

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"

	"github.com/vmware/govmomi/vim25/progress"
)

type HasFault interface {
	Fault() error
}

type DecodeError struct {
	bytes []byte
	err   error
}

func (f *DecodeError) Error() string {
	return f.err.Error() + ", xml as follow:\r\n" + string(f.bytes)
}

type FaultError struct {
	bytes []byte
	err   error
}

func (f *FaultError) Error() string {
	return f.err.Error() + ", xml as follow:\r\n" + string(f.bytes)
}

type RoundTripper interface {
	RoundTrip(action string, reqBody interface{}, resBody HasFault, cached *bytes.Buffer) error
}

var cn uint64 // Client counter

type Client struct {
	rn uint64 // Request counter

	http.Client

	u        url.URL
	insecure bool

	cn_str string // Client counter
	cn     uint64 // Client counter
	cached *bytes.Buffer
}

func NewClient(u *url.URL, insecure bool) *Client {
	c := &Client{}
	c.init(u, insecure)
	return c
}

func (c *Client) init(u *url.URL, insecure bool) {
	c.u = *u
	c.insecure = insecure
	c.cn = atomic.AddUint64(&cn, 1)
	c.rn = 0
	c.cached = bytes.NewBuffer(make([]byte, 0, 8*1024*1024))

	c.cn_str = strconv.FormatUint(c.cn, 10)

	if c.u.Scheme == "https" {
		c.Client.Transport = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: c.insecure}}
	}

	c.Jar, _ = cookiejar.New(nil)
	//c.u.User = nil
}

func (c *Client) generateId() string {
	return c.cn_str + "-" + GenerateId()
}

func (c *Client) URL() url.URL {
	return c.u
}

type marshaledClient struct {
	Cookies  []*http.Cookie
	URL      *url.URL
	Insecure bool
}

func (c *Client) MarshalJSON() ([]byte, error) {
	m := marshaledClient{
		Cookies:  c.Jar.Cookies(&c.u),
		URL:      &c.u,
		Insecure: c.insecure,
	}

	return json.Marshal(m)
}

func (c *Client) UnmarshalJSON(b []byte) error {
	var m marshaledClient

	err := json.Unmarshal(b, &m)
	if err != nil {
		return err
	}

	*c = *NewClient(m.URL, m.Insecure)
	c.Jar.SetCookies(m.URL, m.Cookies)
	return nil
}

func (c *Client) RoundTrip(action string, headers map[string]string, reqBody interface{}, resBody HasFault) error {
	var httpreq *http.Request
	var httpres *http.Response
	var dumpWriter io.WriteCloser
	var err error

	num := atomic.AddUint64(&c.rn, 1)
	c.cached.Reset()

	c.cached.WriteString(xml.Header)

	if err = xml.NewEncoder(c.cached).Encode(reqBody); err != nil {
		panic(err)
	}
	rawreqbody := c.cached

	httpreq, err = http.NewRequest(action, c.u.String(), rawreqbody)
	if err != nil {
		panic(err)
	}
	//if pwd, ok := c.u.User.Password(); ok {
	//	httpreq.SetBasicAuth(c.u.User.Username(), pwd)
	//}

	httpreq.Header.Set(`Content-Type`, `text/xml; charset="utf-8"`)
	if 0 != len(headers) {
		for k, v := range headers {
			httpreq.Header.Set(k, v)
		}
	}
	//httpreq.Header.Set(`SOAPAction`, `urn:vim25/5.5`)

	if DebugEnabled() {
		b, _ := httputil.DumpRequest(httpreq, false)
		dumpWriter = DebugNewFile(fmt.Sprintf("%d-%04d.log", c.cn, num))
		defer dumpWriter.Close()
		dumpWriter.Write(b)
		dumpWriter.Write(c.cached.Bytes())
	}

	//tstart := time.Now()
	httpres, err = c.Client.Do(httpreq)
	//tstop := time.Now()

	// if DebugEnabled() {
	// 	now := time.Now().Format("2006-01-02T15-04-05.999999999")
	// 	ms := tstop.Sub(tstart) / time.Millisecond
	// 	fmt.Fprintf(c.log, "%s: %4d took %6dms\n", now, num, ms)
	// }

	if err != nil {
		return err
	}

	//var rawresbody io.Reader = httpres.Body
	defer httpres.Body.Close()

	if httpres.ContentLength <= 0 && (httpres.StatusCode < http.StatusOK || httpres.StatusCode >= http.StatusMultipleChoices) {
		cimError := httpres.Header.Get("CIMError")
		errorDetail := httpres.Header.Get("PGErrorDetail")
		if "" == cimError {
			if "" != errorDetail {
				return errors.New(errorDetail)
			}
			return errors.New(httpres.Status)
		} else {
			if "" != errorDetail {
				return errors.New(cimError + ": " + errorDetail)
			}
			return errors.New(cimError)
		}
	}

	c.cached.Reset()
	if _, err = io.Copy(c.cached, httpres.Body); nil != err {
		return err
	}

	if DebugEnabled() {
		b, _ := httputil.DumpResponse(httpres, false)
		dumpWriter.Write(b)
		dumpWriter.Write(c.cached.Bytes())
	}

	if 200 != httpres.StatusCode {
		if 0 == c.cached.Len() {
			return errors.New(httpres.Status)
		}
		return errors.New(httpres.Status + ":" + c.cached.String())
	}

	dec := xml.NewDecoder(bytes.NewReader(c.cached.Bytes()))
	err = dec.Decode(resBody)
	if err != nil {
		return &DecodeError{bytes: c.cached.Bytes(), err: err}
	}

	if fault := resBody.Fault(); fault != nil {
		return &FaultError{bytes: c.cached.Bytes(), err: fault}
	}

	return nil
}

// ParseURL wraps url.Parse to rewrite the URL.Host field
// In the case of VM guest uploads or NFC lease URLs, a Host
// field with a value of "*" is rewritten to the Client's URL.Host.
func (c *Client) ParseURL(urlStr string) (*url.URL, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	host := strings.Split(u.Host, ":")
	if host[0] == "*" {
		// Also use Client's port, to support port forwarding
		u.Host = c.URL().Host
	}

	return u, nil
}

type Upload struct {
	Type          string
	Method        string
	ContentLength int64
	Progress      progress.Sinker
}

var DefaultUpload = Upload{
	Type:   "application/octet-stream",
	Method: "PUT",
}

// Upload PUTs the local file to the given URL
func (c *Client) Upload(f io.Reader, u *url.URL, param *Upload) error {
	var err error

	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, f, param.ContentLength)
		f = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	req, err := http.NewRequest(param.Method, u.String(), f)
	if err != nil {
		return err
	}

	req.ContentLength = param.ContentLength
	req.Header.Set("Content-Type", param.Type)

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	switch res.StatusCode {
	case http.StatusOK:
	case http.StatusCreated:
	default:
		err = errors.New(res.Status)
	}

	return err
}

// UploadFile PUTs the local file to the given URL
func (c *Client) UploadFile(file string, u *url.URL, param *Upload) error {
	if param == nil {
		p := DefaultUpload // Copy since we set ContentLength
		param = &p
	}

	s, err := os.Stat(file)
	if err != nil {
		return err
	}

	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	param.ContentLength = s.Size()

	return c.Upload(f, u, param)
}

type Download struct {
	Method   string
	Progress progress.Sinker
}

var DefaultDownload = Download{
	Method: "GET",
}

// DownloadFile GETs the given URL to a local file
func (c *Client) DownloadFile(file string, u *url.URL, param *Download) error {
	var err error

	if param == nil {
		param = &DefaultDownload
	}

	fh, err := os.Create(file)
	if err != nil {
		return err
	}
	defer fh.Close()

	req, err := http.NewRequest(param.Method, u.String(), nil)
	if err != nil {
		return err
	}

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	defer res.Body.Close()

	switch res.StatusCode {
	case http.StatusOK:
	default:
		err = errors.New(res.Status)
	}

	if err != nil {
		return err
	}

	var r io.Reader = res.Body
	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, res.Body, res.ContentLength)
		r = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	_, err = io.Copy(fh, r)
	if err != nil {
		return err
	}

	// Assign error before returning so that it gets picked up by the deferred
	// function marking the progress reader as done.
	err = fh.Close()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright © 2009 Fazlul Shahriar <fshahriar@gmail.com>.
// See LICENSE file for license details.

// This package provides the client side interface to MPD (Music Player Daemon).
// The protocol reference can be found at http://www.musicpd.org/doc/protocol/index.html
package mpd

import (
	"bufio";
	"fmt";
	"net";
	"os";
	"strconv";
	"strings";
)

var Chatty bool	// print all conversation with MPD (for debugging)

type Client struct {
	conn	net.Conn;
	rw	*bufio.ReadWriter;
}

type Attrs map[string]string

// Connect connects to MPD listening on address addr (e.g. "127.0.0.1:6600")
// on network network (e.g. "tcp").
func Connect(network, addr string) (c *Client, err os.Error) {
	conn, err := net.Dial(network, "", addr);
	if err != nil {
		return nil, err
	}
	c = new(Client);
	c.rw = bufio.NewReadWriter(bufio.NewReader(conn),
		bufio.NewWriter(conn));
	line, err := c.readLine();
	if err != nil {
		return nil, err
	}
	if line[0:6] != "OK MPD" {
		return nil, os.NewError("no greeting")
	}
	return;
}

// Close terminates the connection with MPD.
func (c *Client) Close() (err os.Error) {
	if c.conn != nil {
		c.writeLine("close");
		err = c.conn.Close();
		c.conn = nil;
	}
	return;
}

func (c *Client) readLine() (line string, err os.Error) {
	line, err = c.rw.ReadString('\n');
	if err != nil {
		return
	}
	if line[len(line)-1] == '\n' {
		line = line[0 : len(line)-1]
	}
	if Chatty {
		fmt.Println("-->", line)
	}
	return;
}

func (c *Client) writeLine(line string) (err os.Error) {
	if Chatty {
		fmt.Println("<--", line)
	}
	_, err = c.rw.Write(strings.Bytes(line + "\n"));
	// TODO: try again if # written != len(buf)
	c.rw.Flush();
	return;
}

func (c *Client) readPlaylist() (pls []Attrs, err os.Error) {
	pls = make([]Attrs, 100);

	n := 0;
	for {
		line, err := c.readLine();
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		if strings.HasPrefix(line, "file:") {	// new song entry begins
			n++;
			if n > len(pls) || n > cap(pls) {
				pls1 := make([]Attrs, 2*cap(pls));
				for k, a := range pls {
					pls1[k] = a
				}
				pls = pls1;
			}
			pls[n-1] = make(Attrs);
		}
		if n == 0 {
			return nil, os.NewError("unexpected: " + line)
		}
		z := strings.Index(line, ": ");
		if z < 0 {
			return nil, os.NewError("can't parse line: " + line)
		}
		key := line[0:z];
		pls[n-1][key] = line[z+2:];
	}
	return pls[0:n], nil;
}

func (c *Client) getAttrs() (attrs Attrs, err os.Error) {
	attrs = make(Attrs);
	for {
		line, err := c.readLine();
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		z := strings.Index(line, ": ");
		if z < 0 {
			return nil, os.NewError("can't parse line: " + line)
		}
		key := line[0:z];
		attrs[key] = line[z+2:];
	}
	return;
}

// CurrentSong returns information about the current song in the playlist.
func (c *Client) CurrentSong() (Attrs, os.Error) {
	c.writeLine("currentsong");
	return c.getAttrs();
}

// Status returns information about the current status of MPD.
func (c *Client) Status() (Attrs, os.Error) {
	c.writeLine("status");
	return c.getAttrs();
}

func (c *Client) readErr() (err os.Error) {
	line, err := c.readLine();
	switch {
	case err != nil:
		return err
	case line == "OK":
		return nil
	case strings.HasPrefix(line, "ACK "):
		return os.NewError(line[4:])
	}
	return os.NewError("unexpected response: " + line);
}

//
// Playback control
//

// Next plays next song in the playlist.
func (c *Client) Next() os.Error {
	c.writeLine("next");
	return c.readErr();
}

// Pause pauses playback if pause is true; resumes playback otherwise.
func (c *Client) Pause(pause bool) os.Error {
	if pause {
		c.writeLine("pause 1")
	} else {
		c.writeLine("pause 0")
	}
	return c.readErr();
}

// Play starts playing the song at playlist position pos. If pos is negative,
// start playing at the current position in the playlist.
func (c *Client) Play(pos int) os.Error {
	if pos < 0 {
		c.writeLine("play")
	} else {
		c.writeLine(fmt.Sprintf("play %d", pos))
	}
	return c.readErr();
}

// PlayId plays the song identified by id. If id is negative, start playing
// at the currect position in playlist.
func (c *Client) PlayId(id int) os.Error {
	if id < 0 {
		c.writeLine("playid")
	} else {
		c.writeLine(fmt.Sprintf("playid %d", id))
	}
	return c.readErr();
}

// Previous plays previous song in the playlist.
func (c *Client) Previous() os.Error {
	c.writeLine("next");
	return c.readErr();
}

// Seek seeks to the position time (in seconds) of the song at playlist position pos.
func (c *Client) Seek(pos, time int) os.Error {
	c.writeLine(fmt.Sprintf("seek %d %d", pos, time));
	return c.readErr();
}

// SeekId is identical to Seek except the song is identified by it's id
// (not position in playlist).
func (c *Client) SeekId(id, time int) os.Error {
	c.writeLine(fmt.Sprintf("seekid %d %d", id, time));
	return c.readErr();
}

// Stop stops playback.
func (c *Client) Stop() os.Error {
	c.writeLine("stop");
	return c.readErr();
}

//
// Playlist related functions
//

// PlaylistInfo returns attributes for songs in the current playlist. If
// both start and end are negative, it does this for all songs in
// playlist. If end is negative but start is positive, it does it for the
// song at position start. If both start and end are positive, it does it
// for positions in range [start, end).
func (c *Client) PlaylistInfo(start, end int) (pls []Attrs, err os.Error) {
	if start < 0 && end >= 0 {
		return nil, os.NewError("negative start index")
	}
	if start >= 0 && end < 0 {
		c.writeLine(fmt.Sprintf("playlistinfo %d", start));
		return c.readPlaylist();
	}
	c.writeLine("playlistinfo");
	pls, err = c.readPlaylist();
	if err != nil || start < 0 || end < 0 {
		return
	}
	return pls[start:end], nil;
}

// Delete deletes songs from playlist. If both start and end are positive,
// it deletes those at positions in range [start, end). If end is negative,
// it deletes the song at position start.
func (c *Client) Delete(start, end int) os.Error {
	if start < 0 {
		return os.NewError("negative start index")
	}
	if end < 0 {
		c.writeLine(fmt.Sprintf("delete %d", start))
	} else {
		c.writeLine(fmt.Sprintf("delete %d %d", start, end))
	}
	return c.readErr();
}

// DeleteId deletes the song identified by id.
func (c *Client) DeleteId(id int) os.Error {
	c.writeLine(fmt.Sprintf("deleteid %d", id));
	return c.readErr();
}

// Add adds the file/directory uri to playlist. Directories add recursively.
func (c *Client) Add(uri string) os.Error {
	c.writeLine(fmt.Sprintf("add %q", uri));
	return c.readErr();
}

// AddId adds the file/directory uri to playlist and returns the identity
// id of the song added. If pos is positive, the song is added to position
// pos.
func (c *Client) AddId(uri string, pos int) (id int, err os.Error) {
	if pos >= 0 {
		c.writeLine(fmt.Sprintf("addid %q %d", uri, pos))
	} else {
		c.writeLine(fmt.Sprintf("addid %q", uri))
	}
	attrs, err := c.getAttrs();
	if err != nil {
		return
	}
	tok, ok := attrs["Id"];
	if !ok {
		return -1, os.NewError("addid did not return Id")
	}
	return strconv.Atoi(tok);
}

// Clear clears the currect playlist.
func (c *Client) Clear() os.Error {
	c.writeLine("clear");
	return c.readErr();
}
<commit_msg>fix typo<commit_after>// Copyright © 2009 Fazlul Shahriar <fshahriar@gmail.com>.
// See LICENSE file for license details.

// This package provides the client side interface to MPD (Music Player Daemon).
// The protocol reference can be found at http://www.musicpd.org/doc/protocol/index.html
package mpd

import (
	"bufio";
	"fmt";
	"net";
	"os";
	"strconv";
	"strings";
)

var Chatty bool	// print all conversation with MPD (for debugging)

type Client struct {
	conn	net.Conn;
	rw	*bufio.ReadWriter;
}

type Attrs map[string]string

// Connect connects to MPD listening on address addr (e.g. "127.0.0.1:6600")
// on network network (e.g. "tcp").
func Connect(network, addr string) (c *Client, err os.Error) {
	conn, err := net.Dial(network, "", addr);
	if err != nil {
		return nil, err
	}
	c = new(Client);
	c.rw = bufio.NewReadWriter(bufio.NewReader(conn),
		bufio.NewWriter(conn));
	line, err := c.readLine();
	if err != nil {
		return nil, err
	}
	if line[0:6] != "OK MPD" {
		return nil, os.NewError("no greeting")
	}
	return;
}

// Close terminates the connection with MPD.
func (c *Client) Close() (err os.Error) {
	if c.conn != nil {
		c.writeLine("close");
		err = c.conn.Close();
		c.conn = nil;
	}
	return;
}

func (c *Client) readLine() (line string, err os.Error) {
	line, err = c.rw.ReadString('\n');
	if err != nil {
		return
	}
	if line[len(line)-1] == '\n' {
		line = line[0 : len(line)-1]
	}
	if Chatty {
		fmt.Println("-->", line)
	}
	return;
}

func (c *Client) writeLine(line string) (err os.Error) {
	if Chatty {
		fmt.Println("<--", line)
	}
	_, err = c.rw.Write(strings.Bytes(line + "\n"));
	// TODO: try again if # written != len(buf)
	c.rw.Flush();
	return;
}

func (c *Client) readPlaylist() (pls []Attrs, err os.Error) {
	pls = make([]Attrs, 100);

	n := 0;
	for {
		line, err := c.readLine();
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		if strings.HasPrefix(line, "file:") {	// new song entry begins
			n++;
			if n > len(pls) || n > cap(pls) {
				pls1 := make([]Attrs, 2*cap(pls));
				for k, a := range pls {
					pls1[k] = a
				}
				pls = pls1;
			}
			pls[n-1] = make(Attrs);
		}
		if n == 0 {
			return nil, os.NewError("unexpected: " + line)
		}
		z := strings.Index(line, ": ");
		if z < 0 {
			return nil, os.NewError("can't parse line: " + line)
		}
		key := line[0:z];
		pls[n-1][key] = line[z+2:];
	}
	return pls[0:n], nil;
}

func (c *Client) getAttrs() (attrs Attrs, err os.Error) {
	attrs = make(Attrs);
	for {
		line, err := c.readLine();
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		z := strings.Index(line, ": ");
		if z < 0 {
			return nil, os.NewError("can't parse line: " + line)
		}
		key := line[0:z];
		attrs[key] = line[z+2:];
	}
	return;
}

// CurrentSong returns information about the current song in the playlist.
func (c *Client) CurrentSong() (Attrs, os.Error) {
	c.writeLine("currentsong");
	return c.getAttrs();
}

// Status returns information about the current status of MPD.
func (c *Client) Status() (Attrs, os.Error) {
	c.writeLine("status");
	return c.getAttrs();
}

func (c *Client) readErr() (err os.Error) {
	line, err := c.readLine();
	switch {
	case err != nil:
		return err
	case line == "OK":
		return nil
	case strings.HasPrefix(line, "ACK "):
		return os.NewError(line[4:])
	}
	return os.NewError("unexpected response: " + line);
}

//
// Playback control
//

// Next plays next song in the playlist.
func (c *Client) Next() os.Error {
	c.writeLine("next");
	return c.readErr();
}

// Pause pauses playback if pause is true; resumes playback otherwise.
func (c *Client) Pause(pause bool) os.Error {
	if pause {
		c.writeLine("pause 1")
	} else {
		c.writeLine("pause 0")
	}
	return c.readErr();
}

// Play starts playing the song at playlist position pos. If pos is negative,
// start playing at the current position in the playlist.
func (c *Client) Play(pos int) os.Error {
	if pos < 0 {
		c.writeLine("play")
	} else {
		c.writeLine(fmt.Sprintf("play %d", pos))
	}
	return c.readErr();
}

// PlayId plays the song identified by id. If id is negative, start playing
// at the currect position in playlist.
func (c *Client) PlayId(id int) os.Error {
	if id < 0 {
		c.writeLine("playid")
	} else {
		c.writeLine(fmt.Sprintf("playid %d", id))
	}
	return c.readErr();
}

// Previous plays previous song in the playlist.
func (c *Client) Previous() os.Error {
	c.writeLine("next");
	return c.readErr();
}

// Seek seeks to the position time (in seconds) of the song at playlist position pos.
func (c *Client) Seek(pos, time int) os.Error {
	c.writeLine(fmt.Sprintf("seek %d %d", pos, time));
	return c.readErr();
}

// SeekId is identical to Seek except the song is identified by it's id
// (not position in playlist).
func (c *Client) SeekId(id, time int) os.Error {
	c.writeLine(fmt.Sprintf("seekid %d %d", id, time));
	return c.readErr();
}

// Stop stops playback.
func (c *Client) Stop() os.Error {
	c.writeLine("stop");
	return c.readErr();
}

//
// Playlist related functions
//

// PlaylistInfo returns attributes for songs in the current playlist. If
// both start and end are negative, it does this for all songs in
// playlist. If end is negative but start is positive, it does it for the
// song at position start. If both start and end are positive, it does it
// for positions in range [start, end).
func (c *Client) PlaylistInfo(start, end int) (pls []Attrs, err os.Error) {
	if start < 0 && end >= 0 {
		return nil, os.NewError("negative start index")
	}
	if start >= 0 && end < 0 {
		c.writeLine(fmt.Sprintf("playlistinfo %d", start));
		return c.readPlaylist();
	}
	c.writeLine("playlistinfo");
	pls, err = c.readPlaylist();
	if err != nil || start < 0 || end < 0 {
		return
	}
	return pls[start:end], nil;
}

// Delete deletes songs from playlist. If both start and end are positive,
// it deletes those at positions in range [start, end). If end is negative,
// it deletes the song at position start.
func (c *Client) Delete(start, end int) os.Error {
	if start < 0 {
		return os.NewError("negative start index")
	}
	if end < 0 {
		c.writeLine(fmt.Sprintf("delete %d", start))
	} else {
		c.writeLine(fmt.Sprintf("delete %d %d", start, end))
	}
	return c.readErr();
}

// DeleteId deletes the song identified by id.
func (c *Client) DeleteId(id int) os.Error {
	c.writeLine(fmt.Sprintf("deleteid %d", id));
	return c.readErr();
}

// Add adds the file/directory uri to playlist. Directories add recursively.
func (c *Client) Add(uri string) os.Error {
	c.writeLine(fmt.Sprintf("add %q", uri));
	return c.readErr();
}

// AddId adds the file/directory uri to playlist and returns the identity
// id of the song added. If pos is positive, the song is added to position
// pos.
func (c *Client) AddId(uri string, pos int) (id int, err os.Error) {
	if pos >= 0 {
		c.writeLine(fmt.Sprintf("addid %q %d", uri, pos))
	} else {
		c.writeLine(fmt.Sprintf("addid %q", uri))
	}
	attrs, err := c.getAttrs();
	if err != nil {
		return
	}
	tok, ok := attrs["Id"];
	if !ok {
		return -1, os.NewError("addid did not return Id")
	}
	return strconv.Atoi(tok);
}

// Clear clears the current playlist.
func (c *Client) Clear() os.Error {
	c.writeLine("clear");
	return c.readErr();
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/html"
)

const mysqlPerms = 0660

// downloadInfoStruct stores information necessary for the client to download and apply objects to the database
type (
	downloadInfoStruct struct {
		db         *sql.DB
		taburl     string
		backurl    string
		schema     string
		table      string
		mysqldir   string
		uid        int
		gid        int
		engine     string
		extensions []string
		version    string
	}
)

// startClient is responsible for retrieving database creation satements and binary table files from a trite server instance.
func startClient(triteURL string, tritePort string, workers uint, dbi *mysqlCredentials) {

	// Make a database connection
	db, err := dbi.connect()
	defer db.Close()

	// Problem connecting to database
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Percona import variable differs between versions
	var ignore string
	var version string
	err = db.QueryRow("show global variables like 'version'").Scan(&ignore, &version)
	checkErr(err)

	var importFlag string
	if strings.HasPrefix(version, "5.1") || strings.HasPrefix(version, "5.5") {
		err = db.QueryRow("show global variables like '%innodb%import%'").Scan(&importFlag, &ignore)
		checkErr(err)

		_, err = db.Exec("set global " + importFlag + "=1")
		checkErr(err)
	} else if strings.HasPrefix(version, "5.6") || strings.HasPrefix(version, "10") {
		// No import flag for 5.6 or MariaDB 10
	} else {
		fmt.Fprintln(os.Stderr, version, "is not supported")
		os.Exit(1)
	}

	// Get MySQL datadir
	var mysqldir string
	err = db.QueryRow("show variables like 'datadir'").Scan(&ignore, &mysqldir)
	checkErr(err)

	// Make sure mysql datadir is writable
	err = ioutil.WriteFile(mysqldir+"/trite_test", []byte("delete\n"), mysqlPerms)
	if err != nil {
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, "The MySQL data directory is not writable as this user!")
		fmt.Fprintln(os.Stderr)
		os.Exit(1)
	} else {
		os.Remove(mysqldir + "/trite_test")
	}

	// URL variables
	taburl := "http://" + triteURL + ":" + tritePort + "/tables/"
	backurl := "http://" + triteURL + ":" + tritePort + "/backups/"

	// Verify server urls are accessible
	urls := []string{taburl, backurl}
	for _, url := range urls {
		_, err = http.Head(url)
		if err != nil {
			fmt.Fprintln(os.Stderr)
			fmt.Fprintln(os.Stderr)
			fmt.Fprintln(os.Stderr, "Problem connecting to", url)
			fmt.Fprintln(os.Stderr, "Check that the server is running, port number is correct or that a firewall is not blocking access")
			os.Exit(1)
		}
	}

	// Parse html and get a list of schemas to transport
	base := getURL(taburl)
	defer base.Body.Close()
	schemas := parseAnchor(base)

	// Loop through all schemas and apply tables
	var active int32 //limit number of concurrent running applyTables()
	wg := new(sync.WaitGroup)
	for _, schema := range schemas {

		// Check if schema exists
		schemaTrimmed := strings.Trim(schema, "/")
		checkSchema(db, schemaTrimmed, taburl+schema+schemaTrimmed+".sql")

		// Parse html and get a list of tables to transport
		tablesDir := getURL(taburl + schema + "/tables")
		defer tablesDir.Body.Close()
		tables := parseAnchor(tablesDir)

		// ignore when path is empty
		if len(tables) > 0 {
			for _, table := range tables {

				downloadInfo := downloadInfoStruct{db: db, taburl: taburl, backurl: backurl, schema: schema, table: table, mysqldir: mysqldir, uid: dbi.uid, gid: dbi.gid, version: version}

				// Infinite loop to keep active go routines to 5 or less
				for {
					if active < int32(workers) {
						break
					} else {
						time.Sleep(1 * time.Second)
					}
				}

				wg.Add(1)
				atomic.AddInt32(&active, 1)
				go downloadTable(downloadInfo, &active, wg)
			}
		}
	}
	wg.Wait()

	// Loop through all schemas again and apply triggers, views, procedures & functions
	fmt.Println()
	objectTypes := []string{"trigger", "view", "procedure", "function"}
	for _, schema := range schemas {
		for _, objectType := range objectTypes {
			applyObjects(db, objectType, schema, taburl)
		}
	}

	// Reset global db variables
	if importFlag != "" {
		_, err = db.Exec("set global " + importFlag + "=0")
	}
}

// getURL is a small http.Get() wrapper
func getURL(u string) *http.Response {
	resp, err := http.Get(u)
	checkErr(err)

	return resp
}

// parseAnchor returns a slice of files and directories from a HTTP response. This function requires the google net/html sub repo.
func parseAnchor(r *http.Response) []string {
	txt := make([]string, 0)
	tok := html.NewTokenizer(r.Body)

	for {
		tt := tok.Next()
		if tt == html.ErrorToken {
			break
		}

		if tt == html.TextToken {
			a := tok.Raw()
			if a[0] != 10 {
				txt = append(txt, string(a))
			}
		}
	}
	return txt
}

// checkSchema creates a schema if it does not already exist
func checkSchema(db *sql.DB, schema string, url string) {
	var exists string
	err := db.QueryRow("show databases like '" + schema + "'").Scan(&exists)

	if err != nil {
		resp := getURL(url)
		defer resp.Body.Close()
		stmt, _ := ioutil.ReadAll(resp.Body)
		_, err = db.Exec(string(stmt))
		checkErr(err)

		fmt.Println("	Created schema", schema)
	}
}

// downloadTables retrieves files from the HTTP server. Files to download is MySQL engine specific.
func downloadTable(downloadInfo downloadInfoStruct, active *int32, wg *sync.WaitGroup) {
	filename, _ := parseFileName(downloadInfo.table)

	// Ensure backup exists and check the engine type
	// Make separate function to determine engine type
	resp, err := http.Head(downloadInfo.backurl + downloadInfo.schema + filename + ".ibd") // Assume InnoDB first
	checkErr(err)

	var engine string
	extensions := make([]string, 0)
	if resp.StatusCode == 200 {
		engine = "InnoDB"

		// 5.1 & 5.5 use .exp - 5.6 uses .cfg but it is ignored. Metadata checks appeared too brittle in testing.
		if strings.HasPrefix(downloadInfo.version, "5.1") || strings.HasPrefix(downloadInfo.version, "5.5") {
			extensions = append(extensions, ".exp")
		}

		extensions = append(extensions, ".ibd")
	} else {
		resp, err := http.Head(downloadInfo.backurl + downloadInfo.schema + filename + ".MYD") // Check for MyISAM
		checkErr(err)

		if resp.StatusCode == 200 {
			engine = "MyISAM"
			extensions = append(extensions, ".MYI")
			extensions = append(extensions, ".MYD")
			extensions = append(extensions, ".frm")
		} else {
			fmt.Println()
			fmt.Println("!!!!!!!!!!!!!!!!!!!!")
			fmt.Println("The .ibd or .MYD file is missing for table", filename)
			fmt.Println("Skipping ...")
			fmt.Println("!!!!!!!!!!!!!!!!!!!!")
			fmt.Println()

			// Need to decrement since applyTables() will never be called
			atomic.AddInt32(active, -1)
			wg.Done()

			return
		}
	}

	// Update downloadInfo struct with engine type and extensions array
	downloadInfo.engine = engine
	downloadInfo.extensions = extensions

	// Loop through and download all files from extensions array
	for _, extension := range extensions {
		tmpfile := downloadInfo.mysqldir + downloadInfo.schema + filename + extension + ".trite"
		urlfile := downloadInfo.backurl + downloadInfo.schema + filename + extension

		// Ensure the .exp exists if we expect it
		// Checking this due to a bug encountered where XtraBackup did not create a tables .exp file
		if extension == ".exp" {
			resp, err := http.Head(downloadInfo.backurl + downloadInfo.schema + filename + ".exp")
			checkErr(err)

			if resp.StatusCode != 200 {
				fmt.Println()
				fmt.Println("!!!!!!!!!!!!!!!!!!!!")
				fmt.Println("The .exp file is missing for table", filename)
				fmt.Println("Skipping ...")
				fmt.Println("!!!!!!!!!!!!!!!!!!!!")
				fmt.Println()

				// Need to decrement since applyTables() will never be called
				atomic.AddInt32(active, -1)
				wg.Done()

				return
			}
		}

		// Request and write file
		fo, err := os.Create(tmpfile)
		checkErr(err)
		defer fo.Close()

		// Chown to mysql user
		os.Chown(tmpfile, downloadInfo.uid, downloadInfo.gid)
		os.Chmod(tmpfile, mysqlPerms)

		// Download files from trite server
		w := bufio.NewWriter(fo)
		ibdresp := getURL(urlfile)
		defer ibdresp.Body.Close()

		sizeServer := ibdresp.ContentLength
		var sizeDown int64
		sizeDown, err = w.ReadFrom(ibdresp.Body) // int of file size returned here
		checkErr(err)
		w.Flush() // Just in case

		// Check if size of file downloaded matches size on server -- Add retry ability
		if sizeDown != sizeServer {
			fmt.Println("\n\nFile download size does not match size on server!")
			fmt.Println(tmpfile, "has been removed.")

			// Remove partial file download
			err = os.Remove(tmpfile)
			checkErr(err)
		}
	}

	// Call applyTables
	applyTables(downloadInfo, active, wg)
}

// applyTables performs all of the database actions required to restore a table
func applyTables(downloadInfo downloadInfoStruct, active *int32, wg *sync.WaitGroup) {
	filename, _ := parseFileName(downloadInfo.table)
	schemaTrimmed := strings.Trim(downloadInfo.schema, "/")

	// Start db transaction
	tx, err := downloadInfo.db.Begin()
	checkErr(err)

	// make the following code work for any settings -- need to preserve before changing so they can be changed back, figure out global vs session and how to handle not setting properly
	_, err = tx.Exec("set session foreign_key_checks=0")
	_, err = tx.Exec("use " + schemaTrimmed)

	switch downloadInfo.engine {
	case "InnoDB":
		// Get table create
		resp := getURL(downloadInfo.taburl + downloadInfo.schema + "tables/" + downloadInfo.table)
		defer resp.Body.Close()
		stmt, _ := ioutil.ReadAll(resp.Body)

		// Drop table if exists
		_, err := tx.Exec("drop table if exists " + addQuotes(filename))
		checkErr(err)

		// Create table
		_, err = tx.Exec(string(stmt))
		checkErr(err)

		// Discard the tablespace
		_, err = tx.Exec("alter table " + addQuotes(filename) + " discard tablespace")
		checkErr(err)

		// Lock the table just in case
		_, err = tx.Exec("lock table " + addQuotes(filename) + " write")
		checkErr(err)

		// Rename happens here
		for _, extension := range downloadInfo.extensions {
			err := os.Rename(downloadInfo.mysqldir+downloadInfo.schema+filename+extension+".trite", downloadInfo.mysqldir+downloadInfo.schema+filename+extension)
			checkErr(err)
		}

		// Import tablespace and analyze otherwise there will be no index statistics
		_, err = tx.Exec("alter table " + addQuotes(filename) + " import tablespace")
		checkErr(err)

		_, err = tx.Exec("analyze local table " + addQuotes(filename))
		checkErr(err)

		// Unlock the table
		_, err = tx.Exec("unlock tables")
		checkErr(err)

		// Commit transaction
		err = tx.Commit()
		checkErr(err)

	case "MyISAM":
		// Drop table if exists
		_, err := tx.Exec("drop table if exists " + addQuotes(filename))
		checkErr(err)

		// Rename happens here
		for _, extension := range downloadInfo.extensions {
			err = os.Rename(downloadInfo.mysqldir+downloadInfo.schema+filename+extension+".trite", downloadInfo.mysqldir+downloadInfo.schema+filename+extension)
			checkErr(err)
		}

		// Commit transaction
		err = tx.Commit()
		checkErr(err)

	default:
		fmt.Println()
		fmt.Println("!!!!!!!!!!!!!!!!!!!!")
		fmt.Println("Backup does not exist or", filename, "is not InnoDB or MyISAM")
		fmt.Println("Skipping ...")
		fmt.Println("!!!!!!!!!!!!!!!!!!!!")
		fmt.Println()
	}

	// Decrement active go routine counter
	fmt.Println(schemaTrimmed + "." + filename + " has been restored")
	atomic.AddInt32(active, -1)
	wg.Done()
}

// applyObjects is a generic function for creating procedures, functions, views and triggers.
func applyObjects(db *sql.DB, objType string, schema string, taburl string) {
	// Start transaction
	tx, err := db.Begin()
	checkErr(err)

	// Use schema
	schemaTrimmed := strings.Trim(schema, "/")
	_, err = tx.Exec("set session foreign_key_checks=0")
	_, err = tx.Exec("use " + schemaTrimmed)

	loc := getURL(taburl + schema + "/" + objType + "s")
	defer loc.Body.Close()
	objects := parseAnchor(loc)
	fmt.Println("Applying", objType+"s for", schemaTrimmed)

	// Only continue if there are objects to create
	if len(objects) > 0 { // ignore when path is empty
		for _, object := range objects {

			filename, _ := parseFileName(object)
			_, err := tx.Exec("drop " + objType + " if exists " + addQuotes(filename))
			resp := getURL(taburl + schema + objType + "s/" + object) // ssssso hacky
			defer resp.Body.Close()
			stmt, _ := ioutil.ReadAll(resp.Body)

			var objInfo createInfoStruct
			err = json.Unmarshal(stmt, &objInfo)
			checkErr(err)

			// Set session level variables to recreate stored code properly
			if objInfo.SqlMode != "" {
				_, err = tx.Exec("set session sql_mode = '" + objInfo.SqlMode + "'")
			}
			if objInfo.CharsetClient != "" {
				_, err = tx.Exec("set session character_set_client = '" + objInfo.CharsetClient + "'")
			}
			if objInfo.Collation != "" {
				_, err = tx.Exec("set session collation_connection = '" + objInfo.Collation + "'")
			}
			if objInfo.DbCollation != "" {
				_, err = tx.Exec("set session collation_database = '" + objInfo.DbCollation + "'")
			}

			// Create object
			_, err = tx.Exec(objInfo.Create)
			checkErr(err)

		}
	}

	// Commit transaction
	err = tx.Commit()
	checkErr(err)
}
<commit_msg>Proper channel concurrency!<commit_after>package main

import (
	"bufio"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"sync"

	"golang.org/x/net/html"
)

const mysqlPerms = 0660

// downloadInfoStruct stores information necessary for the client to download and apply objects to the database
type (
	downloadInfoStruct struct {
		db         *sql.DB
		taburl     string
		backurl    string
		schema     string
		table      string
		mysqldir   string
		uid        int
		gid        int
		engine     string
		extensions []string
		version    string
	}
)

// startClient is responsible for retrieving database creation satements and binary table files from a trite server instance.
func startClient(triteURL string, tritePort string, workers uint, dbi *mysqlCredentials) {

	// Make a database connection
	db, err := dbi.connect()
	defer db.Close()

	// Problem connecting to database
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Percona import variable differs between versions
	var ignore string
	var version string
	err = db.QueryRow("show global variables like 'version'").Scan(&ignore, &version)
	checkErr(err)

	var importFlag string
	if strings.HasPrefix(version, "5.1") || strings.HasPrefix(version, "5.5") {
		err = db.QueryRow("show global variables like '%innodb%import%'").Scan(&importFlag, &ignore)
		checkErr(err)

		_, err = db.Exec("set global " + importFlag + "=1")
		checkErr(err)
	} else if strings.HasPrefix(version, "5.6") || strings.HasPrefix(version, "10") {
		// No import flag for 5.6 or MariaDB 10
	} else {
		fmt.Fprintln(os.Stderr, version, "is not supported")
		os.Exit(1)
	}

	// Get MySQL datadir
	var mysqldir string
	err = db.QueryRow("show variables like 'datadir'").Scan(&ignore, &mysqldir)
	checkErr(err)

	// Make sure mysql datadir is writable
	err = ioutil.WriteFile(mysqldir+"/trite_test", []byte("delete\n"), mysqlPerms)
	if err != nil {
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, "The MySQL data directory is not writable as this user!")
		fmt.Fprintln(os.Stderr)
		os.Exit(1)
	} else {
		os.Remove(mysqldir + "/trite_test")
	}

	// URL variables
	taburl := "http://" + triteURL + ":" + tritePort + "/tables/"
	backurl := "http://" + triteURL + ":" + tritePort + "/backups/"

	// Verify server urls are accessible
	urls := []string{taburl, backurl}
	for _, url := range urls {
		_, err = http.Head(url)
		if err != nil {
			fmt.Fprintln(os.Stderr)
			fmt.Fprintln(os.Stderr)
			fmt.Fprintln(os.Stderr, "Problem connecting to", url)
			fmt.Fprintln(os.Stderr, "Check that the server is running, port number is correct or that a firewall is not blocking access")
			os.Exit(1)
		}
	}

	// Parse html and get a list of schemas to copy
	base := getURL(taburl)
	defer base.Body.Close()
	schemas := parseAnchor(base)

	// Start up download workers
	var wg sync.WaitGroup
	dl := make(chan downloadInfoStruct)
	for i := 0; i < int(workers); i++ {
		go func(i int) {
			for d := range dl {
				downloadTable(d)
				wg.Done()
			}
		}(i)
	}

	// Loop through all schemas and apply tables
	for _, schema := range schemas {

		// Check if schema exists
		schemaTrimmed := strings.Trim(schema, "/")
		checkSchema(db, schemaTrimmed, taburl+schema+schemaTrimmed+".sql")

		// Parse html and get a list of tables to transport
		tablesDir := getURL(taburl + schema + "/tables")
		defer tablesDir.Body.Close()
		tables := parseAnchor(tablesDir)

		// ignore when path is empty
		if len(tables) > 0 {
			for _, table := range tables {
				wg.Add(1)
				downloadInfo := downloadInfoStruct{db: db, taburl: taburl, backurl: backurl, schema: schema, table: table, mysqldir: mysqldir, uid: dbi.uid, gid: dbi.gid, version: version}
				dl <- downloadInfo
			}
		}
	}
	wg.Wait()

	// Loop through all schemas again and apply triggers, views, procedures & functions
	fmt.Println()
	objectTypes := []string{"trigger", "view", "procedure", "function"}
	for _, schema := range schemas {
		for _, objectType := range objectTypes {
			applyObjects(db, objectType, schema, taburl)
		}
	}

	// Reset global db variables
	if importFlag != "" {
		_, err = db.Exec("set global " + importFlag + "=0")
	}
}

// getURL is a small http.Get() wrapper
func getURL(u string) *http.Response {
	resp, err := http.Get(u)
	checkErr(err)

	return resp
}

// parseAnchor returns a slice of files and directories from a HTTP response. This function requires the google net/html sub repo.
func parseAnchor(r *http.Response) []string {
	txt := make([]string, 0)
	tok := html.NewTokenizer(r.Body)

	for {
		tt := tok.Next()
		if tt == html.ErrorToken {
			break
		}

		if tt == html.TextToken {
			a := tok.Raw()
			if a[0] != 10 {
				txt = append(txt, string(a))
			}
		}
	}
	return txt
}

// checkSchema creates a schema if it does not already exist
func checkSchema(db *sql.DB, schema string, url string) {
	var exists string
	err := db.QueryRow("show databases like '" + schema + "'").Scan(&exists)

	if err != nil {
		resp := getURL(url)
		defer resp.Body.Close()
		stmt, _ := ioutil.ReadAll(resp.Body)
		_, err = db.Exec(string(stmt))
		checkErr(err)

		fmt.Println("	Created schema", schema)
	}
}

// downloadTables retrieves files from the HTTP server. Files to download is MySQL engine specific.
func downloadTable(downloadInfo downloadInfoStruct) {
	filename, _ := parseFileName(downloadInfo.table)

	// Ensure backup exists and check the engine type
	// Make separate function to determine engine type
	resp, err := http.Head(downloadInfo.backurl + downloadInfo.schema + filename + ".ibd") // Assume InnoDB first
	checkErr(err)

	var engine string
	extensions := make([]string, 0)
	if resp.StatusCode == 200 {
		engine = "InnoDB"

		// 5.1 & 5.5 use .exp - 5.6 uses .cfg but it is ignored. Metadata checks appeared too brittle in testing.
		if strings.HasPrefix(downloadInfo.version, "5.1") || strings.HasPrefix(downloadInfo.version, "5.5") {
			extensions = append(extensions, ".exp")
		}

		extensions = append(extensions, ".ibd")
	} else {
		resp, err := http.Head(downloadInfo.backurl + downloadInfo.schema + filename + ".MYD") // Check for MyISAM
		checkErr(err)

		if resp.StatusCode == 200 {
			engine = "MyISAM"
			extensions = append(extensions, ".MYI")
			extensions = append(extensions, ".MYD")
			extensions = append(extensions, ".frm")
		} else {
			fmt.Println()
			fmt.Println("!!!!!!!!!!!!!!!!!!!!")
			fmt.Println("The .ibd or .MYD file is missing for table", filename)
			fmt.Println("Skipping ...")
			fmt.Println("!!!!!!!!!!!!!!!!!!!!")
			fmt.Println()

			return
		}
	}

	// Update downloadInfo struct with engine type and extensions array
	downloadInfo.engine = engine
	downloadInfo.extensions = extensions

	// Loop through and download all files from extensions array
	for _, extension := range extensions {
		tmpfile := downloadInfo.mysqldir + downloadInfo.schema + filename + extension + ".trite"
		urlfile := downloadInfo.backurl + downloadInfo.schema + filename + extension

		// Ensure the .exp exists if we expect it
		// Checking this due to a bug encountered where XtraBackup did not create a tables .exp file
		if extension == ".exp" {
			resp, err := http.Head(downloadInfo.backurl + downloadInfo.schema + filename + ".exp")
			checkErr(err)

			if resp.StatusCode != 200 {
				fmt.Println()
				fmt.Println("!!!!!!!!!!!!!!!!!!!!")
				fmt.Println("The .exp file is missing for table", filename)
				fmt.Println("Skipping ...")
				fmt.Println("!!!!!!!!!!!!!!!!!!!!")
				fmt.Println()

				return
			}
		}

		// Request and write file
		fo, err := os.Create(tmpfile)
		checkErr(err)
		defer fo.Close()

		// Chown to mysql user
		os.Chown(tmpfile, downloadInfo.uid, downloadInfo.gid)
		os.Chmod(tmpfile, mysqlPerms)

		// Download files from trite server
		w := bufio.NewWriter(fo)
		ibdresp := getURL(urlfile)
		defer ibdresp.Body.Close()

		sizeServer := ibdresp.ContentLength
		var sizeDown int64
		sizeDown, err = w.ReadFrom(ibdresp.Body) // int of file size returned here
		checkErr(err)
		w.Flush() // Just in case

		// Check if size of file downloaded matches size on server -- Add retry ability
		if sizeDown != sizeServer {
			fmt.Println("\n\nFile download size does not match size on server!")
			fmt.Println(tmpfile, "has been removed.")

			// Remove partial file download
			err = os.Remove(tmpfile)
			checkErr(err)
		}
	}

	// Call applyTables
	applyTables(downloadInfo)
}

// applyTables performs all of the database actions required to restore a table
func applyTables(downloadInfo downloadInfoStruct) {
	filename, _ := parseFileName(downloadInfo.table)
	schemaTrimmed := strings.Trim(downloadInfo.schema, "/")

	// Start db transaction
	tx, err := downloadInfo.db.Begin()
	checkErr(err)

	// make the following code work for any settings -- need to preserve before changing so they can be changed back, figure out global vs session and how to handle not setting properly
	_, err = tx.Exec("set session foreign_key_checks=0")
	_, err = tx.Exec("use " + schemaTrimmed)

	switch downloadInfo.engine {
	case "InnoDB":
		// Get table create
		resp := getURL(downloadInfo.taburl + downloadInfo.schema + "tables/" + downloadInfo.table)
		defer resp.Body.Close()
		stmt, _ := ioutil.ReadAll(resp.Body)

		// Drop table if exists
		_, err := tx.Exec("drop table if exists " + addQuotes(filename))
		checkErr(err)

		// Create table
		_, err = tx.Exec(string(stmt))
		checkErr(err)

		// Discard the tablespace
		_, err = tx.Exec("alter table " + addQuotes(filename) + " discard tablespace")
		checkErr(err)

		// Lock the table just in case
		_, err = tx.Exec("lock table " + addQuotes(filename) + " write")
		checkErr(err)

		// Rename happens here
		for _, extension := range downloadInfo.extensions {
			err := os.Rename(downloadInfo.mysqldir+downloadInfo.schema+filename+extension+".trite", downloadInfo.mysqldir+downloadInfo.schema+filename+extension)
			checkErr(err)
		}

		// Import tablespace and analyze otherwise there will be no index statistics
		_, err = tx.Exec("alter table " + addQuotes(filename) + " import tablespace")
		checkErr(err)

		_, err = tx.Exec("analyze local table " + addQuotes(filename))
		checkErr(err)

		// Unlock the table
		_, err = tx.Exec("unlock tables")
		checkErr(err)

		// Commit transaction
		err = tx.Commit()
		checkErr(err)

	case "MyISAM":
		// Drop table if exists
		_, err := tx.Exec("drop table if exists " + addQuotes(filename))
		checkErr(err)

		// Rename happens here
		for _, extension := range downloadInfo.extensions {
			err = os.Rename(downloadInfo.mysqldir+downloadInfo.schema+filename+extension+".trite", downloadInfo.mysqldir+downloadInfo.schema+filename+extension)
			checkErr(err)
		}

		// Commit transaction
		err = tx.Commit()
		checkErr(err)

	default:
		fmt.Println()
		fmt.Println("!!!!!!!!!!!!!!!!!!!!")
		fmt.Println("Backup does not exist or", filename, "is not InnoDB or MyISAM")
		fmt.Println("Skipping ...")
		fmt.Println("!!!!!!!!!!!!!!!!!!!!")
		fmt.Println()
	}

	// Decrement active go routine counter
	fmt.Println(schemaTrimmed + "." + filename + " has been restored")
}

// applyObjects is a generic function for creating procedures, functions, views and triggers.
func applyObjects(db *sql.DB, objType string, schema string, taburl string) {
	// Start transaction
	tx, err := db.Begin()
	checkErr(err)

	// Use schema
	schemaTrimmed := strings.Trim(schema, "/")
	_, err = tx.Exec("set session foreign_key_checks=0")
	_, err = tx.Exec("use " + schemaTrimmed)

	loc := getURL(taburl + schema + "/" + objType + "s")
	defer loc.Body.Close()
	objects := parseAnchor(loc)
	fmt.Println("Applying", objType+"s for", schemaTrimmed)

	// Only continue if there are objects to create
	if len(objects) > 0 { // ignore when path is empty
		for _, object := range objects {

			filename, _ := parseFileName(object)
			_, err := tx.Exec("drop " + objType + " if exists " + addQuotes(filename))
			resp := getURL(taburl + schema + objType + "s/" + object) // ssssso hacky
			defer resp.Body.Close()
			stmt, _ := ioutil.ReadAll(resp.Body)

			var objInfo createInfoStruct
			err = json.Unmarshal(stmt, &objInfo)
			checkErr(err)

			// Set session level variables to recreate stored code properly
			if objInfo.SqlMode != "" {
				_, err = tx.Exec("set session sql_mode = '" + objInfo.SqlMode + "'")
			}
			if objInfo.CharsetClient != "" {
				_, err = tx.Exec("set session character_set_client = '" + objInfo.CharsetClient + "'")
			}
			if objInfo.Collation != "" {
				_, err = tx.Exec("set session collation_connection = '" + objInfo.Collation + "'")
			}
			if objInfo.DbCollation != "" {
				_, err = tx.Exec("set session collation_database = '" + objInfo.DbCollation + "'")
			}

			// Create object
			_, err = tx.Exec(objInfo.Create)
			checkErr(err)

		}
	}

	// Commit transaction
	err = tx.Commit()
	checkErr(err)
}
<|endoftext|>


<commit_before>package main

import (
	"strconv"
	"strings"

	statsd "github.com/DataDog/datadog-go/statsd"
	log "github.com/Sirupsen/logrus"
)

const sampleRate = 1.0

const (
	routerMsg int = iota
	scalingMsg
	sampleMsg
	metricsMsg
	metricsTag
)

var routerMetricsKeys = []string{"dyno", "method", "status", "path", "host", "code", "desc", "at"}
var sampleMetricsKeys = []string{"source"}
var scalingMetricsKeys = []string{"mailer", "web"}
var customMetricsKeys = []string{"media_type", "output_type", "route"}

type Client struct {
	*statsd.Client
}

func statsdClient(addr string) (*Client, error) {

	c, err := statsd.New(addr)
	return &Client{c}, err
}

func (c *Client) sendToStatsd(in chan *logMetrics) {

	var data *logMetrics
	var ok bool
	for {
		data, ok = <-in

		if !ok { //Exit, channel was closed
			return
		}

		log.WithFields(log.Fields{
			"type":   data.typ,
			"app":    data.app,
			"tags":   data.tags,
			"prefix": data.prefix,
		}).Debug("logMetrics received")

		if data.typ == routerMsg {
			c.sendRouterMsg(data)
		} else if data.typ == sampleMsg {
			c.sendSampleMsg(data)
		} else if data.typ == scalingMsg {
			c.sendScalingMsg(data)
		} else if data.typ == metricsMsg {
			c.sendMetricsMsg(data)
		} else if data.typ == metricsTag {
			c.sendMetricsWithTags(data)
		} else {
			log.WithField("type", data.typ).Warn("Unknown log message")
		}
	}
}

func (c *Client) sendRouterMsg(data *logMetrics) {

	tags := *data.tags
	for _, mk := range routerMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			tags = append(tags, mk+":"+v.Val)
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   *data.tags,
		"prefix": *data.prefix,
	}).Debug("sendRouterMsg")

	conn, err := strconv.ParseFloat(data.metrics["connect"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"err":    err,
			"metric": "connect",
		}).Info("Could not parse metric value")
		return
	}
	serv, err := strconv.ParseFloat(data.metrics["service"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"metric": "service",
			"err":    err,
		}).Info("Could not parse metric value")
		return
	}

	err = c.Histogram(*data.prefix+"heroku.router.request.connect", conn, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	err = c.Histogram(*data.prefix+"heroku.router.request.service", serv, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	if data.metrics["at"].Val == "error" {
		err = c.Count(*data.prefix+"heroku.router.error", 1, tags, 0.1)
		if err != nil {
			log.WithField("error", err).Info("Failed to send Count")
		}
	}
}

func (c *Client) sendSampleMsg(data *logMetrics) {

	tags := *data.tags
	for _, mk := range sampleMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			tags = append(tags, mk+":"+v.Val)
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendSampleMsg")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "sample",
					"metric": k,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendScalingMsg(data *logMetrics) {
	tags := *data.tags

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendScalingMsg")

	for _, mk := range scalingMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+mk, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "scaling",
					"metric": mk,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendMetricsMsg(data *logMetrics) {
	tags := *data.tags

Tags:
	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			if _, err := strconv.Atoi(v.Val); err != nil {
				m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
				for _, mk := range customMetricsKeys {
					if m == mk {
						tags = append(tags, mk+":"+v.Val)
						continue Tags
					}
				}
			}
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendMetricMsg")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			if vnum, err := strconv.ParseFloat(v.Val, 10); err == nil {
				m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
				err = c.Gauge(*data.prefix+"app.metric."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Warning("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "metrics",
					"metric": k,
					"err":    err,
				}).Debug("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendMetricsWithTags(data *logMetrics) {
	tags := *data.tags

Tags:
	for k, v := range data.metrics {
		if strings.Index(k, "tag#") != -1 {
			if _, err := strconv.Atoi(v.Val); err != nil {
				m := strings.Replace(strings.Split(k, "tag#")[1], "_", ".", -1)
				for _, mk := range customMetricsKeys {
					if m == mk {
						tags = append(tags, mk+":"+v.Val)
						continue Tags
					}
				}
			}
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendMetricTag")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			if vnum, err := strconv.ParseFloat(v.Val, 10); err == nil {
				m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
				err = c.Gauge(*data.prefix+"app.metric."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Warning("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "metrics",
					"metric": k,
					"err":    err,
				}).Debug("Could not parse metric value")
			}
		}
	}
}
<commit_msg>Removes unused "metric messages"<commit_after>package main

import (
	"strconv"
	"strings"

	statsd "github.com/DataDog/datadog-go/statsd"
	log "github.com/Sirupsen/logrus"
)

const sampleRate = 1.0

const (
	routerMsg int = iota
	scalingMsg
	sampleMsg
	metricsTag
)

var routerMetricsKeys = []string{"dyno", "method", "status", "path", "host", "code", "desc", "at"}
var sampleMetricsKeys = []string{"source"}
var scalingMetricsKeys = []string{"mailer", "web"}
var customMetricsKeys = []string{"media_type", "output_type", "route"}

type Client struct {
	*statsd.Client
}

func statsdClient(addr string) (*Client, error) {

	c, err := statsd.New(addr)
	return &Client{c}, err
}

func (c *Client) sendToStatsd(in chan *logMetrics) {

	var data *logMetrics
	var ok bool
	for {
		data, ok = <-in

		if !ok { //Exit, channel was closed
			return
		}

		log.WithFields(log.Fields{
			"type":   data.typ,
			"app":    data.app,
			"tags":   data.tags,
			"prefix": data.prefix,
		}).Debug("logMetrics received")

		if data.typ == routerMsg {
			c.sendRouterMsg(data)
		} else if data.typ == sampleMsg {
			c.sendSampleMsg(data)
		} else if data.typ == scalingMsg {
			c.sendScalingMsg(data)
		} else if data.typ == metricsTag {
			c.sendMetricsWithTags(data)
		} else {
			log.WithField("type", data.typ).Warn("Unknown log message")
		}
	}
}

func (c *Client) sendRouterMsg(data *logMetrics) {

	tags := *data.tags
	for _, mk := range routerMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			tags = append(tags, mk+":"+v.Val)
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   *data.tags,
		"prefix": *data.prefix,
	}).Debug("sendRouterMsg")

	conn, err := strconv.ParseFloat(data.metrics["connect"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"err":    err,
			"metric": "connect",
		}).Info("Could not parse metric value")
		return
	}
	serv, err := strconv.ParseFloat(data.metrics["service"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"metric": "service",
			"err":    err,
		}).Info("Could not parse metric value")
		return
	}

	err = c.Histogram(*data.prefix+"heroku.router.request.connect", conn, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	err = c.Histogram(*data.prefix+"heroku.router.request.service", serv, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	if data.metrics["at"].Val == "error" {
		err = c.Count(*data.prefix+"heroku.router.error", 1, tags, 0.1)
		if err != nil {
			log.WithField("error", err).Info("Failed to send Count")
		}
	}
}

func (c *Client) sendSampleMsg(data *logMetrics) {

	tags := *data.tags
	for _, mk := range sampleMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			tags = append(tags, mk+":"+v.Val)
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendSampleMsg")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "sample",
					"metric": k,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendScalingMsg(data *logMetrics) {
	tags := *data.tags

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendScalingMsg")

	for _, mk := range scalingMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+mk, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "scaling",
					"metric": mk,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendMetricsWithTags(data *logMetrics) {
	tags := *data.tags

Tags:
	for k, v := range data.metrics {
		if strings.Index(k, "tag#") != -1 {
			if _, err := strconv.Atoi(v.Val); err != nil {
				m := strings.Replace(strings.Split(k, "tag#")[1], "_", ".", -1)
				for _, mk := range customMetricsKeys {
					if m == mk {
						tags = append(tags, mk+":"+v.Val)
						continue Tags
					}
				}
			}
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendMetricTag")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			if vnum, err := strconv.ParseFloat(v.Val, 10); err == nil {
				m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
				err = c.Gauge(*data.prefix+"app.metric."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Warning("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "metrics",
					"metric": k,
					"err":    err,
				}).Debug("Could not parse metric value")
			}
		}
	}
}
<|endoftext|>