<commit_before>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/service/athena"
	"github.com/hashicorp/terraform/helper/validation"

	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsAthenaWorkgroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsAthenaWorkgroupCreate,
		Read:   resourceAwsAthenaWorkgroupRead,
		Update: resourceAwsAthenaWorkgroupUpdate,
		Delete: resourceAwsAthenaWorkgroupDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"description": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"bytes_scanned_cutoff_per_query": {
				Type:         schema.TypeInt,
				Optional:     true,
				ValidateFunc: validation.IntAtLeast(10485760),
			},
			"enforce_workgroup_configuration": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"publish_cloudwatch_metrics_enabled": {
				Type:     schema.TypeBool,
				Optional: true,
			},
			"output_location": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"encryption_option": {
				Type:     schema.TypeString,
				Optional: true,
				ValidateFunc: validation.StringInSlice([]string{
					athena.EncryptionOptionCseKms,
					athena.EncryptionOptionSseKms,
					athena.EncryptionOptionSseS3,
				}, false),
			},
			"kms_key": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"tags": tagsSchema(),
		},
	}
}

func resourceAwsAthenaWorkgroupCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	name := d.Get("name").(string)

	input := &athena.CreateWorkGroupInput{
		Name: aws.String(name),
	}

	basicConfig := false
	resultConfig := false
	encryptConfig := false

	if v, ok := d.GetOk("description"); ok {
		input.Description = aws.String(v.(string))
	}

	inputConfiguration := &athena.WorkGroupConfiguration{}

	if v, ok := d.GetOk("bytes_scanned_cutoff_per_query"); ok {
		basicConfig = true
		inputConfiguration.BytesScannedCutoffPerQuery = aws.Int64(int64(v.(int)))
	}

	if v, ok := d.GetOk("enforce_workgroup_configuration"); ok {
		basicConfig = true
		inputConfiguration.EnforceWorkGroupConfiguration = aws.Bool(v.(bool))
	}

	if v, ok := d.GetOk("publish_cloudwatch_metrics_enabled"); ok {
		basicConfig = true
		inputConfiguration.PublishCloudWatchMetricsEnabled = aws.Bool(v.(bool))
	}

	resultConfiguration := &athena.ResultConfiguration{}

	if v, ok := d.GetOk("output_location"); ok {
		resultConfig = true
		resultConfiguration.OutputLocation = aws.String(v.(string))
	}

	encryptionConfiguration := &athena.EncryptionConfiguration{}

	if v, ok := d.GetOk("encryption_option"); ok {
		resultConfig = true
		encryptConfig = true
		encryptionConfiguration.EncryptionOption = aws.String(v.(string))

		if v.(string) == athena.EncryptionOptionCseKms || v.(string) == athena.EncryptionOptionSseKms {
			if vkms, ok := d.GetOk("kms_key"); ok {
				encryptionConfiguration.KmsKey = aws.String(vkms.(string))
			} else {
				return fmt.Errorf("KMS Key required but not provided for encryption_option: %s", v.(string))
			}
		}
	}

	if basicConfig {
		input.Configuration = inputConfiguration
	}

	if resultConfig {
		input.Configuration.ResultConfiguration = resultConfiguration

		if encryptConfig {
			input.Configuration.ResultConfiguration.EncryptionConfiguration = encryptionConfiguration
		}
	}

	// Prevent the below error:
	// InvalidRequestException: Tags provided upon WorkGroup creation must not be empty
	if v := d.Get("tags").(map[string]interface{}); len(v) > 0 {
		input.Tags = tagsFromMapAthena(v)
	}

	_, err := conn.CreateWorkGroup(input)

	if err != nil {
		return err
	}

	d.SetId(name)

	return resourceAwsAthenaWorkgroupRead(d, meta)
}

func resourceAwsAthenaWorkgroupRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	input := &athena.GetWorkGroupInput{
		WorkGroup: aws.String(d.Id()),
	}

	resp, err := conn.GetWorkGroup(input)

	if isAWSErr(err, athena.ErrCodeInvalidRequestException, "is not found") {
		log.Printf("[WARN] Athena WorkGroup (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return err
	}

	d.Set("name", *resp.WorkGroup.Name)

	if resp.WorkGroup.Description != nil {
		d.Set("description", *resp.WorkGroup.Description)
	}

	client := meta.(*AWSClient)

	arn := arn.ARN{
		Partition: client.partition,
		Region:    client.region,
		Service:   "athena",
		AccountID: client.accountid,
		Resource:  fmt.Sprintf("workgroup/%s", d.Id()),
	}

	d.Set("arn", arn.String())

	if resp.WorkGroup.Configuration != nil {
		if resp.WorkGroup.Configuration.BytesScannedCutoffPerQuery != nil {
			d.Set("bytes_scanned_cutoff_per_query", *resp.WorkGroup.Configuration.BytesScannedCutoffPerQuery)
		}

		if resp.WorkGroup.Configuration.EnforceWorkGroupConfiguration != nil {
			d.Set("enforce_workgroup_configuration", *resp.WorkGroup.Configuration.EnforceWorkGroupConfiguration)
		}

		if resp.WorkGroup.Configuration.PublishCloudWatchMetricsEnabled != nil {
			d.Set("publish_cloudwatch_metrics_enabled", *resp.WorkGroup.Configuration.PublishCloudWatchMetricsEnabled)
		}

		if resp.WorkGroup.Configuration.ResultConfiguration != nil {
			if resp.WorkGroup.Configuration.ResultConfiguration.OutputLocation != nil {
				d.Set("output_location", *resp.WorkGroup.Configuration.ResultConfiguration.OutputLocation)
			}

			if resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration != nil {
				if resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.EncryptionOption != nil {
					d.Set("encryption_option", *resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.EncryptionOption)
				}

				if resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.KmsKey != nil {
					d.Set("kms_key", *resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.KmsKey)
				}
			}
		}
	}

	err = saveTagsAthena(conn, d, d.Get("arn").(string))

	if isAWSErr(err, athena.ErrCodeInvalidRequestException, "is not found") {
		log.Printf("[WARN] Athena WorkGroup (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}

func resourceAwsAthenaWorkgroupDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	input := &athena.DeleteWorkGroupInput{
		WorkGroup: aws.String(d.Id()),
	}

	_, err := conn.DeleteWorkGroup(input)

	return err
}

func resourceAwsAthenaWorkgroupUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	workGroupUpdate := false
	resultConfigUpdate := false
	configUpdate := false
	encryptionUpdate := false
	removeEncryption := false

	input := &athena.UpdateWorkGroupInput{
		WorkGroup: aws.String(d.Get("name").(string)),
	}

	if d.HasChange("description") {
		workGroupUpdate = true
		input.Description = aws.String(d.Get("description").(string))
	}

	inputConfigurationUpdates := &athena.WorkGroupConfigurationUpdates{}

	if d.HasChange("bytes_scanned_cutoff_per_query") {
		workGroupUpdate = true
		configUpdate = true

		if v, ok := d.GetOk("bytes_scanned_cutoff_per_query"); ok {
			inputConfigurationUpdates.BytesScannedCutoffPerQuery = aws.Int64(int64(v.(int)))
		} else {
			inputConfigurationUpdates.RemoveBytesScannedCutoffPerQuery = aws.Bool(true)
		}
	}

	if d.HasChange("enforce_workgroup_configuration") {
		workGroupUpdate = true
		configUpdate = true

		v := d.Get("enforce_workgroup_configuration")
		inputConfigurationUpdates.EnforceWorkGroupConfiguration = aws.Bool(v.(bool))
	}

	if d.HasChange("publish_cloudwatch_metrics_enabled") {
		workGroupUpdate = true
		configUpdate = true

		v := d.Get("publish_cloudwatch_metrics_enabled")
		inputConfigurationUpdates.PublishCloudWatchMetricsEnabled = aws.Bool(v.(bool))
	}

	resultConfigurationUpdates := &athena.ResultConfigurationUpdates{}

	if d.HasChange("output_location") {
		workGroupUpdate = true
		configUpdate = true
		resultConfigUpdate = true

		if v, ok := d.GetOk("output_location"); ok {
			resultConfigurationUpdates.OutputLocation = aws.String(v.(string))
		} else {
			resultConfigurationUpdates.RemoveOutputLocation = aws.Bool(true)
		}
	}

	encryptionConfiguration := &athena.EncryptionConfiguration{}

	if d.HasChange("encryption_option") {
		workGroupUpdate = true
		configUpdate = true
		resultConfigUpdate = true
		encryptionUpdate = true

		if v, ok := d.GetOk("encryption_option"); ok {
			encryptionConfiguration.EncryptionOption = aws.String(v.(string))

			if v.(string) == athena.EncryptionOptionCseKms || v.(string) == athena.EncryptionOptionSseKms {
				if vkms, ok := d.GetOk("kms_key"); ok {
					encryptionConfiguration.KmsKey = aws.String(vkms.(string))
				} else {
					return fmt.Errorf("KMS Key required but not provided for encryption_option: %s", v.(string))
				}
			}
		} else {
			removeEncryption = true
			resultConfigurationUpdates.RemoveEncryptionConfiguration = aws.Bool(true)
		}
	}

	if workGroupUpdate {
		if configUpdate {
			input.ConfigurationUpdates = inputConfigurationUpdates
		}

		if resultConfigUpdate {
			input.ConfigurationUpdates.ResultConfigurationUpdates = resultConfigurationUpdates

			if encryptionUpdate && !removeEncryption {
				input.ConfigurationUpdates.ResultConfigurationUpdates.EncryptionConfiguration = encryptionConfiguration
			}
		}

		_, err := conn.UpdateWorkGroup(input)

		if err != nil {
			return fmt.Errorf("Error updating Athena WorkGroup (%s): %s", d.Id(), err)
		}
	}

	if d.HasChange("tags") {
		err := setTagsAthena(conn, d, d.Get("arn").(string))

		if err != nil {
			return fmt.Errorf("error updating tags: %s", err)
		}
	}

	return resourceAwsAthenaWorkgroupRead(d, meta)
}
<commit_msg>resource/aws_athena_workgroup: Return error message context<commit_after>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/service/athena"
	"github.com/hashicorp/terraform/helper/validation"

	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsAthenaWorkgroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsAthenaWorkgroupCreate,
		Read:   resourceAwsAthenaWorkgroupRead,
		Update: resourceAwsAthenaWorkgroupUpdate,
		Delete: resourceAwsAthenaWorkgroupDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"description": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"bytes_scanned_cutoff_per_query": {
				Type:         schema.TypeInt,
				Optional:     true,
				ValidateFunc: validation.IntAtLeast(10485760),
			},
			"enforce_workgroup_configuration": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"publish_cloudwatch_metrics_enabled": {
				Type:     schema.TypeBool,
				Optional: true,
			},
			"output_location": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"encryption_option": {
				Type:     schema.TypeString,
				Optional: true,
				ValidateFunc: validation.StringInSlice([]string{
					athena.EncryptionOptionCseKms,
					athena.EncryptionOptionSseKms,
					athena.EncryptionOptionSseS3,
				}, false),
			},
			"kms_key": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"tags": tagsSchema(),
		},
	}
}

func resourceAwsAthenaWorkgroupCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	name := d.Get("name").(string)

	input := &athena.CreateWorkGroupInput{
		Name: aws.String(name),
	}

	basicConfig := false
	resultConfig := false
	encryptConfig := false

	if v, ok := d.GetOk("description"); ok {
		input.Description = aws.String(v.(string))
	}

	inputConfiguration := &athena.WorkGroupConfiguration{}

	if v, ok := d.GetOk("bytes_scanned_cutoff_per_query"); ok {
		basicConfig = true
		inputConfiguration.BytesScannedCutoffPerQuery = aws.Int64(int64(v.(int)))
	}

	if v, ok := d.GetOk("enforce_workgroup_configuration"); ok {
		basicConfig = true
		inputConfiguration.EnforceWorkGroupConfiguration = aws.Bool(v.(bool))
	}

	if v, ok := d.GetOk("publish_cloudwatch_metrics_enabled"); ok {
		basicConfig = true
		inputConfiguration.PublishCloudWatchMetricsEnabled = aws.Bool(v.(bool))
	}

	resultConfiguration := &athena.ResultConfiguration{}

	if v, ok := d.GetOk("output_location"); ok {
		resultConfig = true
		resultConfiguration.OutputLocation = aws.String(v.(string))
	}

	encryptionConfiguration := &athena.EncryptionConfiguration{}

	if v, ok := d.GetOk("encryption_option"); ok {
		resultConfig = true
		encryptConfig = true
		encryptionConfiguration.EncryptionOption = aws.String(v.(string))

		if v.(string) == athena.EncryptionOptionCseKms || v.(string) == athena.EncryptionOptionSseKms {
			if vkms, ok := d.GetOk("kms_key"); ok {
				encryptionConfiguration.KmsKey = aws.String(vkms.(string))
			} else {
				return fmt.Errorf("KMS Key required but not provided for encryption_option: %s", v.(string))
			}
		}
	}

	if basicConfig {
		input.Configuration = inputConfiguration
	}

	if resultConfig {
		input.Configuration.ResultConfiguration = resultConfiguration

		if encryptConfig {
			input.Configuration.ResultConfiguration.EncryptionConfiguration = encryptionConfiguration
		}
	}

	// Prevent the below error:
	// InvalidRequestException: Tags provided upon WorkGroup creation must not be empty
	if v := d.Get("tags").(map[string]interface{}); len(v) > 0 {
		input.Tags = tagsFromMapAthena(v)
	}

	_, err := conn.CreateWorkGroup(input)

	if err != nil {
		return fmt.Errorf("error creating Athena WorkGroup: %s", err)
	}

	d.SetId(name)

	return resourceAwsAthenaWorkgroupRead(d, meta)
}

func resourceAwsAthenaWorkgroupRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	input := &athena.GetWorkGroupInput{
		WorkGroup: aws.String(d.Id()),
	}

	resp, err := conn.GetWorkGroup(input)

	if isAWSErr(err, athena.ErrCodeInvalidRequestException, "is not found") {
		log.Printf("[WARN] Athena WorkGroup (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("error reading Athena WorkGroup (%s): %s", d.Id(), err)
	}

	d.Set("name", *resp.WorkGroup.Name)

	if resp.WorkGroup.Description != nil {
		d.Set("description", *resp.WorkGroup.Description)
	}

	client := meta.(*AWSClient)

	arn := arn.ARN{
		Partition: client.partition,
		Region:    client.region,
		Service:   "athena",
		AccountID: client.accountid,
		Resource:  fmt.Sprintf("workgroup/%s", d.Id()),
	}

	d.Set("arn", arn.String())

	if resp.WorkGroup.Configuration != nil {
		if resp.WorkGroup.Configuration.BytesScannedCutoffPerQuery != nil {
			d.Set("bytes_scanned_cutoff_per_query", *resp.WorkGroup.Configuration.BytesScannedCutoffPerQuery)
		}

		if resp.WorkGroup.Configuration.EnforceWorkGroupConfiguration != nil {
			d.Set("enforce_workgroup_configuration", *resp.WorkGroup.Configuration.EnforceWorkGroupConfiguration)
		}

		if resp.WorkGroup.Configuration.PublishCloudWatchMetricsEnabled != nil {
			d.Set("publish_cloudwatch_metrics_enabled", *resp.WorkGroup.Configuration.PublishCloudWatchMetricsEnabled)
		}

		if resp.WorkGroup.Configuration.ResultConfiguration != nil {
			if resp.WorkGroup.Configuration.ResultConfiguration.OutputLocation != nil {
				d.Set("output_location", *resp.WorkGroup.Configuration.ResultConfiguration.OutputLocation)
			}

			if resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration != nil {
				if resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.EncryptionOption != nil {
					d.Set("encryption_option", *resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.EncryptionOption)
				}

				if resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.KmsKey != nil {
					d.Set("kms_key", *resp.WorkGroup.Configuration.ResultConfiguration.EncryptionConfiguration.KmsKey)
				}
			}
		}
	}

	err = saveTagsAthena(conn, d, d.Get("arn").(string))

	if isAWSErr(err, athena.ErrCodeInvalidRequestException, "is not found") {
		log.Printf("[WARN] Athena WorkGroup (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}

func resourceAwsAthenaWorkgroupDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	input := &athena.DeleteWorkGroupInput{
		WorkGroup: aws.String(d.Id()),
	}

	_, err := conn.DeleteWorkGroup(input)

	if err != nil {
		return fmt.Errorf("error deleting Athena WorkGroup (%s): %s", d.Id(), err)
	}

	return nil
}

func resourceAwsAthenaWorkgroupUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).athenaconn

	workGroupUpdate := false
	resultConfigUpdate := false
	configUpdate := false
	encryptionUpdate := false
	removeEncryption := false

	input := &athena.UpdateWorkGroupInput{
		WorkGroup: aws.String(d.Get("name").(string)),
	}

	if d.HasChange("description") {
		workGroupUpdate = true
		input.Description = aws.String(d.Get("description").(string))
	}

	inputConfigurationUpdates := &athena.WorkGroupConfigurationUpdates{}

	if d.HasChange("bytes_scanned_cutoff_per_query") {
		workGroupUpdate = true
		configUpdate = true

		if v, ok := d.GetOk("bytes_scanned_cutoff_per_query"); ok {
			inputConfigurationUpdates.BytesScannedCutoffPerQuery = aws.Int64(int64(v.(int)))
		} else {
			inputConfigurationUpdates.RemoveBytesScannedCutoffPerQuery = aws.Bool(true)
		}
	}

	if d.HasChange("enforce_workgroup_configuration") {
		workGroupUpdate = true
		configUpdate = true

		v := d.Get("enforce_workgroup_configuration")
		inputConfigurationUpdates.EnforceWorkGroupConfiguration = aws.Bool(v.(bool))
	}

	if d.HasChange("publish_cloudwatch_metrics_enabled") {
		workGroupUpdate = true
		configUpdate = true

		v := d.Get("publish_cloudwatch_metrics_enabled")
		inputConfigurationUpdates.PublishCloudWatchMetricsEnabled = aws.Bool(v.(bool))
	}

	resultConfigurationUpdates := &athena.ResultConfigurationUpdates{}

	if d.HasChange("output_location") {
		workGroupUpdate = true
		configUpdate = true
		resultConfigUpdate = true

		if v, ok := d.GetOk("output_location"); ok {
			resultConfigurationUpdates.OutputLocation = aws.String(v.(string))
		} else {
			resultConfigurationUpdates.RemoveOutputLocation = aws.Bool(true)
		}
	}

	encryptionConfiguration := &athena.EncryptionConfiguration{}

	if d.HasChange("encryption_option") {
		workGroupUpdate = true
		configUpdate = true
		resultConfigUpdate = true
		encryptionUpdate = true

		if v, ok := d.GetOk("encryption_option"); ok {
			encryptionConfiguration.EncryptionOption = aws.String(v.(string))

			if v.(string) == athena.EncryptionOptionCseKms || v.(string) == athena.EncryptionOptionSseKms {
				if vkms, ok := d.GetOk("kms_key"); ok {
					encryptionConfiguration.KmsKey = aws.String(vkms.(string))
				} else {
					return fmt.Errorf("KMS Key required but not provided for encryption_option: %s", v.(string))
				}
			}
		} else {
			removeEncryption = true
			resultConfigurationUpdates.RemoveEncryptionConfiguration = aws.Bool(true)
		}
	}

	if workGroupUpdate {
		if configUpdate {
			input.ConfigurationUpdates = inputConfigurationUpdates
		}

		if resultConfigUpdate {
			input.ConfigurationUpdates.ResultConfigurationUpdates = resultConfigurationUpdates

			if encryptionUpdate && !removeEncryption {
				input.ConfigurationUpdates.ResultConfigurationUpdates.EncryptionConfiguration = encryptionConfiguration
			}
		}

		_, err := conn.UpdateWorkGroup(input)

		if err != nil {
			return fmt.Errorf("error updating Athena WorkGroup (%s): %s", d.Id(), err)
		}
	}

	if d.HasChange("tags") {
		err := setTagsAthena(conn, d, d.Get("arn").(string))

		if err != nil {
			return fmt.Errorf("error updating tags: %s", err)
		}
	}

	return resourceAwsAthenaWorkgroupRead(d, meta)
}
<|endoftext|>


<commit_before>package backend

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/containerops/dockyard/backend/drivers"
   	"github.com/containerops/wrench/setting"
)

func init() {
	drivers.Register("amazons3", InitFunc)
}

func InitFunc() {
	drivers.InjectReflect.Bind("amazons3save", amazons3save)
}

