<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsInternetGateway() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsInternetGatewayCreate,
		Read:   resourceAwsInternetGatewayRead,
		Update: resourceAwsInternetGatewayUpdate,
		Delete: resourceAwsInternetGatewayDelete,

		Schema: map[string]*schema.Schema{
			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"tags": tagsSchema(),
		},
	}
}

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

	// Create the gateway
	log.Printf("[DEBUG] Creating internet gateway")
	var err error
	resp, err := conn.CreateInternetGateway(nil)
	if err != nil {
		return fmt.Errorf("Error creating internet gateway: %s", err)
	}

	// Get the ID and store it
	ig := *resp.InternetGateway
	d.SetId(*ig.InternetGatewayId)
	log.Printf("[INFO] InternetGateway ID: %s", d.Id())

	err = setTags(conn, d)
	if err != nil {
		return err
	}

	// Attach the new gateway to the correct vpc
	return resourceAwsInternetGatewayAttach(d, meta)
}

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

	igRaw, _, err := IGStateRefreshFunc(conn, d.Id())()
	if err != nil {
		return err
	}
	if igRaw == nil {
		// Seems we have lost our internet gateway
		d.SetId("")
		return nil
	}

	ig := igRaw.(*ec2.InternetGateway)
	if len(ig.Attachments) == 0 {
		// Gateway exists but not attached to the VPC
		d.Set("vpc_id", "")
	} else {
		d.Set("vpc_id", ig.Attachments[0].VpcId)
	}

	d.Set("tags", tagsToMap(ig.Tags))

	return nil
}

func resourceAwsInternetGatewayUpdate(d *schema.ResourceData, meta interface{}) error {
	if d.HasChange("vpc_id") {
		// If we're already attached, detach it first
		if err := resourceAwsInternetGatewayDetach(d, meta); err != nil {
			return err
		}

		// Attach the gateway to the new vpc
		if err := resourceAwsInternetGatewayAttach(d, meta); err != nil {
			return err
		}
	}

	conn := meta.(*AWSClient).ec2conn

	if err := setTags(conn, d); err != nil {
		return err
	}

	d.SetPartial("tags")

	return nil
}

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

	// Detach if it is attached
	if err := resourceAwsInternetGatewayDetach(d, meta); err != nil {
		return err
	}

	log.Printf("[INFO] Deleting Internet Gateway: %s", d.Id())

	return resource.Retry(5*time.Minute, func() error {
		_, err := conn.DeleteInternetGateway(&ec2.DeleteInternetGatewayInput{
			InternetGatewayId: aws.String(d.Id()),
		})
		if err == nil {
			return nil
		}

		ec2err, ok := err.(awserr.Error)
		if !ok {
			return err
		}

		switch ec2err.Code() {
		case "InvalidInternetGatewayID.NotFound":
			return nil
		case "DependencyViolation":
			return err // retry
		}

		return resource.RetryError{Err: err}
	})
}

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

	if d.Get("vpc_id").(string) == "" {
		log.Printf(
			"[DEBUG] Not attaching Internet Gateway '%s' as no VPC ID is set",
			d.Id())
		return nil
	}

	log.Printf(
		"[INFO] Attaching Internet Gateway '%s' to VPC '%s'",
		d.Id(),
		d.Get("vpc_id").(string))

	_, err := conn.AttachInternetGateway(&ec2.AttachInternetGatewayInput{
		InternetGatewayId: aws.String(d.Id()),
		VpcId:             aws.String(d.Get("vpc_id").(string)),
	})
	if err != nil {
		return err
	}

	// A note on the states below: the AWS docs (as of July, 2014) say
	// that the states would be: attached, attaching, detached, detaching,
	// but when running, I noticed that the state is usually "available" when
	// it is attached.

	// Wait for it to be fully attached before continuing
	log.Printf("[DEBUG] Waiting for internet gateway (%s) to attach", d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{"detached", "attaching"},
		Target:  []string{"available"},
		Refresh: IGAttachStateRefreshFunc(conn, d.Id(), "available"),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for internet gateway (%s) to attach: %s",
			d.Id(), err)
	}

	return nil
}

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

	// Get the old VPC ID to detach from
	vpcID, _ := d.GetChange("vpc_id")

	if vpcID.(string) == "" {
		log.Printf(
			"[DEBUG] Not detaching Internet Gateway '%s' as no VPC ID is set",
			d.Id())
		return nil
	}

	log.Printf(
		"[INFO] Detaching Internet Gateway '%s' from VPC '%s'",
		d.Id(),
		vpcID.(string))

	// Wait for it to be fully detached before continuing
	log.Printf("[DEBUG] Waiting for internet gateway (%s) to detach", d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{"detaching"},
		Target:  []string{"detached"},
		Refresh: detachIGStateRefreshFunc(conn, d.Id(), vpcID.(string)),
		Timeout: 5 * time.Minute,
		Delay:   10 * time.Second,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for internet gateway (%s) to detach: %s",
			d.Id(), err)
	}

	return nil
}

// InstanceStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// an EC2 instance.
func detachIGStateRefreshFunc(conn *ec2.EC2, instanceID, vpcID string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		_, err := conn.DetachInternetGateway(&ec2.DetachInternetGatewayInput{
			InternetGatewayId: aws.String(instanceID),
			VpcId:             aws.String(vpcID),
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok {
				if ec2err.Code() == "InvalidInternetGatewayID.NotFound" {
					return nil, "Not Found", err
				} else if ec2err.Code() == "Gateway.NotAttached" {
					return "detached", "detached", nil
				} else if ec2err.Code() == "DependencyViolation" {
					return nil, "detaching", nil
				}
			}
		}
		// DetachInternetGateway only returns an error, so if it's nil, assume we're
		// detached
		return "detached", "detached", nil
	}
}

// IGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// an internet gateway.
func IGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeInternetGateways(&ec2.DescribeInternetGatewaysInput{
			InternetGatewayIds: []*string{aws.String(id)},
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok && ec2err.Code() == "InvalidInternetGatewayID.NotFound" {
				resp = nil
			} else {
				log.Printf("[ERROR] Error on IGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		ig := resp.InternetGateways[0]
		return ig, "available", nil
	}
}

// IGAttachStateRefreshFunc returns a resource.StateRefreshFunc that is used
// watch the state of an internet gateway's attachment.
func IGAttachStateRefreshFunc(conn *ec2.EC2, id string, expected string) resource.StateRefreshFunc {
	var start time.Time
	return func() (interface{}, string, error) {
		if start.IsZero() {
			start = time.Now()
		}

		resp, err := conn.DescribeInternetGateways(&ec2.DescribeInternetGatewaysInput{
			InternetGatewayIds: []*string{aws.String(id)},
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok && ec2err.Code() == "InvalidInternetGatewayID.NotFound" {
				resp = nil
			} else {
				log.Printf("[ERROR] Error on IGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		ig := resp.InternetGateways[0]

		if len(ig.Attachments) == 0 {
			// No attachments, we're detached
			return ig, "detached", nil
		}

		return ig, *ig.Attachments[0].State, nil
	}
}
<commit_msg>IGW detatch waiter tweak, 30 not found checks w/15 min max<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsInternetGateway() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsInternetGatewayCreate,
		Read:   resourceAwsInternetGatewayRead,
		Update: resourceAwsInternetGatewayUpdate,
		Delete: resourceAwsInternetGatewayDelete,

		Schema: map[string]*schema.Schema{
			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"tags": tagsSchema(),
		},
	}
}

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

	// Create the gateway
	log.Printf("[DEBUG] Creating internet gateway")
	var err error
	resp, err := conn.CreateInternetGateway(nil)
	if err != nil {
		return fmt.Errorf("Error creating internet gateway: %s", err)
	}

	// Get the ID and store it
	ig := *resp.InternetGateway
	d.SetId(*ig.InternetGatewayId)
	log.Printf("[INFO] InternetGateway ID: %s", d.Id())

	err = setTags(conn, d)
	if err != nil {
		return err
	}

	// Attach the new gateway to the correct vpc
	return resourceAwsInternetGatewayAttach(d, meta)
}

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

	igRaw, _, err := IGStateRefreshFunc(conn, d.Id())()
	if err != nil {
		return err
	}
	if igRaw == nil {
		// Seems we have lost our internet gateway
		d.SetId("")
		return nil
	}

	ig := igRaw.(*ec2.InternetGateway)
	if len(ig.Attachments) == 0 {
		// Gateway exists but not attached to the VPC
		d.Set("vpc_id", "")
	} else {
		d.Set("vpc_id", ig.Attachments[0].VpcId)
	}

	d.Set("tags", tagsToMap(ig.Tags))

	return nil
}

func resourceAwsInternetGatewayUpdate(d *schema.ResourceData, meta interface{}) error {
	if d.HasChange("vpc_id") {
		// If we're already attached, detach it first
		if err := resourceAwsInternetGatewayDetach(d, meta); err != nil {
			return err
		}

		// Attach the gateway to the new vpc
		if err := resourceAwsInternetGatewayAttach(d, meta); err != nil {
			return err
		}
	}

	conn := meta.(*AWSClient).ec2conn

	if err := setTags(conn, d); err != nil {
		return err
	}

	d.SetPartial("tags")

	return nil
}

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

	// Detach if it is attached
	if err := resourceAwsInternetGatewayDetach(d, meta); err != nil {
		return err
	}

	log.Printf("[INFO] Deleting Internet Gateway: %s", d.Id())

	return resource.Retry(5*time.Minute, func() error {
		_, err := conn.DeleteInternetGateway(&ec2.DeleteInternetGatewayInput{
			InternetGatewayId: aws.String(d.Id()),
		})
		if err == nil {
			return nil
		}

		ec2err, ok := err.(awserr.Error)
		if !ok {
			return err
		}

		switch ec2err.Code() {
		case "InvalidInternetGatewayID.NotFound":
			return nil
		case "DependencyViolation":
			return err // retry
		}

		return resource.RetryError{Err: err}
	})
}

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

	if d.Get("vpc_id").(string) == "" {
		log.Printf(
			"[DEBUG] Not attaching Internet Gateway '%s' as no VPC ID is set",
			d.Id())
		return nil
	}

	log.Printf(
		"[INFO] Attaching Internet Gateway '%s' to VPC '%s'",
		d.Id(),
		d.Get("vpc_id").(string))

	_, err := conn.AttachInternetGateway(&ec2.AttachInternetGatewayInput{
		InternetGatewayId: aws.String(d.Id()),
		VpcId:             aws.String(d.Get("vpc_id").(string)),
	})
	if err != nil {
		return err
	}

	// A note on the states below: the AWS docs (as of July, 2014) say
	// that the states would be: attached, attaching, detached, detaching,
	// but when running, I noticed that the state is usually "available" when
	// it is attached.

	// Wait for it to be fully attached before continuing
	log.Printf("[DEBUG] Waiting for internet gateway (%s) to attach", d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{"detached", "attaching"},
		Target:  []string{"available"},
		Refresh: IGAttachStateRefreshFunc(conn, d.Id(), "available"),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for internet gateway (%s) to attach: %s",
			d.Id(), err)
	}

	return nil
}

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

	// Get the old VPC ID to detach from
	vpcID, _ := d.GetChange("vpc_id")

	if vpcID.(string) == "" {
		log.Printf(
			"[DEBUG] Not detaching Internet Gateway '%s' as no VPC ID is set",
			d.Id())
		return nil
	}

	log.Printf(
		"[INFO] Detaching Internet Gateway '%s' from VPC '%s'",
		d.Id(),
		vpcID.(string))

	// Wait for it to be fully detached before continuing
	log.Printf("[DEBUG] Waiting for internet gateway (%s) to detach", d.Id())
	stateConf := &resource.StateChangeConf{
		Pending:        []string{"detaching"},
		Target:         []string{"detached"},
		Refresh:        detachIGStateRefreshFunc(conn, d.Id(), vpcID.(string)),
		Timeout:        15 * time.Minute,
		Delay:          10 * time.Second,
		NotFoundChecks: 30,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for internet gateway (%s) to detach: %s",
			d.Id(), err)
	}

	return nil
}

// InstanceStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// an EC2 instance.
func detachIGStateRefreshFunc(conn *ec2.EC2, instanceID, vpcID string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		_, err := conn.DetachInternetGateway(&ec2.DetachInternetGatewayInput{
			InternetGatewayId: aws.String(instanceID),
			VpcId:             aws.String(vpcID),
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok {
				if ec2err.Code() == "InvalidInternetGatewayID.NotFound" {
					return nil, "Not Found", err
				} else if ec2err.Code() == "Gateway.NotAttached" {
					return "detached", "detached", nil
				} else if ec2err.Code() == "DependencyViolation" {
					return nil, "detaching", nil
				}
			}
		}
		// DetachInternetGateway only returns an error, so if it's nil, assume we're
		// detached
		return "detached", "detached", nil
	}
}

// IGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// an internet gateway.
func IGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeInternetGateways(&ec2.DescribeInternetGatewaysInput{
			InternetGatewayIds: []*string{aws.String(id)},
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok && ec2err.Code() == "InvalidInternetGatewayID.NotFound" {
				resp = nil
			} else {
				log.Printf("[ERROR] Error on IGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		ig := resp.InternetGateways[0]
		return ig, "available", nil
	}
}

// IGAttachStateRefreshFunc returns a resource.StateRefreshFunc that is used
// watch the state of an internet gateway's attachment.
func IGAttachStateRefreshFunc(conn *ec2.EC2, id string, expected string) resource.StateRefreshFunc {
	var start time.Time
	return func() (interface{}, string, error) {
		if start.IsZero() {
			start = time.Now()
		}

		resp, err := conn.DescribeInternetGateways(&ec2.DescribeInternetGatewaysInput{
			InternetGatewayIds: []*string{aws.String(id)},
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok && ec2err.Code() == "InvalidInternetGatewayID.NotFound" {
				resp = nil
			} else {
				log.Printf("[ERROR] Error on IGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		ig := resp.InternetGateways[0]

		if len(ig.Attachments) == 0 {
			// No attachments, we're detached
			return ig, "detached", nil
		}

		return ig, *ig.Attachments[0].State, nil
	}
}
<|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 bigquery

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"cloud.google.com/go/civil"
	"cloud.google.com/go/internal/pretty"

	bq "google.golang.org/api/bigquery/v2"
)

func (fs *FieldSchema) GoString() string {
	if fs == nil {
		return "<nil>"
	}

	return fmt.Sprintf("{Name:%s Description:%s Repeated:%t Required:%t Type:%s Schema:%s}",
		fs.Name,
		fs.Description,
		fs.Repeated,
		fs.Required,
		fs.Type,
		fmt.Sprintf("%#v", fs.Schema),
	)
}

func bqTableFieldSchema(desc, name, typ, mode string) *bq.TableFieldSchema {
	return &bq.TableFieldSchema{
		Description: desc,
		Name:        name,
		Mode:        mode,
		Type:        typ,
	}
}

func fieldSchema(desc, name, typ string, repeated, required bool) *FieldSchema {
	return &FieldSchema{
		Description: desc,
		Name:        name,
		Repeated:    repeated,
		Required:    required,
		Type:        FieldType(typ),
	}
}

func TestSchemaConversion(t *testing.T) {
	testCases := []struct {
		schema   Schema
		bqSchema *bq.TableSchema
	}{
		{
			// required
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "STRING", "REQUIRED"),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "STRING", false, true),
			},
		},
		{
			// repeated
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "STRING", "REPEATED"),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "STRING", true, false),
			},
		},
		{
			// nullable, string
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "STRING", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "STRING", false, false),
			},
		},
		{
			// integer
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "INTEGER", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "INTEGER", false, false),
			},
		},
		{
			// float
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "FLOAT", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "FLOAT", false, false),
			},
		},
		{
			// boolean
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "BOOLEAN", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "BOOLEAN", false, false),
			},
		},
		{
			// timestamp
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "TIMESTAMP", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "TIMESTAMP", false, false),
			},
		},
		{
			// civil times
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "f1", "TIME", ""),
					bqTableFieldSchema("desc", "f2", "DATE", ""),
					bqTableFieldSchema("desc", "f3", "DATETIME", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "f1", "TIME", false, false),
				fieldSchema("desc", "f2", "DATE", false, false),
				fieldSchema("desc", "f3", "DATETIME", false, false),
			},
		},
		{
			// nested
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					{
						Description: "An outer schema wrapping a nested schema",
						Name:        "outer",
						Mode:        "REQUIRED",
						Type:        "RECORD",
						Fields: []*bq.TableFieldSchema{
							bqTableFieldSchema("inner field", "inner", "STRING", ""),
						},
					},
				},
			},
			schema: Schema{
				&FieldSchema{
					Description: "An outer schema wrapping a nested schema",
					Name:        "outer",
					Required:    true,
					Type:        "RECORD",
					Schema: []*FieldSchema{
						{
							Description: "inner field",
							Name:        "inner",
							Type:        "STRING",
						},
					},
				},
			},
		},
	}

	for _, tc := range testCases {
		bqSchema := tc.schema.asTableSchema()
		if !reflect.DeepEqual(bqSchema, tc.bqSchema) {
			t.Errorf("converting to TableSchema: got:\n%v\nwant:\n%v",
				pretty.Value(bqSchema), pretty.Value(tc.bqSchema))
		}
		schema := convertTableSchema(tc.bqSchema)
		if !reflect.DeepEqual(schema, tc.schema) {
			t.Errorf("converting to Schema: got:\n%v\nwant:\n%v", schema, tc.schema)
		}
	}
}

type allStrings struct {
	String    string
	ByteSlice []byte
}

type allSignedIntegers struct {
	Int64 int64
	Int32 int32
	Int16 int16
	Int8  int8
	Int   int
}

type allUnsignedIntegers struct {
	Uint32 uint32
	Uint16 uint16
	Uint8  uint8
}

type allFloat struct {
	Float64 float64
	Float32 float32
	// NOTE: Complex32 and Complex64 are unsupported by BigQuery
}

type allBoolean struct {
	Bool bool
}

type allTime struct {
	Timestamp time.Time
	Time      civil.Time
	Date      civil.Date
	DateTime  civil.DateTime
}

func reqField(name, typ string) *FieldSchema {
	return &FieldSchema{
		Name:     name,
		Type:     FieldType(typ),
		Required: true,
	}
}

func TestSimpleInference(t *testing.T) {
	testCases := []struct {
		in   interface{}
		want Schema
	}{
		{
			in: allSignedIntegers{},
			want: Schema{
				reqField("Int64", "INTEGER"),
				reqField("Int32", "INTEGER"),
				reqField("Int16", "INTEGER"),
				reqField("Int8", "INTEGER"),
				reqField("Int", "INTEGER"),
			},
		},
		{
			in: allUnsignedIntegers{},
			want: Schema{
				reqField("Uint32", "INTEGER"),
				reqField("Uint16", "INTEGER"),
				reqField("Uint8", "INTEGER"),
			},
		},
		{
			in: allFloat{},
			want: Schema{
				reqField("Float64", "FLOAT"),
				reqField("Float32", "FLOAT"),
			},
		},
		{
			in: allBoolean{},
			want: Schema{
				reqField("Bool", "BOOLEAN"),
			},
		},
		{
			in: &allBoolean{},
			want: Schema{
				reqField("Bool", "BOOLEAN"),
			},
		},
		{
			in: allTime{},
			want: Schema{
				reqField("Timestamp", "TIMESTAMP"),
				reqField("Time", "TIME"),
				reqField("Date", "DATE"),
				reqField("DateTime", "DATETIME"),
			},
		},
		{
			in: allStrings{},
			want: Schema{
				reqField("String", "STRING"),
				reqField("ByteSlice", "BYTES"),
			},
		},
	}
	for i, tc := range testCases {
		got, err := InferSchema(tc.in)
		if err != nil {
			t.Fatalf("%d: error inferring TableSchema: %v", i, err)
		}
		if !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i,
				pretty.Value(got), pretty.Value(tc.want))
		}
	}
}

type containsNested struct {
	hidden    string
	NotNested int
	Nested    struct {
		Inside int
	}
}

type containsDoubleNested struct {
	NotNested int
	Nested    struct {
		InsideNested struct {
			Inside int
		}
	}
}

type ptrNested struct {
	Ptr *struct{ Inside int }
}

func TestNestedInference(t *testing.T) {
	testCases := []struct {
		in   interface{}
		want Schema
	}{
		{
			in: containsNested{},
			want: Schema{
				reqField("NotNested", "INTEGER"),
				&FieldSchema{
					Name:     "Nested",
					Required: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
			},
		},
		{
			in: containsDoubleNested{},
			want: Schema{
				reqField("NotNested", "INTEGER"),
				&FieldSchema{
					Name:     "Nested",
					Required: true,
					Type:     "RECORD",
					Schema: Schema{
						{
							Name:     "InsideNested",
							Required: true,
							Type:     "RECORD",
							Schema:   Schema{reqField("Inside", "INTEGER")},
						},
					},
				},
			},
		},
		{
			in: ptrNested{},
			want: Schema{
				&FieldSchema{
					Name:     "Ptr",
					Required: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
			},
		},
	}

	for i, tc := range testCases {
		got, err := InferSchema(tc.in)
		if err != nil {
			t.Fatalf("%d: error inferring TableSchema: %v", i, err)
		}
		if !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i, got, tc.want)
		}
	}
}

type simpleRepeated struct {
	NotRepeated       []byte
	RepeatedByteSlice [][]byte
	Repeated          []int
}

type simpleNestedRepeated struct {
	NotRepeated int
	Repeated    []struct {
		Inside int
	}
}

type moreRepeated struct {
	Array    [5]int
	PtrArray *[5]int
	PtrSlice *[]int
}

func repField(name, typ string) *FieldSchema {
	return &FieldSchema{
		Name:     name,
		Type:     FieldType(typ),
		Repeated: true,
	}
}

func TestRepeatedInference(t *testing.T) {
	testCases := []struct {
		in   interface{}
		want Schema
	}{
		{
			in: simpleRepeated{},
			want: Schema{
				reqField("NotRepeated", "BYTES"),
				repField("RepeatedByteSlice", "BYTES"),
				repField("Repeated", "INTEGER"),
			},
		},
		{
			in: simpleNestedRepeated{},
			want: Schema{
				reqField("NotRepeated", "INTEGER"),
				&FieldSchema{
					Name:     "Repeated",
					Repeated: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
			},
		},
	}

	for i, tc := range testCases {
		got, err := InferSchema(tc.in)
		if err != nil {
			t.Fatalf("%d: error inferring TableSchema: %v", i, err)
		}
		if !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i, got, tc.want)
		}
	}
}

type Embedded struct {
	Embedded int
}

type embedded struct {
	Embedded2 int
}

type nestedEmbedded struct {
	Embedded
	embedded
}

func TestEmbeddedInference(t *testing.T) {
	got, err := InferSchema(nestedEmbedded{})
	if err != nil {
		t.Fatal(err)
	}
	want := Schema{
		reqField("Embedded", "INTEGER"),
		reqField("Embedded2", "INTEGER"),
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", pretty.Value(got), pretty.Value(want))
	}
}

func TestSchemaErrors(t *testing.T) {
	testCases := []struct {
		in  interface{}
		err error
	}{
		{
			in:  []byte{},
			err: errNoStruct,
		},
		{
			in:  new(int),
			err: errNoStruct,
		},
		{
			in:  struct{ Uint uint }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Uint64 uint64 }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Uintptr uintptr }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Complex complex64 }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Map map[string]int }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Chan chan bool }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Ptr *int }{},
			err: errNoStruct,
		},
		{
			in:  struct{ Interface interface{} }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ MultiDimensional [][]int }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ MultiDimensional [][][]byte }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ ChanSlice []chan bool }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ NestedChan struct{ Chan []chan bool } }{},
			err: errUnsupportedFieldType,
		},
	}
	for _, tc := range testCases {
		want := tc.err
		_, got := InferSchema(tc.in)
		if !reflect.DeepEqual(got, want) {
			t.Errorf("%#v: got:\n%#v\nwant:\n%#v", tc.in, got, want)
		}
	}
}
<commit_msg>bigquery: improve tests of repeated-field inference<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 bigquery

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"cloud.google.com/go/civil"
	"cloud.google.com/go/internal/pretty"

	bq "google.golang.org/api/bigquery/v2"
)