func amazons3save(file string) (url string, err error) {

	var key string

	for _, key = range strings.Split(file, "/") {

	}

	fin, err := os.Open(file)
	if err != nil {
		return "", err
	}
	defer fin.Close()
	var fi os.FileInfo
	fi, err = fin.Stat()
	if err != nil {
		return "", err
	}
	filesize := fi.Size()

	requstUrl := "http://" + setting.Bucket + "." + setting.Endpoint + "/" + key
	r, _ := http.NewRequest("PUT", requstUrl, fin)
	r.ContentLength = int64(filesize)
	r.Header.Set("Date", time.Now().UTC().Format(http.TimeFormat))
	r.Header.Set("X-Amz-Acl", "public-read")

	fmt.Println(requstUrl)

	amazons3Sign(r, key, setting.AccessKeyID, setting.AccessKeysecret)
	_, err = http.DefaultClient.Do(r)
	if err != nil {
		return "", err
	}

	url = "http://" + setting.Endpoint + "/" + setting.Bucket + "/" + key
	return url, nil

}

func amazons3Sign(r *http.Request, key string, accessKeyId string, accessKeySecret string) {

	plainText := amazons3cloudMakePlainText(r, key)
	h := hmac.New(sha1.New, []byte(accessKeySecret))
	h.Write([]byte(plainText))
	sign := base64.StdEncoding.EncodeToString(h.Sum(nil))
	r.Header.Set("Authorization", "AWS "+accessKeyId+":"+sign)
}

func amazons3cloudMakePlainText(r *http.Request, key string) (plainText string) {

	plainText = r.Method + "\n"
	plainText += r.Header.Get("content-md5") + "\n"
	plainText += r.Header.Get("content-type") + "\n"
	if _, ok := r.Header["X-Amz-Date"]; !ok {
		plainText += r.Header.Get("date") + "\n"
	}

	amzHeader := getAmzHeaders(r)
	if amzHeader != "" {
		plainText += amzHeader + "\n"
	}

	plainText += "/" + setting.Bucket + "/" + key
	return
}

func getAmzHeaders(r *http.Request) (amzHeader string) {
	var keys []string
	for k, _ := range r.Header {
		if strings.HasPrefix(strings.ToLower(k), "x-amz-") {
			keys = append(keys, k)
		}
	}

	sort.Strings(keys)
	var a []string
	for _, k := range keys {
		v := r.Header[k]
		a = append(a, strings.ToLower(k)+":"+strings.Join(v, ","))
	}
	for _, h := range a {

		return h
	}
	return ""
}
<commit_msg>fix amazone s3 put object acl option bug<commit_after>package backend

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/containerops/dockyard/backend/drivers"
   	"github.com/containerops/wrench/setting"
)

func init() {
	drivers.Register("amazons3", InitFunc)
}

func InitFunc() {
	drivers.InjectReflect.Bind("amazons3save", amazons3save)
}

func amazons3save(file string) (url string, err error) {

	var key string

	for _, key = range strings.Split(file, "/") {

	}

	fin, err := os.Open(file)
	if err != nil {
		return "", err
	}
	defer fin.Close()
	var fi os.FileInfo
	fi, err = fin.Stat()
	if err != nil {
		return "", err
	}
	filesize := fi.Size()

	requstUrl := "http://" + setting.Bucket + "." + setting.Endpoint + "/" + key
	r, _ := http.NewRequest("PUT", requstUrl, fin)
	r.ContentLength = int64(filesize)
	r.Header.Set("Date", time.Now().UTC().Format(http.TimeFormat))
	r.Header.Set("X-Amz-Acl", "public-write")

	fmt.Println(requstUrl)

	amazons3Sign(r, key, setting.AccessKeyID, setting.AccessKeysecret)
	_, err = http.DefaultClient.Do(r)
	if err != nil {
		return "", err
	}

	url = "http://" + setting.Endpoint + "/" + setting.Bucket + "/" + key
	return url, nil

}

func amazons3Sign(r *http.Request, key string, accessKeyId string, accessKeySecret string) {

	plainText := amazons3cloudMakePlainText(r, key)
	h := hmac.New(sha1.New, []byte(accessKeySecret))
	h.Write([]byte(plainText))
	sign := base64.StdEncoding.EncodeToString(h.Sum(nil))
	r.Header.Set("Authorization", "AWS "+accessKeyId+":"+sign)
}

func amazons3cloudMakePlainText(r *http.Request, key string) (plainText string) {

	plainText = r.Method + "\n"
	plainText += r.Header.Get("content-md5") + "\n"
	plainText += r.Header.Get("content-type") + "\n"
	if _, ok := r.Header["X-Amz-Date"]; !ok {
		plainText += r.Header.Get("date") + "\n"
	}

	amzHeader := getAmzHeaders(r)
	if amzHeader != "" {
		plainText += amzHeader + "\n"
	}

	plainText += "/" + setting.Bucket + "/" + key
	return
}

func getAmzHeaders(r *http.Request) (amzHeader string) {
	var keys []string
	for k, _ := range r.Header {
		if strings.HasPrefix(strings.ToLower(k), "x-amz-") {
			keys = append(keys, k)
		}
	}

	sort.Strings(keys)
	var a []string
	for _, k := range keys {
		v := r.Header[k]
		a = append(a, strings.ToLower(k)+":"+strings.Join(v, ","))
	}
	for _, h := range a {

		return h
	}
	return ""
}
<|endoftext|>


<commit_before>package profile

import (
	"errors"

	"github.com/fatih/structs"
	"github.com/godbus/dbus"
	"github.com/muka/go-bluetooth/bluez"
)

// NewGattCharacteristic1 create a new GattCharacteristic1 client
func NewGattCharacteristic1(path string) (*GattCharacteristic1, error) {
	g := new(GattCharacteristic1)
	g.client = bluez.NewClient(
		&bluez.Config{
			Name:  "org.bluez",
			Iface: bluez.GattCharacteristic1Interface,
			Path:  path,
			Bus:   bluez.SystemBus,
		},
	)

	g.Path = path
	g.Properties = new(GattCharacteristic1Properties)

	_, err := g.GetProperties()
	if err != nil {
		return nil, err
	}

	return g, nil
}

// GattCharacteristic1 client
type GattCharacteristic1 struct {
	client     *bluez.Client
	Properties *GattCharacteristic1Properties
	channel    chan *dbus.Signal
	Path       string
}

// GattCharacteristic1Properties exposed properties for GattCharacteristic1
type GattCharacteristic1Properties struct {
	Value          []byte `dbus:"emit"`
	Notifying      bool
	NotifyAcquired bool
	WriteAcquired  bool
	Service        dbus.ObjectPath
	UUID           string
	Flags          []string
	Descriptors    []dbus.ObjectPath
}

//ToMap serialize properties
func (d *GattCharacteristic1Properties) ToMap() (map[string]interface{}, error) {
	if !d.Service.IsValid() {
		return nil, errors.New("GattCharacteristic1Properties: Service ObjectPath is not valid")
	}
	for i := 0; i < len(d.Descriptors); i++ {
		if d.Descriptors[i].IsValid() {
			return nil, errors.New("GattCharacteristic1Properties: Descriptors contains an ObjectPath that is not valid")
		}
	}
	return structs.Map(d), nil
}

// Close the connection
func (d *GattCharacteristic1) Close() {
	d.client.Disconnect()
}

//Register for changes signalling
func (d *GattCharacteristic1) Register() (chan *dbus.Signal, error) {
	if d.channel == nil {
		channel, err := d.client.Register(d.client.Config.Path, bluez.PropertiesInterface)
		if err != nil {
			return nil, err
		}
		d.channel = channel
	}
	return d.channel, nil
}

//Unregister for changes signalling
func (d *GattCharacteristic1) Unregister(signal chan *dbus.Signal) error {
	if d.channel != nil {
		close(d.channel)
	}
	return d.client.Unregister(d.client.Config.Path, bluez.PropertiesInterface, signal)
}

//GetProperties load all available properties
func (d *GattCharacteristic1) GetProperties() (*GattCharacteristic1Properties, error) {
	err := d.client.GetProperties(d.Properties)
	return d.Properties, err
}

//GetProperty load a single property
func (d *GattCharacteristic1) GetProperty(name string) (interface{}, error) {
	val, err := d.client.GetProperty(name)
	if err != nil {
		return nil, err
	}
	return val.Value(), nil
}

//ReadValue read a value from a characteristic
func (d *GattCharacteristic1) ReadValue(options map[string]dbus.Variant) ([]byte, error) {
	var b []byte
	err := d.client.Call("ReadValue", 0, options).Store(&b)
	return b, err
}

//WriteValue write a value to a characteristic
func (d *GattCharacteristic1) WriteValue(b []byte, options map[string]dbus.Variant) error {
	err := d.client.Call("WriteValue", 0, b, options).Store()
	return err
}

//StartNotify start notifications
func (d *GattCharacteristic1) StartNotify() error {
	return d.client.Call("StartNotify", 0).Store()
}

//StopNotify stop notifications
func (d *GattCharacteristic1) StopNotify() error {
	return d.client.Call("StopNotify", 0).Store()
}

//AcquireWrite acquire file descriptor and MTU for writing [experimental]
func (d *GattCharacteristic1) AcquireWrite() (dbus.UnixFD, uint16, error) {
	var fd dbus.UnixFD
	var mtu uint16
	err := d.client.Call("AcquireWrite", 0).Store(&fd, &mtu)
	return fd, mtu, err
}

//AcquireNotify acquire file descriptor and MTU for notify [experimental]
func (d *GattCharacteristic1) AcquireNotify() (dbus.UnixFD, uint16, error) {
	var fd dbus.UnixFD
	var mtu uint16
	err := d.client.Call("AcquireNotify", 0).Store(&fd, &mtu)
	return fd, mtu, err
}
<commit_msg>fix check on valid descriptor, @KSBhate fixes #56<commit_after>package profile

import (
	"errors"

	"github.com/fatih/structs"
	"github.com/godbus/dbus"
	"github.com/muka/go-bluetooth/bluez"
)

// NewGattCharacteristic1 create a new GattCharacteristic1 client
func NewGattCharacteristic1(path string) (*GattCharacteristic1, error) {
	g := new(GattCharacteristic1)
	g.client = bluez.NewClient(
		&bluez.Config{
			Name:  "org.bluez",
			Iface: bluez.GattCharacteristic1Interface,
			Path:  path,
			Bus:   bluez.SystemBus,
		},
	)

	g.Path = path
	g.Properties = new(GattCharacteristic1Properties)

	_, err := g.GetProperties()
	if err != nil {
		return nil, err
	}

	return g, nil
}

// GattCharacteristic1 client
type GattCharacteristic1 struct {
	client     *bluez.Client
	Properties *GattCharacteristic1Properties
	channel    chan *dbus.Signal
	Path       string
}