func (fs *FieldSchema) GoString() string {
	if fs == nil {
		return "<nil>"
	}

	return fmt.Sprintf("{Name:%s Description:%s Repeated:%t Required:%t Type:%s Schema:%s}",
		fs.Name,
		fs.Description,
		fs.Repeated,
		fs.Required,
		fs.Type,
		fmt.Sprintf("%#v", fs.Schema),
	)
}

func bqTableFieldSchema(desc, name, typ, mode string) *bq.TableFieldSchema {
	return &bq.TableFieldSchema{
		Description: desc,
		Name:        name,
		Mode:        mode,
		Type:        typ,
	}
}

func fieldSchema(desc, name, typ string, repeated, required bool) *FieldSchema {
	return &FieldSchema{
		Description: desc,
		Name:        name,
		Repeated:    repeated,
		Required:    required,
		Type:        FieldType(typ),
	}
}

func TestSchemaConversion(t *testing.T) {
	testCases := []struct {
		schema   Schema
		bqSchema *bq.TableSchema
	}{
		{
			// required
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "STRING", "REQUIRED"),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "STRING", false, true),
			},
		},
		{
			// repeated
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "STRING", "REPEATED"),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "STRING", true, false),
			},
		},
		{
			// nullable, string
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "STRING", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "STRING", false, false),
			},
		},
		{
			// integer
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "INTEGER", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "INTEGER", false, false),
			},
		},
		{
			// float
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "FLOAT", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "FLOAT", false, false),
			},
		},
		{
			// boolean
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "BOOLEAN", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "BOOLEAN", false, false),
			},
		},
		{
			// timestamp
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "name", "TIMESTAMP", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "name", "TIMESTAMP", false, false),
			},
		},
		{
			// civil times
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					bqTableFieldSchema("desc", "f1", "TIME", ""),
					bqTableFieldSchema("desc", "f2", "DATE", ""),
					bqTableFieldSchema("desc", "f3", "DATETIME", ""),
				},
			},
			schema: Schema{
				fieldSchema("desc", "f1", "TIME", false, false),
				fieldSchema("desc", "f2", "DATE", false, false),
				fieldSchema("desc", "f3", "DATETIME", false, false),
			},
		},
		{
			// nested
			bqSchema: &bq.TableSchema{
				Fields: []*bq.TableFieldSchema{
					{
						Description: "An outer schema wrapping a nested schema",
						Name:        "outer",
						Mode:        "REQUIRED",
						Type:        "RECORD",
						Fields: []*bq.TableFieldSchema{
							bqTableFieldSchema("inner field", "inner", "STRING", ""),
						},
					},
				},
			},
			schema: Schema{
				&FieldSchema{
					Description: "An outer schema wrapping a nested schema",
					Name:        "outer",
					Required:    true,
					Type:        "RECORD",
					Schema: []*FieldSchema{
						{
							Description: "inner field",
							Name:        "inner",
							Type:        "STRING",
						},
					},
				},
			},
		},
	}

	for _, tc := range testCases {
		bqSchema := tc.schema.asTableSchema()
		if !reflect.DeepEqual(bqSchema, tc.bqSchema) {
			t.Errorf("converting to TableSchema: got:\n%v\nwant:\n%v",
				pretty.Value(bqSchema), pretty.Value(tc.bqSchema))
		}
		schema := convertTableSchema(tc.bqSchema)
		if !reflect.DeepEqual(schema, tc.schema) {
			t.Errorf("converting to Schema: got:\n%v\nwant:\n%v", schema, tc.schema)
		}
	}
}

type allStrings struct {
	String    string
	ByteSlice []byte
}

type allSignedIntegers struct {
	Int64 int64
	Int32 int32
	Int16 int16
	Int8  int8
	Int   int
}

type allUnsignedIntegers struct {
	Uint32 uint32
	Uint16 uint16
	Uint8  uint8
}

type allFloat struct {
	Float64 float64
	Float32 float32
	// NOTE: Complex32 and Complex64 are unsupported by BigQuery
}

type allBoolean struct {
	Bool bool
}

type allTime struct {
	Timestamp time.Time
	Time      civil.Time
	Date      civil.Date
	DateTime  civil.DateTime
}

func reqField(name, typ string) *FieldSchema {
	return &FieldSchema{
		Name:     name,
		Type:     FieldType(typ),
		Required: true,
	}
}

func TestSimpleInference(t *testing.T) {
	testCases := []struct {
		in   interface{}
		want Schema
	}{
		{
			in: allSignedIntegers{},
			want: Schema{
				reqField("Int64", "INTEGER"),
				reqField("Int32", "INTEGER"),
				reqField("Int16", "INTEGER"),
				reqField("Int8", "INTEGER"),
				reqField("Int", "INTEGER"),
			},
		},
		{
			in: allUnsignedIntegers{},
			want: Schema{
				reqField("Uint32", "INTEGER"),
				reqField("Uint16", "INTEGER"),
				reqField("Uint8", "INTEGER"),
			},
		},
		{
			in: allFloat{},
			want: Schema{
				reqField("Float64", "FLOAT"),
				reqField("Float32", "FLOAT"),
			},
		},
		{
			in: allBoolean{},
			want: Schema{
				reqField("Bool", "BOOLEAN"),
			},
		},
		{
			in: &allBoolean{},
			want: Schema{
				reqField("Bool", "BOOLEAN"),
			},
		},
		{
			in: allTime{},
			want: Schema{
				reqField("Timestamp", "TIMESTAMP"),
				reqField("Time", "TIME"),
				reqField("Date", "DATE"),
				reqField("DateTime", "DATETIME"),
			},
		},
		{
			in: allStrings{},
			want: Schema{
				reqField("String", "STRING"),
				reqField("ByteSlice", "BYTES"),
			},
		},
	}
	for i, tc := range testCases {
		got, err := InferSchema(tc.in)
		if err != nil {
			t.Fatalf("%d: error inferring TableSchema: %v", i, err)
		}
		if !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i,
				pretty.Value(got), pretty.Value(tc.want))
		}
	}
}

type containsNested struct {
	hidden    string
	NotNested int
	Nested    struct {
		Inside int
	}
}

type containsDoubleNested struct {
	NotNested int
	Nested    struct {
		InsideNested struct {
			Inside int
		}
	}
}

type ptrNested struct {
	Ptr *struct{ Inside int }
}

func TestNestedInference(t *testing.T) {
	testCases := []struct {
		in   interface{}
		want Schema
	}{
		{
			in: containsNested{},
			want: Schema{
				reqField("NotNested", "INTEGER"),
				&FieldSchema{
					Name:     "Nested",
					Required: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
			},
		},
		{
			in: containsDoubleNested{},
			want: Schema{
				reqField("NotNested", "INTEGER"),
				&FieldSchema{
					Name:     "Nested",
					Required: true,
					Type:     "RECORD",
					Schema: Schema{
						{
							Name:     "InsideNested",
							Required: true,
							Type:     "RECORD",
							Schema:   Schema{reqField("Inside", "INTEGER")},
						},
					},
				},
			},
		},
		{
			in: ptrNested{},
			want: Schema{
				&FieldSchema{
					Name:     "Ptr",
					Required: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
			},
		},
	}

	for i, tc := range testCases {
		got, err := InferSchema(tc.in)
		if err != nil {
			t.Fatalf("%d: error inferring TableSchema: %v", i, err)
		}
		if !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i,
				pretty.Value(got), pretty.Value(tc.want))
		}
	}
}

type repeated struct {
	NotRepeated       []byte
	RepeatedByteSlice [][]byte
	Slice             []int
	Array             [5]bool
}

type nestedRepeated struct {
	NotRepeated int
	Repeated    []struct {
		Inside int
	}
	RepeatedPtr []*struct{ Inside int }
}

func repField(name, typ string) *FieldSchema {
	return &FieldSchema{
		Name:     name,
		Type:     FieldType(typ),
		Repeated: true,
	}
}

func TestRepeatedInference(t *testing.T) {
	testCases := []struct {
		in   interface{}
		want Schema
	}{
		{
			in: repeated{},
			want: Schema{
				reqField("NotRepeated", "BYTES"),
				repField("RepeatedByteSlice", "BYTES"),
				repField("Slice", "INTEGER"),
				repField("Array", "BOOLEAN"),
			},
		},
		{
			in: nestedRepeated{},
			want: Schema{
				reqField("NotRepeated", "INTEGER"),
				{
					Name:     "Repeated",
					Repeated: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
				{
					Name:     "RepeatedPtr",
					Repeated: true,
					Type:     "RECORD",
					Schema:   Schema{reqField("Inside", "INTEGER")},
				},
			},
		},
	}

	for i, tc := range testCases {
		got, err := InferSchema(tc.in)
		if err != nil {
			t.Fatalf("%d: error inferring TableSchema: %v", i, err)
		}
		if !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i,
				pretty.Value(got), pretty.Value(tc.want))
		}
	}
}

type Embedded struct {
	Embedded int
}

type embedded struct {
	Embedded2 int
}

type nestedEmbedded struct {
	Embedded
	embedded
}

func TestEmbeddedInference(t *testing.T) {
	got, err := InferSchema(nestedEmbedded{})
	if err != nil {
		t.Fatal(err)
	}
	want := Schema{
		reqField("Embedded", "INTEGER"),
		reqField("Embedded2", "INTEGER"),
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", pretty.Value(got), pretty.Value(want))
	}
}

func TestSchemaErrors(t *testing.T) {
	testCases := []struct {
		in  interface{}
		err error
	}{
		{
			in:  []byte{},
			err: errNoStruct,
		},
		{
			in:  new(int),
			err: errNoStruct,
		},
		{
			in:  struct{ Uint uint }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Uint64 uint64 }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Uintptr uintptr }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Complex complex64 }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Map map[string]int }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Chan chan bool }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ Ptr *int }{},
			err: errNoStruct,
		},
		{
			in:  struct{ Interface interface{} }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ MultiDimensional [][]int }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ MultiDimensional [][][]byte }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ ChanSlice []chan bool }{},
			err: errUnsupportedFieldType,
		},
		{
			in:  struct{ NestedChan struct{ Chan []chan bool } }{},
			err: errUnsupportedFieldType,
		},
	}
	for _, tc := range testCases {
		want := tc.err
		_, got := InferSchema(tc.in)
		if !reflect.DeepEqual(got, want) {
			t.Errorf("%#v: got:\n%#v\nwant:\n%#v", tc.in, got, want)
		}
	}
}
<|endoftext|>