// GattCharacteristic1Properties exposed properties for GattCharacteristic1
type GattCharacteristic1Properties struct {
	Value          []byte `dbus:"emit"`
	Notifying      bool
	NotifyAcquired bool
	WriteAcquired  bool
	Service        dbus.ObjectPath
	UUID           string
	Flags          []string
	Descriptors    []dbus.ObjectPath
}

//ToMap serialize properties
func (d *GattCharacteristic1Properties) ToMap() (map[string]interface{}, error) {
	if !d.Service.IsValid() {
		return nil, errors.New("GattCharacteristic1Properties: Service ObjectPath is not valid")
	}
	for i := 0; i < len(d.Descriptors); i++ {
		if !d.Descriptors[i].IsValid() {
			return nil, errors.New("GattCharacteristic1Properties: Descriptors contains an ObjectPath that is not valid")
		}
	}
	return structs.Map(d), nil
}

// Close the connection
func (d *GattCharacteristic1) Close() {
	d.client.Disconnect()
}

//Register for changes signalling
func (d *GattCharacteristic1) Register() (chan *dbus.Signal, error) {
	if d.channel == nil {
		channel, err := d.client.Register(d.client.Config.Path, bluez.PropertiesInterface)
		if err != nil {
			return nil, err
		}
		d.channel = channel
	}
	return d.channel, nil
}

//Unregister for changes signalling
func (d *GattCharacteristic1) Unregister(signal chan *dbus.Signal) error {
	if d.channel != nil {
		close(d.channel)
	}
	return d.client.Unregister(d.client.Config.Path, bluez.PropertiesInterface, signal)
}

//GetProperties load all available properties
func (d *GattCharacteristic1) GetProperties() (*GattCharacteristic1Properties, error) {
	err := d.client.GetProperties(d.Properties)
	return d.Properties, err
}

//GetProperty load a single property
func (d *GattCharacteristic1) GetProperty(name string) (interface{}, error) {
	val, err := d.client.GetProperty(name)
	if err != nil {
		return nil, err
	}
	return val.Value(), nil
}

//ReadValue read a value from a characteristic
func (d *GattCharacteristic1) ReadValue(options map[string]dbus.Variant) ([]byte, error) {
	var b []byte
	err := d.client.Call("ReadValue", 0, options).Store(&b)
	return b, err
}

//WriteValue write a value to a characteristic
func (d *GattCharacteristic1) WriteValue(b []byte, options map[string]dbus.Variant) error {
	err := d.client.Call("WriteValue", 0, b, options).Store()
	return err
}

//StartNotify start notifications
func (d *GattCharacteristic1) StartNotify() error {
	return d.client.Call("StartNotify", 0).Store()
}

//StopNotify stop notifications
func (d *GattCharacteristic1) StopNotify() error {
	return d.client.Call("StopNotify", 0).Store()
}

//AcquireWrite acquire file descriptor and MTU for writing [experimental]
func (d *GattCharacteristic1) AcquireWrite() (dbus.UnixFD, uint16, error) {
	var fd dbus.UnixFD
	var mtu uint16
	err := d.client.Call("AcquireWrite", 0).Store(&fd, &mtu)
	return fd, mtu, err
}

//AcquireNotify acquire file descriptor and MTU for notify [experimental]
func (d *GattCharacteristic1) AcquireNotify() (dbus.UnixFD, uint16, error) {
	var fd dbus.UnixFD
	var mtu uint16
	err := d.client.Call("AcquireNotify", 0).Store(&fd, &mtu)
	return fd, mtu, err
}
<|endoftext|>


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

package tabletserver

import (
	"fmt"
	"math/rand"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysqlconn/fakesqldb"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletenv"
)

func TestTxPoolExecuteRollback(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	defer txPool.Rollback(ctx, transactionID)
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
}

func TestTxPoolRollbackNonBusy(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	txid1, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	_, err = txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	conn1, err := txPool.Get(txid1, "for query")
	if err != nil {
		t.Fatal(err)
	}
	// This should rollback only txid2.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 1 {
		t.Errorf("txPool.activePool.Size(): %d, want 1", sz)
	}
	conn1.Recycle()
	// This should rollback txid1.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 0 {
		t.Errorf("txPool.activePool.Size(): %d, want 0", sz)
	}
}

func TestTxPoolTransactionKiller(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	// make sure transaction killer will run frequent enough
	txPool.SetTimeout(1 * time.Millisecond)
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	killCount := tabletenv.KillStats.Counts()["Transactions"]
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	txConn.Recycle()
	// transaction killer should kill the query
	txPool.WaitForEmpty()
	killCountDiff := tabletenv.KillStats.Counts()["Transactions"] - killCount
	if killCountDiff != 1 {
		t.Fatalf("query: %s should be killed by transaction killer", sql)
	}
}

func TestTxPoolBeginAfterConnPoolClosed(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.SetTimeout(time.Duration(10))
	txPool.Open(db.ConnParams(), db.ConnParams())
	txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	if err == nil {
		t.Fatalf("expect to get an error")
	}
	terr, ok := err.(*tabletenv.TabletError)
	if !ok || terr != tabletenv.ErrConnPoolClosed {
		t.Fatalf("get error: %v, but expect: %v", terr, tabletenv.ErrConnPoolClosed)
	}
}

func TestTxPoolBeginWithPoolConnectionError(t *testing.T) {
	db := fakesqldb.New(t)
	db.Close()
	//	db.EnableConnFail()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "errno 2003"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolBeginWithExecError(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddRejectedQuery("begin", errRejected)
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolRollbackFail(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddRejectedQuery("rollback", errRejected)

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	err = txPool.Rollback(ctx, transactionID)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolGetConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	_, err := txPool.Get(12345, "for query")
	want := "not_in_tx: Transaction 12345: not found"
	if err == nil || err.Error() != want {
		t.Errorf("Get: %v, want %s", err, want)
	}
}

func TestTxPoolExecFailDueToConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	txPool.Begin(ctx)
	sql := "alter table test_table add test_column int"

	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	db.EnableConnFail()
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err == nil {
		t.Fatalf("exec should fail because of a conn error")
	}
}

func newTxPool() *TxPool {
	randID := rand.Int63()
	poolName := fmt.Sprintf("TestTransactionPool-%d", randID)
	transactionCap := 300
	transactionTimeout := time.Duration(30 * time.Second)
	idleTimeout := time.Duration(30 * time.Second)
	return NewTxPool(
		poolName,
		transactionCap,
		transactionTimeout,
		idleTimeout,
		DummyChecker,
	)
}
<commit_msg>tabletserver: Create separate test for stray transactions check.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"math/rand"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysqlconn/fakesqldb"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletenv"
)

func TestTxPoolExecuteRollback(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	defer txPool.Rollback(ctx, transactionID)
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
}

func TestTxPoolRollbackNonBusy(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	txid1, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	_, err = txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	conn1, err := txPool.Get(txid1, "for query")
	if err != nil {
		t.Fatal(err)
	}
	// This should rollback only txid2.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 1 {
		t.Errorf("txPool.activePool.Size(): %d, want 1", sz)
	}
	conn1.Recycle()
	// This should rollback txid1.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 0 {
		t.Errorf("txPool.activePool.Size(): %d, want 0", sz)
	}
}

func TestTxPoolTransactionKiller(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	// make sure transaction killer will run frequent enough
	txPool.SetTimeout(1 * time.Millisecond)
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	killCount := tabletenv.KillStats.Counts()["Transactions"]
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	txConn.Recycle()
	// transaction killer should kill the query
	txPool.WaitForEmpty()
	killCountDiff := tabletenv.KillStats.Counts()["Transactions"] - killCount
	if killCountDiff != 1 {
		t.Fatalf("query: %s should be killed by transaction killer", sql)
	}
}

func TestTxPoolBeginAfterConnPoolClosed(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.SetTimeout(time.Duration(10))
	txPool.Open(db.ConnParams(), db.ConnParams())
	txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	if err == nil {
		t.Fatalf("expect to get an error")
	}
	terr, ok := err.(*tabletenv.TabletError)
	if !ok || terr != tabletenv.ErrConnPoolClosed {
		t.Fatalf("get error: %v, but expect: %v", terr, tabletenv.ErrConnPoolClosed)
	}
}

func TestTxPoolBeginWithPoolConnectionError(t *testing.T) {
	db := fakesqldb.New(t)
	db.Close()
	//	db.EnableConnFail()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "errno 2003"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolBeginWithExecError(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddRejectedQuery("begin", errRejected)
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolRollbackFail(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddRejectedQuery("rollback", errRejected)

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	err = txPool.Rollback(ctx, transactionID)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolGetConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	_, err := txPool.Get(12345, "for query")
	want := "not_in_tx: Transaction 12345: not found"
	if err == nil || err.Error() != want {
		t.Errorf("Get: %v, want %s", err, want)
	}
}

func TestTxPoolExecFailDueToConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	sql := "alter table test_table add test_column int"

	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	db.EnableConnFail()
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err == nil {
		t.Fatalf("exec should fail because of a conn error")
	}
}

func TestTxPoolCloseKillsStrayTransactions(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())

	// Start stray transaction.
	_, err := txPool.Begin(context.Background())
	if err != nil {
		t.Fatal(err)
	}

	// Close kills stray transaction.
	txPool.Close()
	if got, want := tabletenv.InternalErrors.Counts()["StrayTransactions"], int64(1); got != want {
		t.Fatalf("internal error count for stray transactions not increased: got = %v, want = %v", got, want)
	}
	if got, want := txPool.conns.Capacity(), int64(0); got != want {
		t.Fatalf("resource pool was not closed. capacity: got = %v, want = %v", got, want)
	}
}

func newTxPool() *TxPool {
	randID := rand.Int63()
	poolName := fmt.Sprintf("TestTransactionPool-%d", randID)
	transactionCap := 300
	transactionTimeout := time.Duration(30 * time.Second)
	idleTimeout := time.Duration(30 * time.Second)
	return NewTxPool(
		poolName,
		transactionCap,
		transactionTimeout,
		idleTimeout,
		DummyChecker,
	)
}
<|endoftext|>


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

package planbuilder

import (
	"errors"
	"fmt"

	"github.com/youtube/vitess/go/vt/sqlparser"
)

// symtab contains the symbols for a SELECT
// statement.
type symtab struct {
	tables     []*tableAlias
	Colsyms    []*colsym
	FirstRoute *routeBuilder
	Externs    []*sqlparser.ColName
	Outer      *symtab
	VSchema    *VSchema
}

// colsym contains symbol info about a select expression.
type colsym struct {
	Alias      sqlparser.SQLName
	Route      *routeBuilder
	symtab     *symtab
	Underlying colref
	Vindex     Vindex
}

func newColsym(st *symtab) *colsym {
	return &colsym{
		Route:  st.FirstRoute,
		symtab: st,
	}
}

// colref uniquely identifies a column reference.
type colref struct {
	metadata interface{}
	name     sqlparser.SQLName
}

func newColref(col *sqlparser.ColName) colref {
	if col.Metadata == nil {
		panic("unexpected")
	}
	return colref{
		metadata: col.Metadata,
		name:     col.Name,
	}
}

// tableAlias is part of symtab.
// It represnts a table alias in a FROM clause.
// TODO(sougou): Update comments after the struct is finalized.
type tableAlias struct {
	Alias  sqlparser.SQLName
	Route  *routeBuilder
	symtab *symtab
	// Keyspace points to the keyspace to which this
	// alias belongs.
	Keyspace *Keyspace
	// CoVindexes is the list of column Vindexes for this alias.
	ColVindexes []*ColVindex
}

// FindVindex returns the vindex if one was found for the column.
func (t *tableAlias) FindVindex(name sqlparser.SQLName) Vindex {
	for _, colVindex := range t.ColVindexes {
		if string(name) == colVindex.Col {
			return colVindex.Vindex
		}
	}
	return nil
}

// newSymtab creates a new symtab initialized
// to contain the provided table alias.
func newSymtab(alias sqlparser.SQLName, table *Table, route *routeBuilder, vschema *VSchema) *symtab {
	st := &symtab{
		FirstRoute: route,
		VSchema:    vschema,
	}
	st.tables = []*tableAlias{{
		Alias:       alias,
		Route:       route,
		symtab:      st,
		Keyspace:    table.Keyspace,
		ColVindexes: table.ColVindexes,
	}}
	return st
}

// Add merges the new symbol table into the current one
// without merging their routes. This means that the new symbols
// will belong to different routes
func (st *symtab) Add(newsyms *symtab) error {
	for _, t := range newsyms.tables {
		if found := st.findTable(t.Alias); found != nil {
			return errors.New("duplicate symbols")
		}
		t.symtab = st
		st.tables = append(st.tables, t)
	}
	return nil
}

func (st *symtab) findTable(alias sqlparser.SQLName) *tableAlias {
	for i, t := range st.tables {
		if t.Alias == alias {
			return st.tables[i]
		}
	}
	return nil
}

// Merge merges the new symbol table into the current one and makes
// all the tables part of the specified routeBuilder.
func (st *symtab) Merge(newsyms *symtab, route *routeBuilder) error {
	for _, t := range st.tables {
		t.Route = route
	}
	for _, t := range newsyms.tables {
		if found := st.findTable(t.Alias); found != nil {
			return errors.New("duplicate symbols")
		}
		t.Route = route
		t.symtab = st
		st.tables = append(st.tables, t)
	}
	return nil
}

// SetRHS removes the ColVindexes from the aliases signifying
// that they cannot be used to make routing decisions. This is
// called if the table is in the RHS of a LEFT JOIN.
func (st *symtab) SetRHS() {
	for _, t := range st.tables {
		t.ColVindexes = nil
	}
}

// Find returns the routeBuilder for the symbol referenced by col.
// If a reference is found, the column's Metadata is set to point
// it. Subsequent searches will reuse this meatadata.
// If autoResolve is true, and there is only one table in the symbol table,
// then an unqualified reference is assumed to be implicitly against
// that table. The table info doesn't contain the full list of columns.
// So, any column reference is presumed valid. If a Colsyms scope is
// present, then the table scope is not searched. If a symbol is found
// in the current symtab, then isLocal is set to true. Otherwise, the
// search is continued in the outer symtab. If so, isLocal will be set
// to false. If the symbol was not found, an error is returned.
// isLocal must be checked before you can push-down (or pull-out)
// a construct.
func (st *symtab) Find(col *sqlparser.ColName, autoResolve bool) (route *routeBuilder, isLocal bool, err error) {
	switch m := col.Metadata.(type) {
	case *colsym:
		return m.Route, m.symtab == st, nil
	case *tableAlias:
		return m.Route, m.symtab == st, nil
	}
	if len(st.Colsyms) != 0 {
		name := sqlparser.SQLName(sqlparser.String(col))
		for _, colsym := range st.Colsyms {
			if name == colsym.Alias {
				col.Metadata = colsym
				return colsym.Route, true, nil
			}
		}
		st.Externs = append(st.Externs, col)
		if st.Outer != nil {
			// autoResolve only allowed for innermost scope.
			route, _, err = st.Outer.Find(col, false)
			return route, false, err
		}
		return nil, false, fmt.Errorf("symbol %s not found", sqlparser.String(col))
	}
	qualifier := col.Qualifier
	if qualifier == "" && autoResolve {
		if len(st.tables) != 1 {
			return nil, false, fmt.Errorf("symbol %s not found", sqlparser.String(col))
		}
		for _, t := range st.tables {
			qualifier = t.Alias
			break
		}
	}
	alias := st.findTable(qualifier)
	if alias == nil {
		st.Externs = append(st.Externs, col)
		if st.Outer != nil {
			// autoResolve only allowed for innermost scope.
			route, _, err = st.Outer.Find(col, false)
			return route, false, err
		}
		return nil, false, fmt.Errorf("symbol %s not found", sqlparser.String(col))
	}
	col.Metadata = alias
	return alias.Route, true, nil
}

// Vindex returns the vindex if the expression has an associated Vindex,
// but only if it's within the scope of routeBuilder. Because of this
// restriction, this function can be used to make push decisions.
func (st *symtab) Vindex(expr sqlparser.Expr, scope *routeBuilder, autoResolve bool) Vindex {
	col, ok := expr.(*sqlparser.ColName)
	if !ok {
		return nil
	}
	if col.Metadata == nil {
		_, _, err := st.Find(col, autoResolve)
		if err != nil {
			return nil
		}
	}
	switch meta := col.Metadata.(type) {
	case *colsym:
		if scope != meta.Route {
			return nil
		}
		return meta.Vindex
	case *tableAlias:
		if scope != meta.Route {
			return nil
		}
		return meta.FindVindex(col.Name)
	}
	return nil
}

// Reroute re-points the specified route to the new one.
func (st *symtab) Reroute(o, n *routeBuilder) {
	for _, t := range st.tables {
		if t.Route == o {
			t.Route = n
		}
	}
	for _, c := range st.Colsyms {
		if c.Route == o {
			c.Route = n
		}
	}
}
<commit_msg>v3: remove obsolete FirstRoute from symtab<commit_after>// Copyright 2016, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package planbuilder

import (
	"errors"
	"fmt"

	"github.com/youtube/vitess/go/vt/sqlparser"
)

// symtab contains the symbols for a SELECT
// statement.
type symtab struct {
	tables  []*tableAlias
	Colsyms []*colsym
	Externs []*sqlparser.ColName
	Outer   *symtab
	VSchema *VSchema
}

// colsym contains symbol info about a select expression.
type colsym struct {
	Alias      sqlparser.SQLName
	Route      *routeBuilder
	symtab     *symtab
	Underlying colref
	Vindex     Vindex
}

func newColsym(st *symtab) *colsym {
	return &colsym{
		symtab: st,
	}
}

// colref uniquely identifies a column reference.
type colref struct {
	metadata interface{}
	name     sqlparser.SQLName
}

func newColref(col *sqlparser.ColName) colref {
	if col.Metadata == nil {
		panic("unexpected")
	}
	return colref{
		metadata: col.Metadata,
		name:     col.Name,
	}
}

// tableAlias is part of symtab.
// It represnts a table alias in a FROM clause.
// TODO(sougou): Update comments after the struct is finalized.
type tableAlias struct {
	Alias  sqlparser.SQLName
	Route  *routeBuilder
	symtab *symtab
	// Keyspace points to the keyspace to which this
	// alias belongs.
	Keyspace *Keyspace
	// CoVindexes is the list of column Vindexes for this alias.
	ColVindexes []*ColVindex
}

// FindVindex returns the vindex if one was found for the column.
func (t *tableAlias) FindVindex(name sqlparser.SQLName) Vindex {
	for _, colVindex := range t.ColVindexes {
		if string(name) == colVindex.Col {
			return colVindex.Vindex
		}
	}
	return nil
}

// newSymtab creates a new symtab initialized
// to contain the provided table alias.
func newSymtab(alias sqlparser.SQLName, table *Table, route *routeBuilder, vschema *VSchema) *symtab {
	st := &symtab{
		VSchema: vschema,
	}
	st.tables = []*tableAlias{{
		Alias:       alias,
		Route:       route,
		symtab:      st,
		Keyspace:    table.Keyspace,
		ColVindexes: table.ColVindexes,
	}}
	return st
}

// Add merges the new symbol table into the current one
// without merging their routes. This means that the new symbols
// will belong to different routes
func (st *symtab) Add(newsyms *symtab) error {
	for _, t := range newsyms.tables {
		if found := st.findTable(t.Alias); found != nil {
			return errors.New("duplicate symbols")
		}
		t.symtab = st
		st.tables = append(st.tables, t)
	}
	return nil
}

func (st *symtab) findTable(alias sqlparser.SQLName) *tableAlias {
	for i, t := range st.tables {
		if t.Alias == alias {
			return st.tables[i]
		}
	}
	return nil
}

// Merge merges the new symbol table into the current one and makes
// all the tables part of the specified routeBuilder.
func (st *symtab) Merge(newsyms *symtab, route *routeBuilder) error {
	for _, t := range st.tables {
		t.Route = route
	}
	for _, t := range newsyms.tables {
		if found := st.findTable(t.Alias); found != nil {
			return errors.New("duplicate symbols")
		}
		t.Route = route
		t.symtab = st
		st.tables = append(st.tables, t)
	}
	return nil
}

// SetRHS removes the ColVindexes from the aliases signifying
// that they cannot be used to make routing decisions. This is
// called if the table is in the RHS of a LEFT JOIN.
func (st *symtab) SetRHS() {
	for _, t := range st.tables {
		t.ColVindexes = nil
	}
}