<commit_before>/*
	Copyright (c) 2014-2015, Percona LLC and/or its affiliates. All rights reserved.

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU Affero General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Affero General Public License for more details.

	You should have received a copy of the GNU Affero General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package event

import (
	"github.com/percona/go-mysql/log"
	"time"
)

const (
	MAX_EXAMPLE_BYTES = 1024 * 10
)

// A QueryClass represents all events with the same fingerprint and class ID.
// This is only enforced by convention, so be careful not to mix events from
// different classes.
type QueryClass struct {
	Id           string   // 16-character hex checksum of fingerprint
	Fingerprint  string   // canonical form of query: values replaced with "?"
	Metrics      *Metrics // statistics for each metric, e.g. max Query_time
	TotalQueries uint64   // total number of queries in class
	Example      Example  `json:",omitempty"` // example query with max Query_time
	lastDb       string
	example      bool
}

// An Example is a real query and its database, timestamp, and Query_time.
// If the query is larger than MAX_EXAMPLE_BYTES, it is truncated and "..."
// is appended.
type Example struct {
	QueryTime float64 // Query_time
	Db        string  // Schema: <db> or USE <db>
	Query     string  // truncated to MAX_EXAMPLE_BYTES
	Ts        string  `json:",omitempty"` // in MySQL time zone
}

// NewQueryClass returns a new QueryClass for the class ID and fingerprint.
// If example is true, the query with the greatest Query_time is saved.
func NewQueryClass(classId string, fingerprint string, example bool) *QueryClass {
	class := &QueryClass{
		Id:           classId,
		Fingerprint:  fingerprint,
		Metrics:      NewMetrics(),
		TotalQueries: 0,
		example:      example,
	}
	return class
}

// AddEvent adds an event to the query class.
func (c *QueryClass) AddEvent(e *log.Event) {
	c.TotalQueries++
	c.Metrics.AddEvent(e)
	// Save last db seen for this query. This helps ensure the example query
	// has a db.
	if e.Db != "" {
		c.lastDb = e.Db
	}
	if c.example {
		if n, ok := e.TimeMetrics["Query_time"]; ok {
			if float64(n) > c.Example.QueryTime {
				c.Example.QueryTime = float64(n)
				if e.Db != "" {
					c.Example.Db = e.Db
				} else {
					c.Example.Db = c.lastDb
				}
				if len(e.Query) > MAX_EXAMPLE_BYTES {
					c.Example.Query = e.Query[0:MAX_EXAMPLE_BYTES-3] + "..."
				} else {
					c.Example.Query = e.Query
				}
				if e.Ts != "" {
					if t, err := time.Parse("060102 15:04:05", e.Ts); err != nil {
						c.Example.Ts = ""
					} else {
						c.Example.Ts = t.Format("2006-01-02 15:04:05")
					}
				} else {
					c.Example.Ts = ""
				}
			}
		}
	}
}

// Finalize calculates all metric statistics. Call this function when done
// adding events to the class.
func (c *QueryClass) Finalize() {
	c.Metrics.Finalize()
}
<commit_msg>Omit query class Example if there's none.<commit_after>/*
	Copyright (c) 2014-2015, Percona LLC and/or its affiliates. All rights reserved.

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU Affero General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Affero General Public License for more details.

	You should have received a copy of the GNU Affero General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package event

import (
	"github.com/percona/go-mysql/log"
	"time"
)

const (
	MAX_EXAMPLE_BYTES = 1024 * 10
)

// A QueryClass represents all events with the same fingerprint and class ID.
// This is only enforced by convention, so be careful not to mix events from
// different classes.
type QueryClass struct {
	Id           string   // 16-character hex checksum of fingerprint
	Fingerprint  string   // canonical form of query: values replaced with "?"
	Metrics      *Metrics // statistics for each metric, e.g. max Query_time
	TotalQueries uint64   // total number of queries in class
	Example      *Example `json:",omitempty"` // example query with max Query_time
	lastDb       string
	example      bool
}

// An Example is a real query and its database, timestamp, and Query_time.
// If the query is larger than MAX_EXAMPLE_BYTES, it is truncated and "..."
// is appended.
type Example struct {
	QueryTime float64 // Query_time
	Db        string  // Schema: <db> or USE <db>
	Query     string  // truncated to MAX_EXAMPLE_BYTES
	Ts        string  `json:",omitempty"` // in MySQL time zone
}

// NewQueryClass returns a new QueryClass for the class ID and fingerprint.
// If example is true, the query with the greatest Query_time is saved.
func NewQueryClass(classId string, fingerprint string, example bool) *QueryClass {
	class := &QueryClass{
		Id:           classId,
		Fingerprint:  fingerprint,
		Metrics:      NewMetrics(),
		TotalQueries: 0,
		Example:      &Example{},
		example:      example,
	}
	return class
}

// AddEvent adds an event to the query class.
func (c *QueryClass) AddEvent(e *log.Event) {
	c.TotalQueries++
	c.Metrics.AddEvent(e)
	// Save last db seen for this query. This helps ensure the example query
	// has a db.
	if e.Db != "" {
		c.lastDb = e.Db
	}
	if c.example {
		if n, ok := e.TimeMetrics["Query_time"]; ok {
			if float64(n) > c.Example.QueryTime {
				c.Example.QueryTime = float64(n)
				if e.Db != "" {
					c.Example.Db = e.Db
				} else {
					c.Example.Db = c.lastDb
				}
				if len(e.Query) > MAX_EXAMPLE_BYTES {
					c.Example.Query = e.Query[0:MAX_EXAMPLE_BYTES-3] + "..."
				} else {
					c.Example.Query = e.Query
				}
				if e.Ts != "" {
					if t, err := time.Parse("060102 15:04:05", e.Ts); err != nil {
						c.Example.Ts = ""
					} else {
						c.Example.Ts = t.Format("2006-01-02 15:04:05")
					}
				} else {
					c.Example.Ts = ""
				}
			}
		}
	}
}

// Finalize calculates all metric statistics. Call this function when done
// adding events to the class.
func (c *QueryClass) Finalize() {
	c.Metrics.Finalize()
	if c.Example.QueryTime == 0 {
		c.Example = nil
	}
}
<|endoftext|>


<commit_before>package service

import (
	"bytes"
	"github.com/timeredbull/commandmocker"
	"github.com/timeredbull/tsuru/api/app"
	"github.com/timeredbull/tsuru/db"
	"github.com/timeredbull/tsuru/log"
	"labix.org/v2/mgo/bson"
	. "launchpad.net/gocheck"
	stdlog "log"
	"strings"
)

func (s *ServiceSuite) createService() {
	s.serviceType = &ServiceType{Name: "Mysql", Charm: "mysql"}
	s.serviceType.Create()
	s.service = &Service{ServiceTypeId: s.serviceType.Id, Name: "my_service"}
	s.service.Create()
}

func (s *ServiceSuite) TestGetService(c *C) {
	s.createService()
	anotherService := Service{Name: s.service.Name}
	anotherService.Get()
	c.Assert(anotherService.Name, Equals, s.service.Name)
	c.Assert(anotherService.ServiceTypeId, Equals, s.service.ServiceTypeId)
}

func (s *ServiceSuite) TestAllServices(c *C) {
	st := ServiceType{Name: "mysql", Charm: "mysql"}
	st.Create()
	se := Service{ServiceTypeId: st.Id, Name: "myService"}
	se2 := Service{ServiceTypeId: st.Id, Name: "myOtherService"}
	se.Create()
	se2.Create()

	s_ := Service{}
	results := s_.All()
	c.Assert(len(results), Equals, 2)
}

func (s *ServiceSuite) TestCreateService(c *C) {
	w := bytes.NewBuffer([]byte{})
	l := stdlog.New(w, "", stdlog.LstdFlags)
	log.Target = l
	dir, err := commandmocker.Add("juju", "$*")
	c.Assert(err, IsNil)
	defer commandmocker.Remove(dir)
	s.createService()
	se := Service{Name: s.service.Name}
	se.Get()
	c.Assert(se.ServiceTypeId, Equals, s.serviceType.Id)
	c.Assert(se.Name, Equals, s.service.Name)
	strOut := strings.Replace(w.String(), "\n", "", -1)
	c.Assert(strOut, Matches, ".*deploy --repository=/home/charms mysql my_service")
}

func (s *ServiceSuite) TestDeleteService(c *C) {
	s.createService()
	s.service.Delete()
	qtd, err := db.Session.Services().Find(nil).Count()
	c.Assert(err, IsNil)
	c.Assert(qtd, Equals, 0)
}

func (s *ServiceSuite) TestRetrieveAssociateServiceType(c *C) {
	serviceType := ServiceType{Name: "Mysql", Charm: "mysql"}
	serviceType.Create()

	service := &Service{
		ServiceTypeId: serviceType.Id,
		Name:          "my_service",
	}
	service.Create()
	retrievedServiceType := service.ServiceType()

	c.Assert(serviceType.Id, Equals, retrievedServiceType.Id)
	c.Assert(serviceType.Name, Equals, retrievedServiceType.Name)
	c.Assert(serviceType.Charm, Equals, retrievedServiceType.Charm)
}

func (s *ServiceSuite) TestBindService(c *C) {
	s.createService()
	app := &app.App{Name: "my_app", Framework: "django"}
	app.Create()
	s.service.Bind(app)
	var result ServiceApp
	query := bson.M{
		"service_name": s.service.Name,
		"app_name":     app.Name,
	}
	err := db.Session.ServiceApps().Find(query).One(&result)
	c.Assert(err, IsNil)
	c.Assert(s.service.Name, Equals, result.ServiceName)
	c.Assert(app.Name, Equals, result.AppName)
}

func (s *ServiceSuite) TestUnbindService(c *C) {
	s.createService()
	app := &app.App{Name: "my_app", Framework: "django"}
	app.Create()
	s.service.Bind(app)
	s.service.Unbind(app)
	query := bson.M{
		"service_name": s.service.Name,
		"app_name":     app.Name,
	}
	qtd, err := db.Session.ServiceApps().Find(query).Count()
	c.Assert(err, IsNil)
	c.Assert(qtd, Equals, 0)
}

func (s *ServiceSuite) TestGrantAccessShouldAddTeamToTheService(c *C) {
	s.createService()
	err := s.service.GrantAccess(s.team)
	c.Assert(err, IsNil)
	c.Assert(*s.team, HasAccessTo, *s.service)
}

func (s *ServiceSuite) TestGrantAccessShouldReturnErrorIfTheTeamAlreadyHasAcessToTheService(c *C) {
	s.createService()
	err := s.service.GrantAccess(s.team)
	c.Assert(err, IsNil)
	err = s.service.GrantAccess(s.team)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, "^This team already has access to this service$")
}

func (s *ServiceSuite) TestRevokeAccessShouldRemoveTeamFromService(c *C) {
	s.createService()
	err := s.service.GrantAccess(s.team)
	c.Assert(err, IsNil)
	err = s.service.RevokeAccess(s.team)
	c.Assert(err, IsNil)
	c.Assert(*s.team, Not(HasAccessTo), *s.service)
}

func (s *ServiceSuite) TestRevokeAcessShouldReturnErrorIfTheTeamDoesNotHaveAccessToTheService(c *C) {
	s.createService()
	err := s.service.RevokeAccess(s.team)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, "^This team does not have access to this service$")
}

func (s *ServiceSuite) TestCheckUserPermissionShouldReturnTrueIfTheGivenUserIsMemberOfOneOfTheServicesTeam(c *C) {
	s.createService()
	s.service.GrantAccess(s.team)
	c.Assert(s.service.CheckUserAccess(s.user), Equals, true)
}

func (s *ServiceSuite) TestCheckUserPermissionShouldReturnFalseIfTheGivenUserIsNotMemberOfAnyOfTheServicesTeam(c *C) {
	s.createService()
	c.Assert(s.service.CheckUserAccess(s.user), Equals, false)
}
<commit_msg>api/service: fix intermitent assertion<commit_after>package service

import (
	"bytes"
	"github.com/timeredbull/commandmocker"
	"github.com/timeredbull/tsuru/api/app"
	"github.com/timeredbull/tsuru/db"
	"github.com/timeredbull/tsuru/log"
	"labix.org/v2/mgo/bson"
	. "launchpad.net/gocheck"
	stdlog "log"
	"strings"
)

func (s *ServiceSuite) createService() {
	s.serviceType = &ServiceType{Name: "Mysql", Charm: "mysql"}
	s.serviceType.Create()
	s.service = &Service{ServiceTypeId: s.serviceType.Id, Name: "my_service"}
	s.service.Create()
}

func (s *ServiceSuite) TestGetService(c *C) {
	s.createService()
	anotherService := Service{Name: s.service.Name}
	anotherService.Get()
	c.Assert(anotherService.Name, Equals, s.service.Name)
	c.Assert(anotherService.ServiceTypeId, Equals, s.service.ServiceTypeId)
}

func (s *ServiceSuite) TestAllServices(c *C) {
	st := ServiceType{Name: "mysql", Charm: "mysql"}
	st.Create()
	se := Service{ServiceTypeId: st.Id, Name: "myService"}
	se2 := Service{ServiceTypeId: st.Id, Name: "myOtherService"}
	se.Create()
	se2.Create()

	s_ := Service{}
	results := s_.All()
	c.Assert(len(results), Equals, 2)
}

func (s *ServiceSuite) TestCreateService(c *C) {
	w := bytes.NewBuffer([]byte{})
	l := stdlog.New(w, "", stdlog.LstdFlags)
	log.Target = l
	dir, err := commandmocker.Add("juju", "$*")
	c.Assert(err, IsNil)
	defer commandmocker.Remove(dir)
	s.createService()
	se := Service{Name: s.service.Name}
	se.Get()
	c.Assert(se.ServiceTypeId, Equals, s.serviceType.Id)
	c.Assert(se.Name, Equals, s.service.Name)
	strOut := strings.Replace(w.String(), "\n", "", -1)
	c.Assert(strOut, Matches, ".*deploy --repository=/home/charms mysql my_service.*")
}

func (s *ServiceSuite) TestDeleteService(c *C) {
	s.createService()
	s.service.Delete()
	qtd, err := db.Session.Services().Find(nil).Count()
	c.Assert(err, IsNil)
	c.Assert(qtd, Equals, 0)
}

func (s *ServiceSuite) TestRetrieveAssociateServiceType(c *C) {
	serviceType := ServiceType{Name: "Mysql", Charm: "mysql"}
	serviceType.Create()

	service := &Service{
		ServiceTypeId: serviceType.Id,
		Name:          "my_service",
	}
	service.Create()
	retrievedServiceType := service.ServiceType()

	c.Assert(serviceType.Id, Equals, retrievedServiceType.Id)
	c.Assert(serviceType.Name, Equals, retrievedServiceType.Name)
	c.Assert(serviceType.Charm, Equals, retrievedServiceType.Charm)
}

func (s *ServiceSuite) TestBindService(c *C) {
	s.createService()
	app := &app.App{Name: "my_app", Framework: "django"}
	app.Create()
	s.service.Bind(app)
	var result ServiceApp
	query := bson.M{
		"service_name": s.service.Name,
		"app_name":     app.Name,
	}
	err := db.Session.ServiceApps().Find(query).One(&result)
	c.Assert(err, IsNil)
	c.Assert(s.service.Name, Equals, result.ServiceName)
	c.Assert(app.Name, Equals, result.AppName)
}

func (s *ServiceSuite) TestUnbindService(c *C) {
	s.createService()
	app := &app.App{Name: "my_app", Framework: "django"}
	app.Create()
	s.service.Bind(app)
	s.service.Unbind(app)
	query := bson.M{
		"service_name": s.service.Name,
		"app_name":     app.Name,
	}
	qtd, err := db.Session.ServiceApps().Find(query).Count()
	c.Assert(err, IsNil)
	c.Assert(qtd, Equals, 0)
}

func (s *ServiceSuite) TestGrantAccessShouldAddTeamToTheService(c *C) {
	s.createService()
	err := s.service.GrantAccess(s.team)
	c.Assert(err, IsNil)
	c.Assert(*s.team, HasAccessTo, *s.service)
}

func (s *ServiceSuite) TestGrantAccessShouldReturnErrorIfTheTeamAlreadyHasAcessToTheService(c *C) {
	s.createService()
	err := s.service.GrantAccess(s.team)
	c.Assert(err, IsNil)
	err = s.service.GrantAccess(s.team)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, "^This team already has access to this service$")
}

func (s *ServiceSuite) TestRevokeAccessShouldRemoveTeamFromService(c *C) {
	s.createService()
	err := s.service.GrantAccess(s.team)
	c.Assert(err, IsNil)
	err = s.service.RevokeAccess(s.team)
	c.Assert(err, IsNil)
	c.Assert(*s.team, Not(HasAccessTo), *s.service)
}

func (s *ServiceSuite) TestRevokeAcessShouldReturnErrorIfTheTeamDoesNotHaveAccessToTheService(c *C) {
	s.createService()
	err := s.service.RevokeAccess(s.team)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, "^This team does not have access to this service$")
}

func (s *ServiceSuite) TestCheckUserPermissionShouldReturnTrueIfTheGivenUserIsMemberOfOneOfTheServicesTeam(c *C) {
	s.createService()
	s.service.GrantAccess(s.team)
	c.Assert(s.service.CheckUserAccess(s.user), Equals, true)
}

func (s *ServiceSuite) TestCheckUserPermissionShouldReturnFalseIfTheGivenUserIsNotMemberOfAnyOfTheServicesTeam(c *C) {
	s.createService()
	c.Assert(s.service.CheckUserAccess(s.user), Equals, false)
}
<|endoftext|>


<commit_before>package witness

import (
	"github.com/asuleymanov/golos-go/types"
)

//CurrentMedianHistoryPrice structure for the GetCurrentMedianHistoryPrice function. And additional structure for the functions GetFeedHistory, GetWitnessByAccount, GetWitnesses and GetWitnessByVote.
type CurrentMedianHistoryPrice struct {
	Base  *types.Asset `json:"base"`
	Quote *types.Asset `json:"quote"`
}

//FeedHistory structure for the GetFeedHistory function.
type FeedHistory struct {
	ID                   *types.Int                   `json:"id"`
	CurrentMedianHistory *CurrentMedianHistoryPrice   `json:"current_median_history"`
	PriceHistory         []*CurrentMedianHistoryPrice `json:"price_history"`
}

//WitnessSchedule structure for the GetWitnessSchedule function.
type WitnessSchedule struct {
	ID                            *types.Int             `json:"id"`
	CurrentVirtualTime            string                 `json:"current_virtual_time"`
	NextShuffleBlockNum           *types.Int             `json:"next_shuffle_block_num"`
	CurrentShuffledWitnesses      string                 `json:"current_shuffled_witnesses"`
	NumScheduledWitnesses         *types.Int             `json:"num_scheduled_witnesses"`
	Top19Weight                   *types.Int             `json:"top19_weight"`
	TimeshareWeight               *types.Int             `json:"timeshare_weight"`
	MinerWeight                   *types.Int             `json:"miner_weight"`
	WitnessPayNormalizationFactor *types.Int             `json:"witness_pay_normalization_factor"`
	MedianProps                   *types.ChainProperties `json:"median_props"`
	MajorityVersion               string                 `json:"majority_version"`
}

//Witness structure for the GetWitnessByAccount, GetWitnesses and GetWitnessByVote function.
type Witness struct {
	ID                    *types.Int                 `json:"id"`
	Owner                 string                     `json:"owner"`
	Created               *types.Time                `json:"created"`
	URL                   string                     `json:"url"`
	Votes                 string                     `json:"votes"`
	VirtualLastUpdate     string                     `json:"virtual_last_update"`
	VirtualPosition       string                     `json:"virtual_position"`
	VirtualScheduledTime  string                     `json:"virtual_scheduled_time"`
	TotalMissed           *types.Int                 `json:"total_missed"`
	LastAslot             *types.Int                 `json:"last_aslot"`
	LastConfirmedBlockNum *types.Int                 `json:"last_confirmed_block_num"`
	PowWorker             *types.Int                 `json:"pow_worker"`
	SigningKey            string                     `json:"signing_key"`
	Props                 *types.ChainProperties     `json:"props"`
	SbdExchangeRate       *CurrentMedianHistoryPrice `json:"sbd_exchange_rate"`
	LastSbdExchangeUpdate *types.Time                `json:"last_sbd_exchange_update"`
	LastWork              string                     `json:"last_work"`
	RunningVersion        string                     `json:"running_version"`
	HardforkVersionVote   string                     `json:"hardfork_version_vote"`
	HardforkTimeVote      *types.Time                `json:"hardfork_time_vote"`
}
<commit_msg>Update data.go<commit_after>package witness

import (
	"github.com/asuleymanov/golos-go/types"
)

//CurrentMedianHistoryPrice structure for the GetCurrentMedianHistoryPrice function. And additional structure for the functions GetFeedHistory, GetWitnessByAccount, GetWitnesses and GetWitnessByVote.
type CurrentMedianHistoryPrice struct {
	Base  *types.Asset `json:"base"`
	Quote *types.Asset `json:"quote"`
}

//FeedHistory structure for the GetFeedHistory function.
type FeedHistory struct {
	ID                   *types.Int                   `json:"id"`
	CurrentMedianHistory *CurrentMedianHistoryPrice   `json:"current_median_history"`
	PriceHistory         []*CurrentMedianHistoryPrice `json:"price_history"`
}

//WitnessSchedule structure for the GetWitnessSchedule function.
type WitnessSchedule struct {
	ID                            *types.Int             `json:"id"`
	CurrentVirtualTime            string                 `json:"current_virtual_time"`
	NextShuffleBlockNum           *types.Int             `json:"next_shuffle_block_num"`
	CurrentShuffledWitnesses      string                 `json:"current_shuffled_witnesses"`
	NumScheduledWitnesses         *types.Int             `json:"num_scheduled_witnesses"`
	Top19Weight                   *types.Int             `json:"top19_weight"`
	TimeshareWeight               *types.Int             `json:"timeshare_weight"`
	MinerWeight                   *types.Int             `json:"miner_weight"`
	WitnessPayNormalizationFactor *types.Int             `json:"witness_pay_normalization_factor"`
	MedianProps                   *types.ChainProperties `json:"median_props"`
	MajorityVersion               string                 `json:"majority_version"`
}

//Witness structure for the GetWitnessByAccount, GetWitnesses and GetWitnessByVote function.
type Witness struct {
	ID                    *types.Int                 `json:"id"`
	Owner                 string                     `json:"owner"`
	Created               *types.Time                `json:"created"`
	URL                   string                     `json:"url"`
	Votes                 *types.Int                 `json:"votes"`
	VirtualLastUpdate     *types.Int                 `json:"virtual_last_update"`
	VirtualPosition       *types.Int                 `json:"virtual_position"`
	VirtualScheduledTime  *types.Int                 `json:"virtual_scheduled_time"`
	TotalMissed           *types.Int                 `json:"total_missed"`
	LastAslot             *types.Int                 `json:"last_aslot"`
	LastConfirmedBlockNum *types.Int                 `json:"last_confirmed_block_num"`
	PowWorker             *types.Int                 `json:"pow_worker"`
	SigningKey            string                     `json:"signing_key"`
	Props                 *types.ChainProperties     `json:"props"`
	SbdExchangeRate       *CurrentMedianHistoryPrice `json:"sbd_exchange_rate"`
	LastSbdExchangeUpdate *types.Time                `json:"last_sbd_exchange_update"`
	LastWork              string                     `json:"last_work"`
	RunningVersion        string                     `json:"running_version"`
	HardforkVersionVote   string                     `json:"hardfork_version_vote"`
	HardforkTimeVote      *types.Time                `json:"hardfork_time_vote"`
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

var (
	// Protect access to dup
	mu sync.Mutex
	// Duplicates table
	dup = map[string]bool{}

	// Command-line flags
	seed        = flag.String("seed", "http://golang.org", "seed URL")
	cancelAfter = flag.Duration("cancelafter", 0, "automatically cancel the fetchbot after a given time")
	cancelAtURL = flag.String("cancelat", "", "automatically cancel the fetchbot at a given URL")
	stopAfter   = flag.Duration("stopafter", 0, "automatically stop the fetchbot after a given time")
	stopAtURL   = flag.String("stopat", "", "automatically stop the fetchbot at a given URL")
	memStats    = flag.Duration("memstats", 0, "display memory statistics at a given interval")
)

func main() {
	flag.Parse()

	// Parse the provided seed
	u, err := url.Parse(*seed)
	if err != nil {
		log.Fatal(err)
	}

	// Create the muxer
	mux := fetchbot.NewMux()

	// Handle all errors the same
	mux.HandleErrors(fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	}))

	// Handle GET requests for html responses, to parse the body and enqueue all links as HEAD
	// requests.
	mux.Response().Method("GET").ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			// Process the body to find the links
			doc, err := goquery.NewDocumentFromResponse(res)
			if err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}
			// Enqueue all links as HEAD requests
			enqueueLinks(ctx, doc)
		}))

	// Handle HEAD requests for html responses coming from the source host - we don't want
	// to crawl links from other hosts.
	mux.Response().Method("HEAD").Host(u.Host).ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			if _, err := ctx.Q.SendStringGet(ctx.Cmd.URL().String()); err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			}
		}))

	// Create the Fetcher, handle the logging first, then dispatch to the Muxer
	h := logHandler(mux)
	if *stopAtURL != "" || *cancelAtURL != "" {
		h = stopHandler(*stopAtURL, logHandler(mux))
	}
	f := fetchbot.New(h)
	// First mem stat print must be right after creating the fetchbot
	if *memStats > 0 {
		// Print starting stats
		printMemStats(nil)
		// Run at regular intervals
		runMemStats(f, *memStats)
		// On exit, print ending stats after a GC
		defer func() {
			runtime.GC()
			printMemStats(nil)
		}()
	}
	// Start processing
	q := f.Start()
	if *stopAfter > 0 {
		go func() {
			c := time.After(*stopAfter)
			<-c
			q.Close()
		}()
	}
	// Enqueue the seed, which is the first entry in the dup map
	dup[*seed] = true
	_, err = q.SendStringGet(*seed)
	if err != nil {
		fmt.Printf("[ERR] GET %s - %s\n", *seed, err)
	}
	q.Block()
}

func runMemStats(f *fetchbot.Fetcher, tick time.Duration) {
	var mu sync.Mutex
	var di *fetchbot.DebugInfo

	// Start goroutine to collect fetchbot debug info
	go func() {
		for v := range f.Debug() {
			mu.Lock()
			di = v
			mu.Unlock()
		}
	}()
	// Start ticker goroutine to print mem stats at regular intervals
	go func() {
		c := time.Tick(tick)
		for _ = range c {
			mu.Lock()
			printMemStats(di)
			mu.Unlock()
		}
	}()
}

func printMemStats(di *fetchbot.DebugInfo) {
	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)
	buf := bytes.NewBuffer(nil)
	buf.WriteString(strings.Repeat("=", 72) + "\n")
	buf.WriteString("Memory Profile:\n")
	buf.WriteString(fmt.Sprintf("\tAlloc: %d Kb\n", mem.Alloc/1024))
	buf.WriteString(fmt.Sprintf("\tTotalAlloc: %d Kb\n", mem.TotalAlloc/1024))
	buf.WriteString(fmt.Sprintf("\tNumGC: %d\n", mem.NumGC))
	buf.WriteString(fmt.Sprintf("\tGoroutines: %d\n", runtime.NumGoroutine()))
	if di != nil {
		buf.WriteString(fmt.Sprintf("\tNumHosts: %d\n", di.NumHosts))
	}
	buf.WriteString(strings.Repeat("=", 72))
	fmt.Println(buf.String())
}

// stopHandler stops the fetcher if the stopurl is reached. Otherwise it dispatches
// the call to the wrapped Handler.
func stopHandler(stopurl string, wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if ctx.Cmd.URL().String() == stopurl {
			ctx.Q.Close()
			return
		}
		wrapped.Handle(ctx, res, err)
	})
}

// logHandler prints the fetch information and dispatches the call to the wrapped Handler.
func logHandler(wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			fmt.Printf("[%d] %s %s - %s\n", res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		}
		wrapped.Handle(ctx, res, err)
	})
}

func enqueueLinks(ctx *fetchbot.Context, doc *goquery.Document) {
	mu.Lock()
	doc.Find("a[href]").Each(func(i int, s *goquery.Selection) {
		val, _ := s.Attr("href")
		// Resolve address
		u, err := ctx.Cmd.URL().Parse(val)
		if err != nil {
			fmt.Printf("error: resolve URL %s - %s\n", val, err)
			return
		}
		if !dup[u.String()] {
			if _, err := ctx.Q.SendStringHead(u.String()); err != nil {
				fmt.Printf("error: enqueue head %s - %s\n", u, err)
			} else {
				dup[u.String()] = true
			}
		}
	})
	mu.Unlock()
}
<commit_msg>add cancel behaviour in full example<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

var (
	// Protect access to dup
	mu sync.Mutex
	// Duplicates table
	dup = map[string]bool{}

	// Command-line flags
	seed        = flag.String("seed", "http://golang.org", "seed URL")
	cancelAfter = flag.Duration("cancelafter", 0, "automatically cancel the fetchbot after a given time")
	cancelAtURL = flag.String("cancelat", "", "automatically cancel the fetchbot at a given URL")
	stopAfter   = flag.Duration("stopafter", 0, "automatically stop the fetchbot after a given time")
	stopAtURL   = flag.String("stopat", "", "automatically stop the fetchbot at a given URL")
	memStats    = flag.Duration("memstats", 0, "display memory statistics at a given interval")
)

func main() {
	flag.Parse()

	// Parse the provided seed
	u, err := url.Parse(*seed)
	if err != nil {
		log.Fatal(err)
	}

	// Create the muxer
	mux := fetchbot.NewMux()

	// Handle all errors the same
	mux.HandleErrors(fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	}))

	// Handle GET requests for html responses, to parse the body and enqueue all links as HEAD
	// requests.
	mux.Response().Method("GET").ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			// Process the body to find the links
			doc, err := goquery.NewDocumentFromResponse(res)
			if err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}
			// Enqueue all links as HEAD requests
			enqueueLinks(ctx, doc)
		}))

	// Handle HEAD requests for html responses coming from the source host - we don't want
	// to crawl links from other hosts.
	mux.Response().Method("HEAD").Host(u.Host).ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			if _, err := ctx.Q.SendStringGet(ctx.Cmd.URL().String()); err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			}
		}))

	// Create the Fetcher, handle the logging first, then dispatch to the Muxer
	h := logHandler(mux)
	if *stopAtURL != "" || *cancelAtURL != "" {
		stopURL := *stopAtURL
		if *cancelAtURL != "" {
			stopURL = *cancelAtURL
		}
		h = stopHandler(stopURL, *cancelAtURL != "", logHandler(mux))
	}
	f := fetchbot.New(h)

	// First mem stat print must be right after creating the fetchbot
	if *memStats > 0 {
		// Print starting stats
		printMemStats(nil)
		// Run at regular intervals
		runMemStats(f, *memStats)
		// On exit, print ending stats after a GC
		defer func() {
			runtime.GC()
			printMemStats(nil)
		}()
	}

	// Start processing
	q := f.Start()

	// if a stop or cancel is requested after some duration, launch the goroutine
	// that will stop or cancel.
	if *stopAfter > 0 || *cancelAfter > 0 {
		after := *stopAfter
		stopFunc := q.Close
		if *cancelAfter != 0 {
			after = *cancelAfter
			stopFunc = q.Cancel
		}

		go func() {
			c := time.After(after)
			<-c
			stopFunc()
		}()
	}

	// Enqueue the seed, which is the first entry in the dup map
	dup[*seed] = true
	_, err = q.SendStringGet(*seed)
	if err != nil {
		fmt.Printf("[ERR] GET %s - %s\n", *seed, err)
	}
	q.Block()
}

func runMemStats(f *fetchbot.Fetcher, tick time.Duration) {
	var mu sync.Mutex
	var di *fetchbot.DebugInfo

	// Start goroutine to collect fetchbot debug info
	go func() {
		for v := range f.Debug() {
			mu.Lock()
			di = v
			mu.Unlock()
		}
	}()
	// Start ticker goroutine to print mem stats at regular intervals
	go func() {
		c := time.Tick(tick)
		for _ = range c {
			mu.Lock()
			printMemStats(di)
			mu.Unlock()
		}
	}()
}

func printMemStats(di *fetchbot.DebugInfo) {
	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)
	buf := bytes.NewBuffer(nil)
	buf.WriteString(strings.Repeat("=", 72) + "\n")
	buf.WriteString("Memory Profile:\n")
	buf.WriteString(fmt.Sprintf("\tAlloc: %d Kb\n", mem.Alloc/1024))
	buf.WriteString(fmt.Sprintf("\tTotalAlloc: %d Kb\n", mem.TotalAlloc/1024))
	buf.WriteString(fmt.Sprintf("\tNumGC: %d\n", mem.NumGC))
	buf.WriteString(fmt.Sprintf("\tGoroutines: %d\n", runtime.NumGoroutine()))
	if di != nil {
		buf.WriteString(fmt.Sprintf("\tNumHosts: %d\n", di.NumHosts))
	}
	buf.WriteString(strings.Repeat("=", 72))
	fmt.Println(buf.String())
}

// stopHandler stops the fetcher if the stopurl is reached. Otherwise it dispatches
// the call to the wrapped Handler.
func stopHandler(stopurl string, cancel bool, wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if ctx.Cmd.URL().String() == stopurl {
			fmt.Printf(">>>>> STOP URL %s\n", ctx.Cmd.URL())
			// generally not a good idea to stop/block from a handler goroutine
			// so do it in a separate goroutine
			go func() {
				if cancel {
					ctx.Q.Cancel()
				} else {
					ctx.Q.Close()
				}
			}()
			return
		}
		wrapped.Handle(ctx, res, err)
	})
}

// logHandler prints the fetch information and dispatches the call to the wrapped Handler.
func logHandler(wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			fmt.Printf("[%d] %s %s - %s\n", res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		}
		wrapped.Handle(ctx, res, err)
	})
}

func enqueueLinks(ctx *fetchbot.Context, doc *goquery.Document) {
	mu.Lock()
	doc.Find("a[href]").Each(func(i int, s *goquery.Selection) {
		val, _ := s.Attr("href")
		// Resolve address
		u, err := ctx.Cmd.URL().Parse(val)
		if err != nil {
			fmt.Printf("error: resolve URL %s - %s\n", val, err)
			return
		}
		if !dup[u.String()] {
			if _, err := ctx.Q.SendStringHead(u.String()); err != nil {
				fmt.Printf("error: enqueue head %s - %s\n", u, err)
			} else {
				dup[u.String()] = true
			}
		}
	})
	mu.Unlock()
}
<|endoftext|>


<commit_before>package v2

import (
	"encoding/json"
	"fmt"
	"github.com/laincloud/lainlet/api"
	"github.com/laincloud/lainlet/auth"
	"github.com/laincloud/lainlet/watcher"
	"github.com/laincloud/lainlet/watcher/container"
	"net/http"
)

// node watcher api, /lain/nodes/nodes
type GeneralContainers struct {
	Data map[string]container.Info // data type return by configwatcher
}

func (gc *GeneralContainers) Decode(r []byte) error {
	return json.Unmarshal(r, &gc.Data)
}

func (gc *GeneralContainers) Encode() ([]byte, error) {
	return json.Marshal(gc.Data)
}

func (gc *GeneralContainers) URI() string {
	return "/containers"
}

func (gc *GeneralContainers) WatcherName() string {
	return watcher.CONTAINER
}

func (gc *GeneralContainers) Make(data map[string]interface{}) (api.API, bool, error) {
	ret := &GeneralContainers{
		Data: make(map[string]container.Info),
	}
	for k, v := range data {
		ret.Data[k] = v.(container.Info)
	}
	return ret, true, nil
}

func (gc *GeneralContainers) Key(r *http.Request) (string, error) {
	if !auth.IsSuper(r.RemoteAddr) {
		return "", fmt.Errorf("authorize failed, super required")
	}
	target := api.GetString(r, "nodename", "*")
	if target != "*" {
		target = fixPrefix(target)
	}
	return target, nil
}
<commit_msg>fix interface conversion by nil to container.Info<commit_after>package v2

import (
	"encoding/json"
	"fmt"
	"github.com/laincloud/lainlet/api"
	"github.com/laincloud/lainlet/auth"
	"github.com/laincloud/lainlet/watcher"
	"github.com/laincloud/lainlet/watcher/container"
	"net/http"
)

// node watcher api, /lain/nodes/nodes
type GeneralContainers struct {
	Data map[string]container.Info // data type return by configwatcher
}

func (gc *GeneralContainers) Decode(r []byte) error {
	return json.Unmarshal(r, &gc.Data)
}

func (gc *GeneralContainers) Encode() ([]byte, error) {
	return json.Marshal(gc.Data)
}

func (gc *GeneralContainers) URI() string {
	return "/containers"
}

func (gc *GeneralContainers) WatcherName() string {
	return watcher.CONTAINER
}

func (gc *GeneralContainers) Make(data map[string]interface{}) (api.API, bool, error) {
	ret := &GeneralContainers{
		Data: make(map[string]container.Info),
	}
	for k, v := range data {
		if nil == v {
			continue
		}
		ret.Data[k] = v.(container.Info)
	}
	return ret, true, nil
}

func (gc *GeneralContainers) Key(r *http.Request) (string, error) {
	if !auth.IsSuper(r.RemoteAddr) {
		return "", fmt.Errorf("authorize failed, super required")
	}
	target := api.GetString(r, "nodename", "*")
	if target != "*" {
		target = fixPrefix(target)
	}
	return target, nil
}
<|endoftext|>


<commit_before>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package commands

import (
	"log"
	"strings"

	"github.com/nelsam/gxui"
)

type LicenseHeaderUpdate struct {
}

func NewLicenseHeaderUpdate() *LicenseHeaderUpdate {
	return &LicenseHeaderUpdate{}
}

func (u *LicenseHeaderUpdate) Start(gxui.Control) gxui.Control {
	return nil
}

func (u *LicenseHeaderUpdate) Name() string {
	return "goimports"
}

func (u *LicenseHeaderUpdate) Next() gxui.Focusable {
	return nil
}

func (u *LicenseHeaderUpdate) Exec(on interface{}) (executed, consume bool) {
	finder, ok := on.(ProjectFinder)
	if !ok {
		return false, false
	}
	editor := finder.CurrentEditor()
	if editor == nil {
		return true, true
	}
	licenseHeader := strings.TrimSpace(finder.Project().LicenseHeader()) + "\n\n"

	// License headers should be the first comment block on a file,
	// should have an empty line following them, and should not
	// contain build tag comments
	firstCommentBlockEnd := 0
	for i := 0; i < editor.Controller().LineCount(); i++ {
		line := editor.Controller().LineRunes(i)
		if !strings.HasPrefix(string(line), "//") {
			if len(line) == 0 {
				firstCommentBlockEnd = editor.Controller().LineStart(i + 1)
			}
			break
		}
	}
	runes := editor.Controller().TextRunes()
	commentText := string(runes[:firstCommentBlockEnd])
	if strings.Contains(commentText, "// +build ") {
		commentText = ""
		firstCommentBlockEnd = 0
	}
	if commentText == licenseHeader {
		log.Printf("License header update: license is already set correctly")
		return true, true
	}
	runes = append(runes[:len(licenseHeader)], runes[len(commentText):]...)
	copy(runes, []rune(licenseHeader))

	edits := []gxui.TextBoxEdit{
		{
			At:    0,
			Delta: len(licenseHeader) - len(commentText),
			Old:   []rune(commentText),
			New:   []rune(licenseHeader),
		},
	}
	editor.Controller().SetTextEdits(runes, edits)
	return true, true
}
<commit_msg>Make license header update work better for empty licenses<commit_after>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package commands

import (
	"log"
	"strings"

	"github.com/nelsam/gxui"
)

type LicenseHeaderUpdate struct {
}

func NewLicenseHeaderUpdate() *LicenseHeaderUpdate {
	return &LicenseHeaderUpdate{}
}

func (u *LicenseHeaderUpdate) Start(gxui.Control) gxui.Control {
	return nil
}

func (u *LicenseHeaderUpdate) Name() string {
	return "goimports"
}

func (u *LicenseHeaderUpdate) Next() gxui.Focusable {
	return nil
}

func (u *LicenseHeaderUpdate) Exec(on interface{}) (executed, consume bool) {
	finder, ok := on.(ProjectFinder)
	if !ok {
		return false, false
	}
	editor := finder.CurrentEditor()
	if editor == nil {
		return true, true
	}
	licenseHeader := strings.TrimSpace(finder.Project().LicenseHeader())
	if licenseHeader != "" {
		licenseHeader += "\n\n"
	}

	// License headers should be the first comment block on a file,
	// should have an empty line following them, and should not
	// contain build tag comments
	firstCommentBlockEnd := 0
	for i := 0; i < editor.Controller().LineCount(); i++ {
		line := editor.Controller().LineRunes(i)
		if !strings.HasPrefix(string(line), "//") {
			if len(line) == 0 {
				firstCommentBlockEnd = editor.Controller().LineStart(i + 1)
			}
			break
		}
	}
	runes := editor.Controller().TextRunes()
	commentText := string(runes[:firstCommentBlockEnd])
	if strings.Contains(commentText, "// +build ") {
		commentText = ""
		firstCommentBlockEnd = 0
	}
	if commentText == licenseHeader {
		log.Printf("License header update: license is already set correctly")
		return true, true
	}
	runes = append(runes[:len(licenseHeader)], runes[len(commentText):]...)
	copy(runes, []rune(licenseHeader))

	edits := []gxui.TextBoxEdit{
		{
			At:    0,
			Delta: len(licenseHeader) - len(commentText),
			Old:   []rune(commentText),
			New:   []rune(licenseHeader),
		},
	}
	editor.Controller().SetTextEdits(runes, edits)
	return true, true
}
<|endoftext|>


<commit_before>package kana

import (
	"strings"
	"unicode"
)

var consonants []string = []string{"b", "d", "f", "g", "h", "j", "k", "l", "m", "p", "r", "s", "t", "w", "z"}

var kanaToRomajiTrie *Trie
var romajiToHiraganaTrie *Trie
var romajiToKatakanaTrie *Trie

func Initialize() {
	/*
		Build the Hiragana + Katakana trie.

		Because there is no overlap between the hiragana and katakana sets,
		they both use the same trie without conflict. Nice bonus!
	*/
	kanaToRomajiTrie = newTrie()
	romajiToHiraganaTrie = newTrie()
	romajiToKatakanaTrie = newTrie()

	tables := []string{HiraganaTable, KatakanaTable}
	for t, table := range tables {
		rows := strings.Split(table, "\n")
		colNames := strings.Split(string(rows[0]), "\t")[1:]
		for _, row := range rows[1:] {
			cols := strings.Split(string(row), "\t")
			rowName := cols[0]
			for i, kana := range cols[1:] {
				value := rowName + colNames[i]
				kanas := strings.Split(kana, "/")
				for _, singleKana := range kanas {
					if singleKana != "" {
						// add to tries
						kanaToRomajiTrie.insert(singleKana, value)
						if t == 0 {
							romajiToHiraganaTrie.insert(value, singleKana)
						} else if t == 1 {
							romajiToKatakanaTrie.insert(value, singleKana)
						}
					}
				}
			}
		}
	}
}