// Find returns the routeBuilder for the symbol referenced by col.
// If a reference is found, the column's Metadata is set to point
// it. Subsequent searches will reuse this meatadata.
// If autoResolve is true, and there is only one table in the symbol table,
// then an unqualified reference is assumed to be implicitly against
// that table. The table info doesn't contain the full list of columns.
// So, any column reference is presumed valid. If a Colsyms scope is
// present, then the table scope is not searched. If a symbol is found
// in the current symtab, then isLocal is set to true. Otherwise, the
// search is continued in the outer symtab. If so, isLocal will be set
// to false. If the symbol was not found, an error is returned.
// isLocal must be checked before you can push-down (or pull-out)
// a construct.
func (st *symtab) Find(col *sqlparser.ColName, autoResolve bool) (route *routeBuilder, isLocal bool, err error) {
	switch m := col.Metadata.(type) {
	case *colsym:
		return m.Route, m.symtab == st, nil
	case *tableAlias:
		return m.Route, m.symtab == st, nil
	}
	if len(st.Colsyms) != 0 {
		name := sqlparser.SQLName(sqlparser.String(col))
		for _, colsym := range st.Colsyms {
			if name == colsym.Alias {
				col.Metadata = colsym
				return colsym.Route, true, nil
			}
		}
		st.Externs = append(st.Externs, col)
		if st.Outer != nil {
			// autoResolve only allowed for innermost scope.
			route, _, err = st.Outer.Find(col, false)
			return route, false, err
		}
		return nil, false, fmt.Errorf("symbol %s not found", sqlparser.String(col))
	}
	qualifier := col.Qualifier
	if qualifier == "" && autoResolve {
		if len(st.tables) != 1 {
			return nil, false, fmt.Errorf("symbol %s not found", sqlparser.String(col))
		}
		for _, t := range st.tables {
			qualifier = t.Alias
			break
		}
	}
	alias := st.findTable(qualifier)
	if alias == nil {
		st.Externs = append(st.Externs, col)
		if st.Outer != nil {
			// autoResolve only allowed for innermost scope.
			route, _, err = st.Outer.Find(col, false)
			return route, false, err
		}
		return nil, false, fmt.Errorf("symbol %s not found", sqlparser.String(col))
	}
	col.Metadata = alias
	return alias.Route, true, nil
}

// Vindex returns the vindex if the expression has an associated Vindex,
// but only if it's within the scope of routeBuilder. Because of this
// restriction, this function can be used to make push decisions.
func (st *symtab) Vindex(expr sqlparser.Expr, scope *routeBuilder, autoResolve bool) Vindex {
	col, ok := expr.(*sqlparser.ColName)
	if !ok {
		return nil
	}
	if col.Metadata == nil {
		_, _, err := st.Find(col, autoResolve)
		if err != nil {
			return nil
		}
	}
	switch meta := col.Metadata.(type) {
	case *colsym:
		if scope != meta.Route {
			return nil
		}
		return meta.Vindex
	case *tableAlias:
		if scope != meta.Route {
			return nil
		}
		return meta.FindVindex(col.Name)
	}
	return nil
}

// Reroute re-points the specified route to the new one.
func (st *symtab) Reroute(o, n *routeBuilder) {
	for _, t := range st.tables {
		if t.Route == o {
			t.Route = n
		}
	}
	for _, c := range st.Colsyms {
		if c.Route == o {
			c.Route = n
		}
	}
}
<|endoftext|>


<commit_before>package session

import (
	"fmt"
	"testing"
)

func TestRoundTrip(t *testing.T) {
	sd := sessionData{
		Id:    generateSessionId(),
		Store: make(map[string]string),
	}

	sd.Set("TEST", "Fabulous")
	result := sd.Get("TEST")
	if result != "Fabulous" {
		t.Errorf("Expected Fabulous, got %s", result)
	}

	data, err := sd.Unmarshal()
	if err != nil {
		t.Error(err)
	}

	newSd := sessionData{}
	err = newSd.Marshal(data)
	if err != nil {
		t.Error(err)
	}

	result = newSd.Get("TEST")
	if result != "Fabulous" {
		t.Errorf("Expected Fabulous, got %s", result)
	}

	newSd.Clear()

	result = newSd.Get("TEST")
	if result != "Fabulous" {
		t.Errorf("Expected Fabulous, got %s", result)
	}
}
<commit_msg>Test non-existent keys<commit_after>package session

import (
	"testing"
)

func TestRoundTrip(t *testing.T) {
	sd := sessionData{
		Id:    generateSessionId(),
		Store: make(map[string]string),
	}

	sd.Set("TEST", "Fabulous")
	result := sd.Get("TEST")
	if result != "Fabulous" {
		t.Errorf("Expected Fabulous, got %s", result)
	}

	result = sd.Get("NONEXISTANT")
	if result != "" {
		t.Errorf("Expected '', got %s", result)
	}

	data, err := sd.Unmarshal()
	if err != nil {
		t.Error(err)
	}

	newSd := sessionData{}
	err = newSd.Marshal(data)
	if err != nil {
		t.Error(err)
	}

	result = newSd.Get("TEST")
	if result != "Fabulous" {
		t.Errorf("Expected Fabulous, got %s", result)
	}

	newSd.Clear()

	result = newSd.Get("TEST")
	if result != "Fabulous" {
		t.Errorf("Expected Fabulous, got %s", result)
	}
}
<|endoftext|>


<commit_before>package kingpin

import (
	"os"
	"path/filepath"
)

var (
	// CommandLine is the default Kingpin parser.
	CommandLine = New(filepath.Base(os.Args[0]), "")
)

// Command adds a new command to the default parser.
func Command(name, help string) *CmdClause {
	return CommandLine.Command(name, help)
}

// Flag adds a new flag to the default parser.
func Flag(name, help string) *FlagClause {
	return CommandLine.Flag(name, help)
}

// Arg adds a new argument to the top-level of the default parser.
func Arg(name, help string) *ArgClause {
	return CommandLine.Arg(name, help)
}

// Parse and return the selected command. Will call the termination handler if
// an error is encountered.
func Parse() string {
	selected := MustParse(CommandLine.Parse(os.Args[1:]))
	if selected == "" && CommandLine.cmdGroup.have() {
		Usage()
		CommandLine.terminate(0)
	}
	return selected
}

// Errorf prints an error message to stderr.
func Errorf(format string, args ...interface{}) {
	CommandLine.Errorf(format, args...)
}

// Fatalf prints an error message to stderr and exits.
func Fatalf(format string, args ...interface{}) {
	CommandLine.Fatalf(format, args...)
}

// FatalIfError prints an error and exits if err is not nil. The error is printed
// with the given prefix.
func FatalIfError(err error, prefix string) {
	CommandLine.FatalIfError(err, prefix)
}

// FatalUsage prints an error message followed by usage information, then
// exits with a non-zero status.
func FatalUsage(format string, args ...interface{}) {
	CommandLine.FatalUsage(format, args...)
}

// FatalUsageContext writes a printf formatted error message to stderr, then
// usage information for the given ParseContext, before exiting.
func FatalUsageContext(context *ParseContext, format string, args ...interface{}) {
	CommandLine.FatalUsageContext(context, format, args...)
}

// Usage prints usage to stderr.
func Usage() {
	CommandLine.Usage(os.Args[1:])
}

// Set global usage template to use (defaults to DefaultUsageTemplate).
func UsageTemplate(template string) *Application {
	return CommandLine.UsageTemplate(template)
}

// MustParse can be used with app.Parse(args) to exit with an error if parsing fails.
func MustParse(command string, err error) string {
	if err != nil {
		Fatalf("%s, try --help", err)
	}
	return command
}

// Version adds a flag for displaying the application version number.
func Version(version string) *Application {
	return CommandLine.Version(version)
}
<commit_msg>Missed a translation of FatalIfError to varargs.<commit_after>package kingpin

import (
	"os"
	"path/filepath"
)

var (
	// CommandLine is the default Kingpin parser.
	CommandLine = New(filepath.Base(os.Args[0]), "")
)

// Command adds a new command to the default parser.
func Command(name, help string) *CmdClause {
	return CommandLine.Command(name, help)
}

// Flag adds a new flag to the default parser.
func Flag(name, help string) *FlagClause {
	return CommandLine.Flag(name, help)
}

// Arg adds a new argument to the top-level of the default parser.
func Arg(name, help string) *ArgClause {
	return CommandLine.Arg(name, help)
}

// Parse and return the selected command. Will call the termination handler if
// an error is encountered.
func Parse() string {
	selected := MustParse(CommandLine.Parse(os.Args[1:]))
	if selected == "" && CommandLine.cmdGroup.have() {
		Usage()
		CommandLine.terminate(0)
	}
	return selected
}

// Errorf prints an error message to stderr.
func Errorf(format string, args ...interface{}) {
	CommandLine.Errorf(format, args...)
}

// Fatalf prints an error message to stderr and exits.
func Fatalf(format string, args ...interface{}) {
	CommandLine.Fatalf(format, args...)
}

// FatalIfError prints an error and exits if err is not nil. The error is printed
// with the given prefix.
func FatalIfError(err error, format string, args ...interface{}) {
	CommandLine.FatalIfError(err, format, args...)
}

// FatalUsage prints an error message followed by usage information, then
// exits with a non-zero status.
func FatalUsage(format string, args ...interface{}) {
	CommandLine.FatalUsage(format, args...)
}

// FatalUsageContext writes a printf formatted error message to stderr, then
// usage information for the given ParseContext, before exiting.
func FatalUsageContext(context *ParseContext, format string, args ...interface{}) {
	CommandLine.FatalUsageContext(context, format, args...)
}

// Usage prints usage to stderr.
func Usage() {
	CommandLine.Usage(os.Args[1:])
}

// Set global usage template to use (defaults to DefaultUsageTemplate).
func UsageTemplate(template string) *Application {
	return CommandLine.UsageTemplate(template)
}

// MustParse can be used with app.Parse(args) to exit with an error if parsing fails.
func MustParse(command string, err error) string {
	if err != nil {
		Fatalf("%s, try --help", err)
	}
	return command
}

// Version adds a flag for displaying the application version number.
func Version(version string) *Application {
	return CommandLine.Version(version)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

func main() {

	port := flag.String("port", "80", "Specifies the port to listen for incoming connections")
	useTls := flag.Bool("tls", false, "Tells gobble to listen for secure connections (ie. https)")
	tlsPort := flag.String("tlsPort", "443", "Specifies the port to listen for incoming secure connections")
	tlsCert := flag.String("tlsCert", "cert.pem", "Specifies the path to the x509 certificate")
	tlsKey := flag.String("tlsKey", "key.pem", "Specifies the path to the private key corresponding to the x509 certificate")
	usernameFlag := flag.String("username", "", "Specify a username to protect against unauthorized reading of your requests")
	passwordFlag := flag.String("password", "", "Specify a password to protect against unauthorized reading of your requests")

	homeDir := flag.String("dir", "public", "Speci	fies the root directory which all directories and requests will be stored under")
	flag.Parse()

	r := gin.Default()

	if *usernameFlag != "" {
		r.GET("/*path", gin.BasicAuth(gin.Accounts{*usernameFlag: *passwordFlag}), showFiles)
	} else {
		r.GET("/*path", showFiles)

	}
	r.POST("/*path", statusCodeHandler, handlePost)

	err := os.MkdirAll(*homeDir, 0744)
	if err != nil {
		panic("Unable to create home directory: " + err.Error())
	}
	err = os.Chdir(*homeDir)
	if err != nil {
		panic("Unable to switch to home directory: " + err.Error())
	}

	if *useTls == true {
		go func(tlsPort *string, tlsCert *string, tlsKey *string) {
			log.Println("Starting secure server on port " + *tlsPort)
			log.Fatal(r.RunTLS(":"+*tlsPort, *tlsCert, *tlsKey))
		}(tlsPort, tlsCert, tlsKey)
	}

	log.Println("Starting server on port " + *port)
	log.Fatal(r.Run(":" + *port))
}

func showFiles(c *gin.Context) {

	t := template.Must(template.New("index").Parse(`{{define "index"}}
		{{range .Files}}
		{{.Name}}

		{{end}}
		{{end}}`))

	path := c.Param("path")
	//Sanitize the path input and then add the '.' to keep the links relative to the working directory
	//This is necessary to keep badly maliciously formatted paths from escaping the working directory
	path = "." + filepath.Clean("/"+path)

	if info, err := os.Stat(path); err == nil {
		if info.IsDir() {
			files, err := ioutil.ReadDir(path)
			if err != nil {
				panic("Unable to read directory")
			}

			templData := struct {
				Path  string
				Files []os.FileInfo
			}{
				info.Name(),
				files,
			}

			t.Execute(c.Writer, templData)
		} else {
			f, _ := ioutil.ReadFile(path)
			c.Data(http.StatusOK, "text/plain", f)
		}
	} else {
		c.String(http.StatusNotFound, http.StatusText(http.StatusNotFound))
	}
}

func handlePost(c *gin.Context) {
	t := time.Now()
	dir := c.Query("dir")
	if dir != "" {
		//Make sure the requested directory is around
		dir = dir + "/" + t.Format("2006-01-02")
		err := os.MkdirAll(dir, 0744)
		if err != nil {
			panic("unable to create dir")
		}
	} else {
		//No directory requested so we give them the default
		dir = t.Format("2006-01-02")
		err := os.MkdirAll(dir, 0744)
		if err != nil {
			panic("unable to create dir")
		}
	}

	//Create file which is named after the create time
	fo, err := os.Create("./" + dir + "/" + t.Format("15.04.05.0000"))
	if err != nil {
		panic(err)
	}
	// close fo on exit and check for its returned error
	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	writer := bufio.NewWriter(fo)
	defer writer.Flush()

	fmt.Fprintln(writer, "POST ", c.FullPath())
	//Write headers to file
	for k, v := range c.Request.Header {
		fmt.Fprintln(writer, k+":", strings.Join(v, ","))
	}

	fmt.Fprintln(writer)
	//Write request body to file
	io.Copy(writer, c.Request.Body)
	c.String(http.StatusOK, fo.Name()[1:])
}

func statusCodeHandler(c *gin.Context) {
	if c.Query("status_code") != "" {
		status, err := strconv.Atoi(c.Query("status_code"))
		if err == nil {
			c.Status(status)
		} else {
			log.Println("Invalid number in status_code field")
		}
	}
	c.Next()
}
<commit_msg>Removed default gin logger.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

func main() {

	port := flag.String("port", "80", "Specifies the port to listen for incoming connections")
	useTls := flag.Bool("tls", false, "Tells gobble to listen for secure connections (ie. https)")
	tlsPort := flag.String("tlsPort", "443", "Specifies the port to listen for incoming secure connections")
	tlsCert := flag.String("tlsCert", "cert.pem", "Specifies the path to the x509 certificate")
	tlsKey := flag.String("tlsKey", "key.pem", "Specifies the path to the private key corresponding to the x509 certificate")
	usernameFlag := flag.String("username", "", "Specify a username to protect against unauthorized reading of your requests")
	passwordFlag := flag.String("password", "", "Specify a password to protect against unauthorized reading of your requests")

	homeDir := flag.String("dir", "public", "Speci	fies the root directory which all directories and requests will be stored under")
	flag.Parse()

	r := gin.New()
	r.Use(gin.Recovery())

	if *usernameFlag != "" {
		r.GET("/*path", gin.BasicAuth(gin.Accounts{*usernameFlag: *passwordFlag}), showFiles)
	} else {
		r.GET("/*path", showFiles)

	}
	r.POST("/*path", statusCodeHandler, handlePost)

	err := os.MkdirAll(*homeDir, 0744)
	if err != nil {
		panic("Unable to create home directory: " + err.Error())
	}
	err = os.Chdir(*homeDir)
	if err != nil {
		panic("Unable to switch to home directory: " + err.Error())
	}

	if *useTls == true {
		go func(tlsPort *string, tlsCert *string, tlsKey *string) {
			log.Println("Starting secure server on port " + *tlsPort)
			log.Fatal(r.RunTLS(":"+*tlsPort, *tlsCert, *tlsKey))
		}(tlsPort, tlsCert, tlsKey)
	}

	log.Println("Starting server on port " + *port)
	log.Fatal(r.Run(":" + *port))
}

func showFiles(c *gin.Context) {

	t := template.Must(template.New("index").Parse(`{{define "index"}}
		{{range .Files}}
		{{.Name}}

		{{end}}
		{{end}}`))

	path := c.Param("path")
	//Sanitize the path input and then add the '.' to keep the links relative to the working directory
	//This is necessary to keep badly maliciously formatted paths from escaping the working directory
	path = "." + filepath.Clean("/"+path)

	if info, err := os.Stat(path); err == nil {
		if info.IsDir() {
			files, err := ioutil.ReadDir(path)
			if err != nil {
				panic("Unable to read directory")
			}

			templData := struct {
				Path  string
				Files []os.FileInfo
			}{
				info.Name(),
				files,
			}

			t.Execute(c.Writer, templData)
		} else {
			f, _ := ioutil.ReadFile(path)
			c.Data(http.StatusOK, "text/plain", f)
		}
	} else {
		c.String(http.StatusNotFound, http.StatusText(http.StatusNotFound))
	}
}

func handlePost(c *gin.Context) {
	t := time.Now()
	dir := c.Query("dir")
	if dir != "" {
		//Make sure the requested directory is around
		dir = dir + "/" + t.Format("2006-01-02")
		err := os.MkdirAll(dir, 0744)
		if err != nil {
			panic("unable to create dir")
		}
	} else {
		//No directory requested so we give them the default
		dir = t.Format("2006-01-02")
		err := os.MkdirAll(dir, 0744)
		if err != nil {
			panic("unable to create dir")
		}
	}

	//Create file which is named after the create time
	fo, err := os.Create("./" + dir + "/" + t.Format("15.04.05.0000"))
	if err != nil {
		panic(err)
	}
	// close fo on exit and check for its returned error
	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	writer := bufio.NewWriter(fo)
	defer writer.Flush()

	fmt.Fprintln(writer, "POST ", c.FullPath())
	//Write headers to file
	for k, v := range c.Request.Header {
		fmt.Fprintln(writer, k+":", strings.Join(v, ","))
	}

	fmt.Fprintln(writer)
	//Write request body to file
	io.Copy(writer, c.Request.Body)
	c.String(http.StatusOK, fo.Name()[1:])
}

func statusCodeHandler(c *gin.Context) {
	if c.Query("status_code") != "" {
		status, err := strconv.Atoi(c.Query("status_code"))
		if err == nil {
			c.Status(status)
		} else {
			log.Println("Invalid number in status_code field")
		}
	}
	c.Next()
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/sean-duffy/xlsx"
	"strconv"
)

func main() {

	c := []xlsx.Column{
		xlsx.Column{Name: "Col1", Width: 10},
		xlsx.Column{Name: "Col2", Width: 10},
	}

	sh := xlsx.NewSheetWithColumns(c)
	sh.Title = "MySheet"

	for i := 0; i < 10; i++ {

		r := sh.NewRow()

		r.Cells[0] = xlsx.Cell{
			Type:  xlsx.CellTypeNumber,
			Value: strconv.Itoa(i + 1),
		}
		r.Cells[1] = xlsx.Cell{
			Type:  xlsx.CellTypeNumber,
			Value: "1",
		}

		sh.AppendRow(r)
	}

	err := sh.SaveToFile("test.xlsx")
	_ = err
}
<commit_msg>Change large example to use streaming<commit_after>package main

import (
	"bufio"
	"github.com/sean-duffy/xlsx"
	"os"
	"strconv"
)

func main() {

	outputfile, err := os.Create("test.xlsx")

	w := bufio.NewWriter(outputfile)
	ww := xlsx.NewWorkbookWriter(w)

	c := []xlsx.Column{
		xlsx.Column{Name: "Col1", Width: 10},
		xlsx.Column{Name: "Col2", Width: 10},
	}

	sh := xlsx.NewSheetWithColumns(c)
	sh.Title = "MySheet"

	sw, err := ww.NewSheetWriter(&sh)

	for i := 0; i < 1000000; i++ {

		r := sh.NewRow()

		r.Cells[0] = xlsx.Cell{
			Type:  xlsx.CellTypeNumber,
			Value: strconv.Itoa(i + 1),
		}
		r.Cells[1] = xlsx.Cell{
			Type:  xlsx.CellTypeNumber,
			Value: "1",
		}

		err = sw.WriteRows([]xlsx.Row{r})
	}

	err = ww.Close()
	defer w.Flush()

	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Landonia Ltd. All rights reserved.

package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"

	"github.com/landonia/golog"
	"github.com/landonia/gomost/proxy"
)

var (
	logger = golog.New("gomost.Main")
)

// bootstrap the application
func main() {

	// Get access to the settings
	configPath := flag.String("c", "", "The configuration file")
	flag.Parse()
	var config proxy.Configuration
	var err error
	if *configPath != "" {

		// parse the config if it is available
		config, err = proxy.ParseFileConfig(*configPath)
	} else {

		// otherwise create a basic config that will host the static files from
		// the current directory
		config = proxy.Configuration{
			StaticDir: ".",
		}
	}
	if err != nil {
		logger.Fatal("Could not parse configuration: %s", err.Error())
	}

	// Default the local host bind address
	if config.Host == "" {
		config.Host = ":8080"
	}
	golog.LogLevel(config.LogLevel)

	// initialise the server
	p, err := proxy.Setup(config)
	if err != nil {
		logger.Fatal("Could not start Gomost server: %s", err.Error())
	}

	// Wait for a shutdown signal
	go func() {
		sigs := make(chan os.Signal, 1)
		signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
		go func() {
			<-sigs
			p.Shutdown()
		}()
	}()

	// Handle any requests
	if err = p.Service(); err != nil {
		logger.Fatal("Error shutting down Gomost server: %s", err.Error())
	}
}
<commit_msg>Add in log statement to indicate when shutdown signals are received<commit_after>// Copyright 2016 Landonia Ltd. All rights reserved.

package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"

	"github.com/landonia/golog"
	"github.com/landonia/gomost/proxy"
)