func KanaToRomaji(kana string) (romaji string) {
	romaji = kanaToRomajiTrie.convert(kana)

	// do some post-processing for the tsu and stripe characters
	// maybe a bit of a hacky solution - how can we improve?
	// (they act more like punctuation)
	tsus := []string{"っ", "ッ"}
	for _, tsu := range tsus {
		for i := strings.Index(romaji, tsu); i > -1; i = strings.Index(romaji, tsu) {
			rune_romaji := []rune(romaji)
			if len(rune_romaji) > i+2 {
				// TODO: should check if following letter is consonant
				followingLetter := string(rune_romaji[i+1 : i+2])
				romaji = strings.Replace(romaji, tsu, followingLetter, 1)
			} else {
				romaji = strings.Replace(romaji, tsu, "", 1)
			}
		}
	}

	line := "ー"
	for i := strings.Index(romaji, line); i > -1; i = strings.Index(romaji, line) {
		if i > 0 {
			romaji = strings.Replace(romaji, line, "-", 1)
		} else {
			romaji = strings.Replace(romaji, line, "", 1)
		}
	}

	return romaji
}

func replace_tsus(romaji string, tsu string) (result string) {
	result = romaji
	for _, consonant := range consonants {
		result = strings.Replace(result, consonant+consonant, tsu+consonant, -1)
	}
	return result
}