var (
	logger = golog.New("gomost.Main")
)

// bootstrap the application
func main() {

	// Get access to the settings
	configPath := flag.String("c", "", "The configuration file")
	flag.Parse()
	var config proxy.Configuration
	var err error
	if *configPath != "" {

		// parse the config if it is available
		config, err = proxy.ParseFileConfig(*configPath)
	} else {

		// otherwise create a basic config that will host the static files from
		// the current directory
		config = proxy.Configuration{
			StaticDir: ".",
		}
	}
	if err != nil {
		logger.Fatal("Could not parse configuration: %s", err.Error())
	}

	// Default the local host bind address
	if config.Host == "" {
		config.Host = ":8080"
	}
	golog.LogLevel(config.LogLevel)

	// initialise the server
	p, err := proxy.Setup(config)
	if err != nil {
		logger.Fatal("Could not start Gomost server: %s", err.Error())
	}

	// Wait for a shutdown signal
	go func() {
		sigs := make(chan os.Signal, 1)
		signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
		go func() {
			<-sigs
			logger.Info("Received exit signal - shutting down")
			p.Shutdown()
		}()
	}()

	// Handle any requests
	if err = p.Service(); err != nil {
		logger.Fatal("Error shutting down Gomost server: %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import  (
        "fmt"
        "github.com/spf13/cobra"
        "github.com/rakyll/globalconf"
        "../go-notetxt"
        "flag"
        "os/user"
        "time"
        "strings"
        "strconv"
)

func main() {

        conf, _ := globalconf.New("gonote")

        var flagNotedir = flag.String("dir", "", "Location of the note.txt directory.")
        var dir string

        var today bool

        var cmdAdd = &cobra.Command{
            Use:   "add [title] [tag]",
            Short: "Add a note.",
            Long:  `Add a note and tag it.`,
            Run: func(cmd *cobra.Command, args []string) {
                if len(args) < 1 && !today {
                        fmt.Println("I need something to add.")
                        return
                }

                var text string
                t := time.Now().Local()

                if today {
                        text = fmt.Sprintf("Daily journal, date %s", t.Format("02. 01. 2006"))
                } else {
                        text = strings.Join(args, " ")
                }

                file, err := notetxt.CreateNote(text, t.Format("2006/01/"), dir)
                if err != nil {
                        panic(err);
                }

                notetxt.OpenFileInEditor(file)
            },
        }
        cmdAdd.Flags().BoolVarP(&today, "today", "T", false,
                                 "Add today's journal entry.")

        var cmdList = &cobra.Command{
            Use:   "list",
            Short: "List notes.",
            Long:  `List all valid note files in the directory.`,
            Run: func(cmd *cobra.Command, args []string) {
                notes, err := notetxt.ParseDir(dir)
                if err != nil {
                    panic(err)
                }

                needle := strings.Join(args, " ")

                for i, note := range notes {
                    if note.Matches(needle) {
                        fmt.Printf("%d %s - %v\n", i, note.Name, note.Tags)
                    }
                }
            },
        }

        var cmdEdit = &cobra.Command{
            Use:   "edit",
            Short: "Edit notes.",
            Long:  `Edit a note identified by either an ID or a selector.`,
            Run: func(cmd *cobra.Command, args []string) {
                noteid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Notes matching your query:\n")
                        cmdList.Run(cmd, args)
                }

                notes, err := notetxt.ParseDir(dir)
                if err != nil {
                    panic(err)
                }

                if noteid > len(notes) || noteid < 0 {
                        fmt.Printf("Invalid note ID (%v)\n", noteid)
                        return
                }

                notetxt.OpenFileInEditor(notes[noteid].Filename)

            },
        }

        var cmdTag = &cobra.Command{
            Use:   "tag <noteid> <tag-name>",
            Short: "Attaches a tag to a note.",
            Long:  `Tags a note with a one or more tags.`,
            Run: func(cmd *cobra.Command, args []string) {
                if len(args) < 2 {
                        fmt.Printf("Too few arguments.")
                }

                notes, err := notetxt.ParseDir(dir)
                if err != nil {
                    panic(err)
                }

                noteid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                if noteid > len(notes) || noteid < 0 {
                        fmt.Printf("Invalid note ID (%v)\n", noteid)
                        return
                }

                file := notes[noteid].Filename
                tag := args[1]
                err = notetxt.TagNote(file, tag, dir)
                if err != nil {
                        panic(err)
                }

            },
        }

        var GonoterCmd = &cobra.Command{
            Use:   "gonote",
            Short: "gonote is a go implementation of note.txt specification.",
            Long: `A small, fast and fun implementation of note.txt`,
            Run: func(cmd *cobra.Command, args []string) {
            },
        }

        GonoterCmd.PersistentFlags().StringVarP(&dir, "directory", "", "",
                                     "Location of the note.txt directory.")

        conf.ParseAll()
        if dir == "" {
                if *flagNotedir == "" {
                        usr, err := user.Current()
                        if err != nil {
                                panic(err)
                        }

                        dir = usr.HomeDir + "/notes"
                } else {
                        dir = *flagNotedir
                }
        }

        GonoterCmd.AddCommand(cmdAdd)
        GonoterCmd.AddCommand(cmdList)
        GonoterCmd.AddCommand(cmdTag)
        GonoterCmd.AddCommand(cmdEdit)
        GonoterCmd.Execute()
}
<commit_msg>fix: Fixed `edit` command behaviour<commit_after>package main

import  (
        "fmt"
        "github.com/spf13/cobra"
        "github.com/rakyll/globalconf"
        "../go-notetxt"
        "flag"
        "os/user"
        "time"
        "strings"
        "strconv"
)

func main() {

        conf, _ := globalconf.New("gonote")

        var flagNotedir = flag.String("dir", "", "Location of the note.txt directory.")
        var dir string

        var today bool

        var cmdAdd = &cobra.Command{
            Use:   "add [title] [tag]",
            Short: "Add a note.",
            Long:  `Add a note and tag it.`,
            Run: func(cmd *cobra.Command, args []string) {
                if len(args) < 1 && !today {
                        fmt.Println("I need something to add.")
                        return
                }

                var text string
                t := time.Now().Local()

                if today {
                        text = fmt.Sprintf("Daily journal, date %s", t.Format("02. 01. 2006"))
                } else {
                        text = strings.Join(args, " ")
                }

                file, err := notetxt.CreateNote(text, t.Format("2006/01/"), dir)
                if err != nil {
                        panic(err);
                }

                notetxt.OpenFileInEditor(file)
            },
        }
        cmdAdd.Flags().BoolVarP(&today, "today", "T", false,
                                 "Add today's journal entry.")

        var cmdList = &cobra.Command{
            Use:   "list",
            Short: "List notes.",
            Long:  `List all valid note files in the directory.`,
            Run: func(cmd *cobra.Command, args []string) {
                notes, err := notetxt.ParseDir(dir)
                if err != nil {
                    panic(err)
                }

                needle := strings.Join(args, " ")

                for i, note := range notes {
                    if note.Matches(needle) {
                        fmt.Printf("%d %s - %v\n", i, note.Name, note.Tags)
                    }
                }
            },
        }

        var cmdEdit = &cobra.Command{
            Use:   "edit",
            Short: "Edit notes.",
            Long:  `Edit a note identified by either an ID or a selector.`,
            Run: func(cmd *cobra.Command, args []string) {
                noteid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Notes matching your query:\n")
                        cmdList.Run(cmd, args)
                        return
                }

                notes, err := notetxt.ParseDir(dir)
                if err != nil {
                    panic(err)
                }

                if noteid > len(notes) || noteid < 0 {
                        fmt.Printf("Invalid note ID (%v)\n", noteid)
                        return
                }

                notetxt.OpenFileInEditor(notes[noteid].Filename)

            },
        }

        var cmdTag = &cobra.Command{
            Use:   "tag <noteid> <tag-name>",
            Short: "Attaches a tag to a note.",
            Long:  `Tags a note with a one or more tags.`,
            Run: func(cmd *cobra.Command, args []string) {
                if len(args) < 2 {
                        fmt.Printf("Too few arguments.")
                }

                notes, err := notetxt.ParseDir(dir)
                if err != nil {
                    panic(err)
                }

                noteid, err := strconv.Atoi(args[0])
                if err != nil {
                        fmt.Printf("Do you really consider that a number? %v\n", err)
                        return
                }

                if noteid > len(notes) || noteid < 0 {
                        fmt.Printf("Invalid note ID (%v)\n", noteid)
                        return
                }

                file := notes[noteid].Filename
                tag := args[1]
                err = notetxt.TagNote(file, tag, dir)
                if err != nil {
                        panic(err)
                }

            },
        }

        var GonoterCmd = &cobra.Command{
            Use:   "gonote",
            Short: "gonote is a go implementation of note.txt specification.",
            Long: `A small, fast and fun implementation of note.txt`,
            Run: func(cmd *cobra.Command, args []string) {
            },
        }

        GonoterCmd.PersistentFlags().StringVarP(&dir, "directory", "", "",
                                     "Location of the note.txt directory.")

        conf.ParseAll()
        if dir == "" {
                if *flagNotedir == "" {
                        usr, err := user.Current()
                        if err != nil {
                                panic(err)
                        }

                        dir = usr.HomeDir + "/notes"
                } else {
                        dir = *flagNotedir
                }
        }

        GonoterCmd.AddCommand(cmdAdd)
        GonoterCmd.AddCommand(cmdList)
        GonoterCmd.AddCommand(cmdTag)
        GonoterCmd.AddCommand(cmdEdit)
        GonoterCmd.Execute()
}
<|endoftext|>