func RomajiToHiragana(romaji string) (hiragana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	romaji = replace_tsus(romaji, "っ")
	hiragana = romajiToHiraganaTrie.convert(romaji)
	return hiragana
}

func RomajiToKatakana(romaji string) (katakana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	// convert double consonants to little tsus first
	romaji = replace_tsus(romaji, "ッ")
	katakana = romajiToKatakanaTrie.convert(romaji)
	return katakana
}

func IsLatin(s string) bool {
	isLatin := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isLatin = isLatin && unicode.IsOneOf([]*unicode.RangeTable{unicode.Latin, unicode.ASCII_Hex_Digit, unicode.White_Space, unicode.Hyphen}, r)
		if !isLatin {
			return isLatin
		}
	}
	return isLatin
}

func IsKana(s string) bool {
	isKana := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKana = isKana && unicode.IsOneOf([]*unicode.RangeTable{unicode.Hiragana, unicode.Katakana, unicode.Hyphen, unicode.Diacritic}, r)
		if !isKana {
			return isKana
		}
	}
	return isKana
}

func IsKanji(s string) bool {
	isKanji := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKanji = isKanji && unicode.IsOneOf([]*unicode.RangeTable{unicode.Ideographic}, r)
		if !isKanji {
			return isKanji
		}
	}
	return isKanji
}

func replaceAll(haystack string, needles []string, replacements []string) (replaced string) {
	replaced = haystack
	for i := range needles {
		replaced = strings.Replace(replaced, needles[i], replacements[i], -1)
	}
	return replaced
}

func NormalizeRomaji(s string) (romaji string) {
	// transform romaji input to one specific standard form,
	// which should be as close as possible to hiragana so that
	// this library gives correct output when transforming to
	// hiragana / katakana

	romaji = s
	romaji = strings.ToLower(romaji)
	romaji = replaceAll(
		romaji,
		[]string{"ā", "ē", "ī", "ō", "ū", "ee", "uu"},
		[]string{"a-", "ei", "ii", "oo", "u-", "ei", "u-"},
	)

	return romaji
}
<commit_msg>Fix bug in replacing tsus<commit_after>package kana

import (
	"strings"
	"unicode"
	"unicode/utf8"
)

var consonants []string = []string{"b", "d", "f", "g", "h", "j", "k", "l", "m", "p", "r", "s", "t", "w", "z"}

var kanaToRomajiTrie *Trie
var romajiToHiraganaTrie *Trie
var romajiToKatakanaTrie *Trie

func Initialize() {
	/*
		Build the Hiragana + Katakana trie.

		Because there is no overlap between the hiragana and katakana sets,
		they both use the same trie without conflict. Nice bonus!
	*/
	kanaToRomajiTrie = newTrie()
	romajiToHiraganaTrie = newTrie()
	romajiToKatakanaTrie = newTrie()

	tables := []string{HiraganaTable, KatakanaTable}
	for t, table := range tables {
		rows := strings.Split(table, "\n")
		colNames := strings.Split(string(rows[0]), "\t")[1:]
		for _, row := range rows[1:] {
			cols := strings.Split(string(row), "\t")
			rowName := cols[0]
			for i, kana := range cols[1:] {
				value := rowName + colNames[i]
				kanas := strings.Split(kana, "/")
				for _, singleKana := range kanas {
					if singleKana != "" {
						// add to tries
						kanaToRomajiTrie.insert(singleKana, value)
						if t == 0 {
							romajiToHiraganaTrie.insert(value, singleKana)
						} else if t == 1 {
							romajiToKatakanaTrie.insert(value, singleKana)
						}
					}
				}
			}
		}
	}
}

func KanaToRomaji(kana string) (romaji string) {
	romaji = kanaToRomajiTrie.convert(kana)

	// do some post-processing for the tsu and stripe characters
	// maybe a bit of a hacky solution - how can we improve?
	// (they act more like punctuation)
	tsus := []string{"っ", "ッ"}
	for _, tsu := range tsus {
		if strings.Index(romaji, tsu) > -1 {
			for _, c := range romaji {
				ch := string(c)
				if ch == tsu {
					i := strings.Index(romaji, ch)
					runeSize := len(ch)
					followingLetter, _ := utf8.DecodeRuneInString(romaji[i+runeSize:])
					followingLetterStr := string(followingLetter)
					if followingLetterStr != tsu {
						romaji = strings.Replace(romaji, tsu, followingLetterStr, 1)
					} else {
						romaji = strings.Replace(romaji, tsu, "", 1)
					}
				}
			}
		}
	}

	line := "ー"
	for i := strings.Index(romaji, line); i > -1; i = strings.Index(romaji, line) {
		if i > 0 {
			romaji = strings.Replace(romaji, line, "-", 1)
		} else {
			romaji = strings.Replace(romaji, line, "", 1)
		}
	}
	return romaji
}

func replace_tsus(romaji string, tsu string) (result string) {
	result = romaji
	for _, consonant := range consonants {
		result = strings.Replace(result, consonant+consonant, tsu+consonant, -1)
	}
	return result
}

func RomajiToHiragana(romaji string) (hiragana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	romaji = replace_tsus(romaji, "っ")
	hiragana = romajiToHiraganaTrie.convert(romaji)
	return hiragana
}

func RomajiToKatakana(romaji string) (katakana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	// convert double consonants to little tsus first
	romaji = replace_tsus(romaji, "ッ")
	katakana = romajiToKatakanaTrie.convert(romaji)
	return katakana
}

func IsLatin(s string) bool {
	isLatin := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isLatin = isLatin && unicode.IsOneOf([]*unicode.RangeTable{unicode.Latin, unicode.ASCII_Hex_Digit, unicode.White_Space, unicode.Hyphen}, r)
		if !isLatin {
			return isLatin
		}
	}
	return isLatin
}

func IsKana(s string) bool {
	isKana := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKana = isKana && unicode.IsOneOf([]*unicode.RangeTable{unicode.Hiragana, unicode.Katakana, unicode.Hyphen, unicode.Diacritic}, r)
		if !isKana {
			return isKana
		}
	}
	return isKana
}

func IsKanji(s string) bool {
	isKanji := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKanji = isKanji && unicode.IsOneOf([]*unicode.RangeTable{unicode.Ideographic}, r)
		if !isKanji {
			return isKanji
		}
	}
	return isKanji
}

func replaceAll(haystack string, needles []string, replacements []string) (replaced string) {
	replaced = haystack
	for i := range needles {
		replaced = strings.Replace(replaced, needles[i], replacements[i], -1)
	}
	return replaced
}

func NormalizeRomaji(s string) (romaji string) {
	// transform romaji input to one specific standard form,
	// which should be as close as possible to hiragana so that
	// this library gives correct output when transforming to
	// hiragana / katakana

	romaji = s
	romaji = strings.ToLower(romaji)
	romaji = replaceAll(
		romaji,
		[]string{"ā", "ē", "ī", "ō", "ū", "ee", "uu"},
		[]string{"a-", "ei", "ii", "oo", "u-", "ei", "u-"},
	)

	return romaji
}
<|endoftext|>


<commit_before>package kana

import (
	"strings"
	"unicode"
)

type Kana struct {
	kanaToRomajiTrie     *Trie
	romajiToHiraganaTrie *Trie
	romajiToKatakanaTrie *Trie
}

var consonants []string = []string{"b", "d", "f", "g", "h", "j", "k", "l", "m", "p", "r", "s", "t", "w", "z"}

func NewKana() *Kana {
	/*
	   Build a trie for efficient retrieval of entries
	*/
	kana := &Kana{newTrie(), newTrie(), newTrie()}
	kana.initialize()
	return kana
}

func (k *Kana) initialize() {
	/*
		Build the Hiragana + Katakana trie.

		Because there is no overlap between the hiragana and katakana sets,
		they both use the same trie without conflict. Nice bonus!
	*/
	tables := []string{HiraganaTable, KatakanaTable}
	for t, table := range tables {
		rows := strings.Split(table, "\n")
		colNames := strings.Split(string(rows[0]), "\t")[1:]
		for _, row := range rows[1:] {
			cols := strings.Split(string(row), "\t")
			rowName := cols[0]
			for i, kana := range cols[1:] {
				value := rowName + colNames[i]
				kanas := strings.Split(kana, "/")
				for _, singleKana := range kanas {
					if singleKana != "" {
						// add to tries
						k.kanaToRomajiTrie.insert(singleKana, value)
						if t == 0 {
							k.romajiToHiraganaTrie.insert(value, singleKana)
						} else if t == 1 {
							k.romajiToKatakanaTrie.insert(value, singleKana)
						}
					}
				}
			}
		}
	}
}

func (k Kana) KanaToRomaji(kana string) (romaji string) {
	romaji = k.kanaToRomajiTrie.convert(kana)

	// do some post-processing for the tsu and stripe characters
	// maybe a bit of a hacky solution - how can we improve?
	// (they act more like punctuation)
	tsus := []string{"っ", "ッ"}
	for _, tsu := range tsus {
		for i := strings.Index(romaji, tsu); i > -1; i = strings.Index(romaji, tsu) {
			rune_romaji := []rune(romaji)
			if len(rune_romaji) > i+2 {
				// TODO: should check if following letter is consonant
				followingLetter := string(rune_romaji[i+1 : i+2])
				romaji = strings.Replace(romaji, tsu, followingLetter, 1)
			} else {
				romaji = strings.Replace(romaji, tsu, "", 1)
			}
		}
	}

	line := "ー"
	for i := strings.Index(romaji, line); i > -1; i = strings.Index(romaji, line) {
		rune_romaji := []rune(romaji)
		if i > 0 {
			// TODO: should check if following letter is consonant
			previousLetter := string(rune_romaji[i-1 : i])
			romaji = strings.Replace(romaji, line, previousLetter, 1)
		} else {
			romaji = strings.Replace(romaji, line, "", 1)
		}
	}

	return romaji
}

func replace_tsus(romaji string, tsu string) (result string) {
	result = romaji
	for _, consonant := range consonants {
		result = strings.Replace(result, consonant+consonant, tsu+consonant, -1)
	}
	return result
}

func (k Kana) RomajiToHiragana(romaji string) (hiragana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	romaji = replace_tsus(romaji, "っ")
	hiragana = k.romajiToHiraganaTrie.convert(romaji)
	return hiragana
}

func (k Kana) RomajiToKatakana(romaji string) (katakana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	// convert double consonants to little tsus first
	romaji = replace_tsus(romaji, "ッ")
	katakana = k.romajiToKatakanaTrie.convert(romaji)
	return katakana
}

func (k Kana) IsLatin(s string) bool {
	isLatin := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isLatin = isLatin && unicode.IsOneOf([]*unicode.RangeTable{unicode.Latin, unicode.ASCII_Hex_Digit, unicode.White_Space, unicode.Hyphen}, r)
		if !isLatin {
			return isLatin
		}
	}
	return isLatin
}

func (k Kana) IsKana(s string) bool {
	isKana := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKana = isKana && unicode.IsOneOf([]*unicode.RangeTable{unicode.Hiragana, unicode.Katakana, unicode.Hyphen, unicode.Diacritic}, r)
		if !isKana {
			return isKana
		}
	}
	return isKana
}

func (k Kana) IsKanji(s string) bool {
	isKanji := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKanji = isKanji && unicode.IsOneOf([]*unicode.RangeTable{unicode.Ideographic}, r)
		if !isKanji {
			return isKanji
		}
	}
	return isKanji
}
<commit_msg>Add NormalizeRomaji function<commit_after>package kana

import (
	"strings"
	"unicode"
)

type Kana struct {
	kanaToRomajiTrie     *Trie
	romajiToHiraganaTrie *Trie
	romajiToKatakanaTrie *Trie
}

var consonants []string = []string{"b", "d", "f", "g", "h", "j", "k", "l", "m", "p", "r", "s", "t", "w", "z"}

func NewKana() *Kana {
	/*
	   Build a trie for efficient retrieval of entries
	*/
	kana := &Kana{newTrie(), newTrie(), newTrie()}
	kana.initialize()
	return kana
}

func (k *Kana) initialize() {
	/*
		Build the Hiragana + Katakana trie.

		Because there is no overlap between the hiragana and katakana sets,
		they both use the same trie without conflict. Nice bonus!
	*/
	tables := []string{HiraganaTable, KatakanaTable}
	for t, table := range tables {
		rows := strings.Split(table, "\n")
		colNames := strings.Split(string(rows[0]), "\t")[1:]
		for _, row := range rows[1:] {
			cols := strings.Split(string(row), "\t")
			rowName := cols[0]
			for i, kana := range cols[1:] {
				value := rowName + colNames[i]
				kanas := strings.Split(kana, "/")
				for _, singleKana := range kanas {
					if singleKana != "" {
						// add to tries
						k.kanaToRomajiTrie.insert(singleKana, value)
						if t == 0 {
							k.romajiToHiraganaTrie.insert(value, singleKana)
						} else if t == 1 {
							k.romajiToKatakanaTrie.insert(value, singleKana)
						}
					}
				}
			}
		}
	}
}

func (k Kana) KanaToRomaji(kana string) (romaji string) {
	romaji = k.kanaToRomajiTrie.convert(kana)

	// do some post-processing for the tsu and stripe characters
	// maybe a bit of a hacky solution - how can we improve?
	// (they act more like punctuation)
	tsus := []string{"っ", "ッ"}
	for _, tsu := range tsus {
		for i := strings.Index(romaji, tsu); i > -1; i = strings.Index(romaji, tsu) {
			rune_romaji := []rune(romaji)
			if len(rune_romaji) > i+2 {
				// TODO: should check if following letter is consonant
				followingLetter := string(rune_romaji[i+1 : i+2])
				romaji = strings.Replace(romaji, tsu, followingLetter, 1)
			} else {
				romaji = strings.Replace(romaji, tsu, "", 1)
			}
		}
	}

	line := "ー"
	for i := strings.Index(romaji, line); i > -1; i = strings.Index(romaji, line) {
		rune_romaji := []rune(romaji)
		if i > 0 {
			// TODO: should check if following letter is consonant
			previousLetter := string(rune_romaji[i-1 : i])
			romaji = strings.Replace(romaji, line, previousLetter, 1)
		} else {
			romaji = strings.Replace(romaji, line, "", 1)
		}
	}

	return romaji
}

func replace_tsus(romaji string, tsu string) (result string) {
	result = romaji
	for _, consonant := range consonants {
		result = strings.Replace(result, consonant+consonant, tsu+consonant, -1)
	}
	return result
}

func (k Kana) RomajiToHiragana(romaji string) (hiragana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	romaji = replace_tsus(romaji, "っ")
	hiragana = k.romajiToHiraganaTrie.convert(romaji)
	return hiragana
}

func (k Kana) RomajiToKatakana(romaji string) (katakana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	// convert double consonants to little tsus first
	romaji = replace_tsus(romaji, "ッ")
	katakana = k.romajiToKatakanaTrie.convert(romaji)
	return katakana
}

func (k Kana) IsLatin(s string) bool {
	isLatin := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isLatin = isLatin && unicode.IsOneOf([]*unicode.RangeTable{unicode.Latin, unicode.ASCII_Hex_Digit, unicode.White_Space, unicode.Hyphen}, r)
		if !isLatin {
			return isLatin
		}
	}
	return isLatin
}

func (k Kana) IsKana(s string) bool {
	isKana := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKana = isKana && unicode.IsOneOf([]*unicode.RangeTable{unicode.Hiragana, unicode.Katakana, unicode.Hyphen, unicode.Diacritic}, r)
		if !isKana {
			return isKana
		}
	}
	return isKana
}

func (k Kana) IsKanji(s string) bool {
	isKanji := true
	runeForm := []rune(s)
	for _, r := range runeForm {
		isKanji = isKanji && unicode.IsOneOf([]*unicode.RangeTable{unicode.Ideographic}, r)
		if !isKanji {
			return isKanji
		}
	}
	return isKanji
}

func replaceAll(haystack string, needles []string, replacements []string) (replaced string) {
	replaced = haystack
	for i := range needles {
		replaced = strings.Replace(replaced, needles[i], replacements[i], -1)
	}
	return replaced
}

func (k Kana) NormalizeRomaji(s string) (romaji string) {
	// transform romaji input to one specific standard form,
	// which should be as close as possible to hiragana so that
	// this library gives correct output when transforming to
	// hiragana / katakana

	romaji = s
	romaji = strings.ToLower(romaji)
	romaji = replaceAll(
		romaji,
		[]string{"ā", "ē", "ī", "ō", "ū", "ee", "uu"},
		[]string{"a-", "ei", "ii", "oo", "u-", "ei", "u-"},
	)

	return romaji
}
<|endoftext|>