<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 resourceAwsSubnet() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSubnetCreate,
		Read:   resourceAwsSubnetRead,
		Update: resourceAwsSubnetUpdate,
		Delete: resourceAwsSubnetDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		SchemaVersion: 1,
		MigrateState:  resourceAwsSubnetMigrateState,

		Schema: map[string]*schema.Schema{
			"vpc_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"cidr_block": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ipv6_cidr_block": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"availability_zone": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"map_public_ip_on_launch": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},

			"assign_ipv6_address_on_creation": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},

			"ipv6_cidr_block_association_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

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

	createOpts := &ec2.CreateSubnetInput{
		AvailabilityZone: aws.String(d.Get("availability_zone").(string)),
		CidrBlock:        aws.String(d.Get("cidr_block").(string)),
		VpcId:            aws.String(d.Get("vpc_id").(string)),
	}

	if v, ok := d.GetOk("ipv6_cidr_block"); ok {
		createOpts.Ipv6CidrBlock = aws.String(v.(string))
	}

	var err error
	resp, err := conn.CreateSubnet(createOpts)

	if err != nil {
		return fmt.Errorf("Error creating subnet: %s", err)
	}

	// Get the ID and store it
	subnet := resp.Subnet
	d.SetId(*subnet.SubnetId)
	log.Printf("[INFO] Subnet ID: %s", *subnet.SubnetId)

	// Wait for the Subnet to become available
	log.Printf("[DEBUG] Waiting for subnet (%s) to become available", *subnet.SubnetId)
	stateConf := &resource.StateChangeConf{
		Pending: []string{"pending"},
		Target:  []string{"available"},
		Refresh: SubnetStateRefreshFunc(conn, *subnet.SubnetId),
		Timeout: 10 * time.Minute,
	}

	_, err = stateConf.WaitForState()

	if err != nil {
		return fmt.Errorf(
			"Error waiting for subnet (%s) to become ready: %s",
			d.Id(), err)
	}

	return resourceAwsSubnetUpdate(d, meta)
}

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

	resp, err := conn.DescribeSubnets(&ec2.DescribeSubnetsInput{
		SubnetIds: []*string{aws.String(d.Id())},
	})

	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSubnetID.NotFound" {
			// Update state to indicate the subnet no longer exists.
			d.SetId("")
			return nil
		}
		return err
	}
	if resp == nil {
		return nil
	}

	subnet := resp.Subnets[0]

	d.Set("vpc_id", subnet.VpcId)
	d.Set("availability_zone", subnet.AvailabilityZone)
	d.Set("cidr_block", subnet.CidrBlock)
	d.Set("map_public_ip_on_launch", subnet.MapPublicIpOnLaunch)
	d.Set("assign_ipv6_address_on_creation", subnet.AssignIpv6AddressOnCreation)
	for _, a := range subnet.Ipv6CidrBlockAssociationSet {
		if *a.Ipv6CidrBlockState.State == "associated" { //we can only ever have 1 IPv6 block associated at once
			d.Set("ipv6_cidr_block_association_id", a.AssociationId)
			d.Set("ipv6_cidr_block", a.Ipv6CidrBlock)
			break
		} else {
			d.Set("ipv6_cidr_block_association_id", "") // we blank these out to remove old entries
			d.Set("ipv6_cidr_block", "")
		}
	}
	d.Set("tags", tagsToMap(subnet.Tags))

	return nil
}

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

	d.Partial(true)

	if err := setTags(conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	if d.HasChange("assign_ipv6_address_on_creation") {
		modifyOpts := &ec2.ModifySubnetAttributeInput{
			SubnetId: aws.String(d.Id()),
			AssignIpv6AddressOnCreation: &ec2.AttributeBooleanValue{
				Value: aws.Bool(d.Get("assign_ipv6_address_on_creation").(bool)),
			},
		}

		log.Printf("[DEBUG] Subnet modify attributes: %#v", modifyOpts)

		_, err := conn.ModifySubnetAttribute(modifyOpts)

		if err != nil {
			return err
		} else {
			d.SetPartial("assign_ipv6_address_on_creation")
		}
	}

	if d.HasChange("map_public_ip_on_launch") {
		modifyOpts := &ec2.ModifySubnetAttributeInput{
			SubnetId: aws.String(d.Id()),
			MapPublicIpOnLaunch: &ec2.AttributeBooleanValue{
				Value: aws.Bool(d.Get("map_public_ip_on_launch").(bool)),
			},
		}

		log.Printf("[DEBUG] Subnet modify attributes: %#v", modifyOpts)

		_, err := conn.ModifySubnetAttribute(modifyOpts)

		if err != nil {
			return err
		} else {
			d.SetPartial("map_public_ip_on_launch")
		}
	}

	// We have to be careful here to not go through a change of association if this is a new resource
	// A New resource here would denote that the Update func is called by the Create func
	if d.HasChange("ipv6_cidr_block") && !d.IsNewResource() {
		// We need to handle that we disassociate the IPv6 CIDR block before we try and associate the new one
		// This could be an issue as, we could error out when we try and add the new one
		// We may need to roll back the state and reattach the old one if this is the case

		_, new := d.GetChange("ipv6_cidr_block")

		//Firstly we have to disassociate the old IPv6 CIDR Block
		disassociateOps := &ec2.DisassociateSubnetCidrBlockInput{
			AssociationId: aws.String(d.Get("ipv6_cidr_block_association_id").(string)),
		}

		_, err := conn.DisassociateSubnetCidrBlock(disassociateOps)
		if err != nil {
			return err
		}

		// Wait for the CIDR to become disassociated
		log.Printf(
			"[DEBUG] Waiting for IPv6 CIDR (%s) to become disassociated",
			d.Id())
		stateConf := &resource.StateChangeConf{
			Pending: []string{"disassociating", "associated"},
			Target:  []string{"disassociated"},
			Refresh: SubnetIpv6CidrStateRefreshFunc(conn, d.Id(), d.Get("ipv6_cidr_block_association_id").(string)),
			Timeout: 3 * time.Minute,
		}
		if _, err := stateConf.WaitForState(); err != nil {
			return fmt.Errorf(
				"Error waiting for IPv6 CIDR (%s) to become disassociated: %s",
				d.Id(), err)
		}

		//Now we need to try and associate the new CIDR block
		associatesOpts := &ec2.AssociateSubnetCidrBlockInput{
			SubnetId:      aws.String(d.Id()),
			Ipv6CidrBlock: aws.String(new.(string)),
		}

		resp, err := conn.AssociateSubnetCidrBlock(associatesOpts)
		if err != nil {
			//The big question here is, do we want to try and reassociate the old one??
			//If we have a failure here, then we may be in a situation that we have nothing associated
			return err
		}

		// Wait for the CIDR to become associated
		log.Printf(
			"[DEBUG] Waiting for IPv6 CIDR (%s) to become associated",
			d.Id())
		stateConf = &resource.StateChangeConf{
			Pending: []string{"associating", "disassociated"},
			Target:  []string{"associated"},
			Refresh: SubnetIpv6CidrStateRefreshFunc(conn, d.Id(), *resp.Ipv6CidrBlockAssociation.AssociationId),
			Timeout: 3 * time.Minute,
		}
		if _, err := stateConf.WaitForState(); err != nil {
			return fmt.Errorf(
				"Error waiting for IPv6 CIDR (%s) to become associated: %s",
				d.Id(), err)
		}

		d.SetPartial("ipv6_cidr_block")
	}

	d.Partial(false)

	return resourceAwsSubnetRead(d, meta)
}

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

	log.Printf("[INFO] Deleting subnet: %s", d.Id())
	req := &ec2.DeleteSubnetInput{
		SubnetId: aws.String(d.Id()),
	}

	wait := resource.StateChangeConf{
		Pending:    []string{"pending"},
		Target:     []string{"destroyed"},
		Timeout:    10 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			_, err := conn.DeleteSubnet(req)
			if err != nil {
				if apiErr, ok := err.(awserr.Error); ok {
					if apiErr.Code() == "DependencyViolation" {
						// There is some pending operation, so just retry
						// in a bit.
						return 42, "pending", nil
					}

					if apiErr.Code() == "InvalidSubnetID.NotFound" {
						return 42, "destroyed", nil
					}
				}

				return 42, "failure", err
			}

			return 42, "destroyed", nil
		},
	}

	if _, err := wait.WaitForState(); err != nil {
		return fmt.Errorf("Error deleting subnet: %s", err)
	}

	return nil
}

// SubnetStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch a Subnet.
func SubnetStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeSubnets(&ec2.DescribeSubnetsInput{
			SubnetIds: []*string{aws.String(id)},
		})
		if err != nil {
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSubnetID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on SubnetStateRefresh: %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
		}

		subnet := resp.Subnets[0]
		return subnet, *subnet.State, nil
	}
}

func SubnetIpv6CidrStateRefreshFunc(conn *ec2.EC2, id string, associationId string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		opts := &ec2.DescribeSubnetsInput{
			SubnetIds: []*string{aws.String(id)},
		}
		resp, err := conn.DescribeSubnets(opts)
		if err != nil {
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSubnetID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on SubnetIpv6CidrStateRefreshFunc: %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
		}

		if resp.Subnets[0].Ipv6CidrBlockAssociationSet == nil {
			return nil, "", nil
		}

		for _, association := range resp.Subnets[0].Ipv6CidrBlockAssociationSet {
			if *association.AssociationId == associationId {
				return association, *association.Ipv6CidrBlockState.State, nil
			}
		}

		return nil, "", nil
	}
}
<commit_msg>provider/aws: Set aws_subnet ipv6_cidr_block to computed (#14542)<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 resourceAwsSubnet() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSubnetCreate,
		Read:   resourceAwsSubnetRead,
		Update: resourceAwsSubnetUpdate,
		Delete: resourceAwsSubnetDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		SchemaVersion: 1,
		MigrateState:  resourceAwsSubnetMigrateState,

		Schema: map[string]*schema.Schema{
			"vpc_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"cidr_block": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ipv6_cidr_block": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"availability_zone": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"map_public_ip_on_launch": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},

			"assign_ipv6_address_on_creation": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},

			"ipv6_cidr_block_association_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

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

	createOpts := &ec2.CreateSubnetInput{
		AvailabilityZone: aws.String(d.Get("availability_zone").(string)),
		CidrBlock:        aws.String(d.Get("cidr_block").(string)),
		VpcId:            aws.String(d.Get("vpc_id").(string)),
	}

	if v, ok := d.GetOk("ipv6_cidr_block"); ok {
		createOpts.Ipv6CidrBlock = aws.String(v.(string))
	}

	var err error
	resp, err := conn.CreateSubnet(createOpts)

	if err != nil {
		return fmt.Errorf("Error creating subnet: %s", err)
	}

	// Get the ID and store it
	subnet := resp.Subnet
	d.SetId(*subnet.SubnetId)
	log.Printf("[INFO] Subnet ID: %s", *subnet.SubnetId)

	// Wait for the Subnet to become available
	log.Printf("[DEBUG] Waiting for subnet (%s) to become available", *subnet.SubnetId)
	stateConf := &resource.StateChangeConf{
		Pending: []string{"pending"},
		Target:  []string{"available"},
		Refresh: SubnetStateRefreshFunc(conn, *subnet.SubnetId),
		Timeout: 10 * time.Minute,
	}

	_, err = stateConf.WaitForState()

	if err != nil {
		return fmt.Errorf(
			"Error waiting for subnet (%s) to become ready: %s",
			d.Id(), err)
	}

	return resourceAwsSubnetUpdate(d, meta)
}

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

	resp, err := conn.DescribeSubnets(&ec2.DescribeSubnetsInput{
		SubnetIds: []*string{aws.String(d.Id())},
	})

	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSubnetID.NotFound" {
			// Update state to indicate the subnet no longer exists.
			d.SetId("")
			return nil
		}
		return err
	}
	if resp == nil {
		return nil
	}

	subnet := resp.Subnets[0]

	d.Set("vpc_id", subnet.VpcId)
	d.Set("availability_zone", subnet.AvailabilityZone)
	d.Set("cidr_block", subnet.CidrBlock)
	d.Set("map_public_ip_on_launch", subnet.MapPublicIpOnLaunch)
	d.Set("assign_ipv6_address_on_creation", subnet.AssignIpv6AddressOnCreation)
	for _, a := range subnet.Ipv6CidrBlockAssociationSet {
		if *a.Ipv6CidrBlockState.State == "associated" { //we can only ever have 1 IPv6 block associated at once
			d.Set("ipv6_cidr_block_association_id", a.AssociationId)
			d.Set("ipv6_cidr_block", a.Ipv6CidrBlock)
			break
		} else {
			d.Set("ipv6_cidr_block_association_id", "") // we blank these out to remove old entries
			d.Set("ipv6_cidr_block", "")
		}
	}
	d.Set("tags", tagsToMap(subnet.Tags))

	return nil
}

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

	d.Partial(true)

	if err := setTags(conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	if d.HasChange("assign_ipv6_address_on_creation") {
		modifyOpts := &ec2.ModifySubnetAttributeInput{
			SubnetId: aws.String(d.Id()),
			AssignIpv6AddressOnCreation: &ec2.AttributeBooleanValue{
				Value: aws.Bool(d.Get("assign_ipv6_address_on_creation").(bool)),
			},
		}

		log.Printf("[DEBUG] Subnet modify attributes: %#v", modifyOpts)

		_, err := conn.ModifySubnetAttribute(modifyOpts)

		if err != nil {
			return err
		} else {
			d.SetPartial("assign_ipv6_address_on_creation")
		}
	}

	if d.HasChange("map_public_ip_on_launch") {
		modifyOpts := &ec2.ModifySubnetAttributeInput{
			SubnetId: aws.String(d.Id()),
			MapPublicIpOnLaunch: &ec2.AttributeBooleanValue{
				Value: aws.Bool(d.Get("map_public_ip_on_launch").(bool)),
			},
		}

		log.Printf("[DEBUG] Subnet modify attributes: %#v", modifyOpts)

		_, err := conn.ModifySubnetAttribute(modifyOpts)

		if err != nil {
			return err
		} else {
			d.SetPartial("map_public_ip_on_launch")
		}
	}

	// We have to be careful here to not go through a change of association if this is a new resource
	// A New resource here would denote that the Update func is called by the Create func
	if d.HasChange("ipv6_cidr_block") && !d.IsNewResource() {
		// We need to handle that we disassociate the IPv6 CIDR block before we try and associate the new one
		// This could be an issue as, we could error out when we try and add the new one
		// We may need to roll back the state and reattach the old one if this is the case

		_, new := d.GetChange("ipv6_cidr_block")

		//Firstly we have to disassociate the old IPv6 CIDR Block
		disassociateOps := &ec2.DisassociateSubnetCidrBlockInput{
			AssociationId: aws.String(d.Get("ipv6_cidr_block_association_id").(string)),
		}

		_, err := conn.DisassociateSubnetCidrBlock(disassociateOps)
		if err != nil {
			return err
		}

		// Wait for the CIDR to become disassociated
		log.Printf(
			"[DEBUG] Waiting for IPv6 CIDR (%s) to become disassociated",
			d.Id())
		stateConf := &resource.StateChangeConf{
			Pending: []string{"disassociating", "associated"},
			Target:  []string{"disassociated"},
			Refresh: SubnetIpv6CidrStateRefreshFunc(conn, d.Id(), d.Get("ipv6_cidr_block_association_id").(string)),
			Timeout: 3 * time.Minute,
		}
		if _, err := stateConf.WaitForState(); err != nil {
			return fmt.Errorf(
				"Error waiting for IPv6 CIDR (%s) to become disassociated: %s",
				d.Id(), err)
		}

		//Now we need to try and associate the new CIDR block
		associatesOpts := &ec2.AssociateSubnetCidrBlockInput{
			SubnetId:      aws.String(d.Id()),
			Ipv6CidrBlock: aws.String(new.(string)),
		}

		resp, err := conn.AssociateSubnetCidrBlock(associatesOpts)
		if err != nil {
			//The big question here is, do we want to try and reassociate the old one??
			//If we have a failure here, then we may be in a situation that we have nothing associated
			return err
		}

		// Wait for the CIDR to become associated
		log.Printf(
			"[DEBUG] Waiting for IPv6 CIDR (%s) to become associated",
			d.Id())
		stateConf = &resource.StateChangeConf{
			Pending: []string{"associating", "disassociated"},
			Target:  []string{"associated"},
			Refresh: SubnetIpv6CidrStateRefreshFunc(conn, d.Id(), *resp.Ipv6CidrBlockAssociation.AssociationId),
			Timeout: 3 * time.Minute,
		}
		if _, err := stateConf.WaitForState(); err != nil {
			return fmt.Errorf(
				"Error waiting for IPv6 CIDR (%s) to become associated: %s",
				d.Id(), err)
		}

		d.SetPartial("ipv6_cidr_block")
	}

	d.Partial(false)

	return resourceAwsSubnetRead(d, meta)
}

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

	log.Printf("[INFO] Deleting subnet: %s", d.Id())
	req := &ec2.DeleteSubnetInput{
		SubnetId: aws.String(d.Id()),
	}

	wait := resource.StateChangeConf{
		Pending:    []string{"pending"},
		Target:     []string{"destroyed"},
		Timeout:    10 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			_, err := conn.DeleteSubnet(req)
			if err != nil {
				if apiErr, ok := err.(awserr.Error); ok {
					if apiErr.Code() == "DependencyViolation" {
						// There is some pending operation, so just retry
						// in a bit.
						return 42, "pending", nil
					}

					if apiErr.Code() == "InvalidSubnetID.NotFound" {
						return 42, "destroyed", nil
					}
				}

				return 42, "failure", err
			}

			return 42, "destroyed", nil
		},
	}

	if _, err := wait.WaitForState(); err != nil {
		return fmt.Errorf("Error deleting subnet: %s", err)
	}

	return nil
}

// SubnetStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch a Subnet.
func SubnetStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeSubnets(&ec2.DescribeSubnetsInput{
			SubnetIds: []*string{aws.String(id)},
		})
		if err != nil {
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSubnetID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on SubnetStateRefresh: %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
		}

		subnet := resp.Subnets[0]
		return subnet, *subnet.State, nil
	}
}

func SubnetIpv6CidrStateRefreshFunc(conn *ec2.EC2, id string, associationId string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		opts := &ec2.DescribeSubnetsInput{
			SubnetIds: []*string{aws.String(id)},
		}
		resp, err := conn.DescribeSubnets(opts)
		if err != nil {
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSubnetID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on SubnetIpv6CidrStateRefreshFunc: %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
		}

		if resp.Subnets[0].Ipv6CidrBlockAssociationSet == nil {
			return nil, "", nil
		}

		for _, association := range resp.Subnets[0].Ipv6CidrBlockAssociationSet {
			if *association.AssociationId == associationId {
				return association, *association.Ipv6CidrBlockState.State, nil
			}
		}

		return nil, "", nil
	}
}
<|endoftext|>


<commit_before>// Package vtctld contains all the code to expose a vtctld server
// based on the provided topo.Server.
package vtctld

import (
	"flag"
	"net/http"
	"os"
	"path"
	"strings"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/acl"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

var (
	webDir = flag.String("web_dir", "", "directory from which to serve vtctld web interface resources")
	// webDir2 is a temporary additional dir for a new, in-development UI.
	webDir2             = flag.String("web_dir2", "", "directory from which to serve vtctld2 web interface resources")
	enableRealtimeStats = flag.Bool("enable_realtime_stats", false, "Required for the Realtime Stats view. If set, vtctld will maintain a streaming RPC to each tablet (in all cells) to gather the realtime health stats.")
)

const (
	appPrefix = "/app/"
	// appPrefix2 is a temporary additional path for a new, in-development UI.
	appPrefix2 = "/app2/"
)

// InitVtctld initializes all the vtctld functionnality.
func InitVtctld(ts topo.Server) {
	actionRepo := NewActionRepository(ts)

	// keyspace actions
	actionRepo.RegisterKeyspaceAction("ValidateKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidateKeyspace(ctx, keyspace, false)
		})

	actionRepo.RegisterKeyspaceAction("ValidateSchemaKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidateSchemaKeyspace(ctx, keyspace, nil, false)
		})

	actionRepo.RegisterKeyspaceAction("ValidateVersionKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidateVersionKeyspace(ctx, keyspace)
		})

	actionRepo.RegisterKeyspaceAction("ValidatePermissionsKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidatePermissionsKeyspace(ctx, keyspace)
		})

	// shard actions
	actionRepo.RegisterShardAction("ValidateShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidateShard(ctx, keyspace, shard, false)
		})

	actionRepo.RegisterShardAction("ValidateSchemaShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidateSchemaShard(ctx, keyspace, shard, nil, false)
		})

	actionRepo.RegisterShardAction("ValidateVersionShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidateVersionShard(ctx, keyspace, shard)
		})

	actionRepo.RegisterShardAction("ValidatePermissionsShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidatePermissionsShard(ctx, keyspace, shard)
		})

	// tablet actions
	actionRepo.RegisterTabletAction("Ping", "",
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			ti, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
			if err != nil {
				return "", err
			}
			return "", wr.TabletManagerClient().Ping(ctx, ti.Tablet)
		})

	actionRepo.RegisterTabletAction("RefreshState", acl.ADMIN,
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			ti, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
			if err != nil {
				return "", err
			}
			return "", wr.TabletManagerClient().RefreshState(ctx, ti.Tablet)
		})

	actionRepo.RegisterTabletAction("DeleteTablet", acl.ADMIN,
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			return "", wr.DeleteTablet(ctx, tabletAlias, false)
		})

	actionRepo.RegisterTabletAction("ReloadSchema", acl.ADMIN,
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			return "", wr.ReloadSchema(ctx, tabletAlias)
		})

	// Anything unrecognized gets redirected to the main app page.
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, appPrefix, http.StatusFound)
	})

	// Serve the static files for the vtctld web app.
	http.HandleFunc(appPrefix, func(w http.ResponseWriter, r *http.Request) {
		// Strip the prefix.
		parts := strings.SplitN(r.URL.Path, "/", 3)
		if len(parts) != 3 {
			http.NotFound(w, r)
			return
		}
		rest := parts[2]
		if rest == "" {
			rest = "index.html"
		}
		http.ServeFile(w, r, path.Join(*webDir, rest))
	})

	// Serve the static files for the vtctld2 web app.
	// This is a temporary additional URL for serving the new,
	// in-development UI side-by-side with the current one.
	http.HandleFunc(appPrefix2, func(w http.ResponseWriter, r *http.Request) {
		// Strip the prefix.
		parts := strings.SplitN(r.URL.Path, "/", 3)
		if len(parts) != 3 {
			http.NotFound(w, r)
			return
		}
		rest := parts[2]
		if rest == "" {
			rest = "index.html"
		}
		filePath := path.Join(*webDir2, rest)
		if _, err := os.Stat(filePath); err != nil {
			// The requested file doesn't exist.
			if strings.ContainsAny(rest, "/.") {
				// This looks like a real file path, so return Not Found.
				http.NotFound(w, r)
				return
			}
			// It looks like a virtual route path (for pages within the app).
			// For these, we must serve index.html to initialize the app.
			filePath = path.Join(*webDir2, "index.html")
		}
		http.ServeFile(w, r, filePath)
	})

	var realtimeStats *realtimeStats
	if *enableRealtimeStats {
		var err error
		realtimeStats, err = newRealtimeStats(ts)
		if err != nil {
			log.Errorf("Failed to instantiate RealtimeStats at startup: %v", err)
		}
	}

	// Serve the REST API for the vtctld web app.
	initAPI(context.Background(), ts, actionRepo, realtimeStats)

	// Init redirects for explorers
	initExplorer(ts)

	// Init workflow manager.
	initWorkflowManager(ts)
}
<commit_msg>Switching vtctld default UI to app2.<commit_after>// Package vtctld contains all the code to expose a vtctld server
// based on the provided topo.Server.
package vtctld

import (
	"flag"
	"net/http"
	"os"
	"path"
	"strings"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/acl"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

var (
	webDir = flag.String("web_dir", "", "directory from which to serve vtctld web interface resources")
	// webDir2 is a temporary additional dir for a new, in-development UI.
	webDir2             = flag.String("web_dir2", "", "directory from which to serve vtctld2 web interface resources")
	enableRealtimeStats = flag.Bool("enable_realtime_stats", false, "Required for the Realtime Stats view. If set, vtctld will maintain a streaming RPC to each tablet (in all cells) to gather the realtime health stats.")
)

const (
	appPrefix = "/app/"
	// appPrefix2 is a temporary additional path for a new, in-development UI.
	appPrefix2 = "/app2/"
)

// InitVtctld initializes all the vtctld functionnality.
func InitVtctld(ts topo.Server) {
	actionRepo := NewActionRepository(ts)

	// keyspace actions
	actionRepo.RegisterKeyspaceAction("ValidateKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidateKeyspace(ctx, keyspace, false)
		})

	actionRepo.RegisterKeyspaceAction("ValidateSchemaKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidateSchemaKeyspace(ctx, keyspace, nil, false)
		})

	actionRepo.RegisterKeyspaceAction("ValidateVersionKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidateVersionKeyspace(ctx, keyspace)
		})

	actionRepo.RegisterKeyspaceAction("ValidatePermissionsKeyspace",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace string, r *http.Request) (string, error) {
			return "", wr.ValidatePermissionsKeyspace(ctx, keyspace)
		})

	// shard actions
	actionRepo.RegisterShardAction("ValidateShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidateShard(ctx, keyspace, shard, false)
		})

	actionRepo.RegisterShardAction("ValidateSchemaShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidateSchemaShard(ctx, keyspace, shard, nil, false)
		})

	actionRepo.RegisterShardAction("ValidateVersionShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidateVersionShard(ctx, keyspace, shard)
		})

	actionRepo.RegisterShardAction("ValidatePermissionsShard",
		func(ctx context.Context, wr *wrangler.Wrangler, keyspace, shard string, r *http.Request) (string, error) {
			return "", wr.ValidatePermissionsShard(ctx, keyspace, shard)
		})

	// tablet actions
	actionRepo.RegisterTabletAction("Ping", "",
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			ti, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
			if err != nil {
				return "", err
			}
			return "", wr.TabletManagerClient().Ping(ctx, ti.Tablet)
		})

	actionRepo.RegisterTabletAction("RefreshState", acl.ADMIN,
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			ti, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
			if err != nil {
				return "", err
			}
			return "", wr.TabletManagerClient().RefreshState(ctx, ti.Tablet)
		})

	actionRepo.RegisterTabletAction("DeleteTablet", acl.ADMIN,
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			return "", wr.DeleteTablet(ctx, tabletAlias, false)
		})

	actionRepo.RegisterTabletAction("ReloadSchema", acl.ADMIN,
		func(ctx context.Context, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, r *http.Request) (string, error) {
			return "", wr.ReloadSchema(ctx, tabletAlias)
		})

	// Anything unrecognized gets redirected to the main app2 page.
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, appPrefix2, http.StatusFound)
	})

	// Serve the static files for the vtctld web app.
	http.HandleFunc(appPrefix, func(w http.ResponseWriter, r *http.Request) {
		// Strip the prefix.
		parts := strings.SplitN(r.URL.Path, "/", 3)
		if len(parts) != 3 {
			http.NotFound(w, r)
			return
		}
		rest := parts[2]
		if rest == "" {
			rest = "index.html"
		}
		http.ServeFile(w, r, path.Join(*webDir, rest))
	})

	// Serve the static files for the vtctld2 web app.
	// This is a temporary additional URL for serving the new,
	// in-development UI side-by-side with the current one.
	http.HandleFunc(appPrefix2, func(w http.ResponseWriter, r *http.Request) {
		// Strip the prefix.
		parts := strings.SplitN(r.URL.Path, "/", 3)
		if len(parts) != 3 {
			http.NotFound(w, r)
			return
		}
		rest := parts[2]
		if rest == "" {
			rest = "index.html"
		}
		filePath := path.Join(*webDir2, rest)
		if _, err := os.Stat(filePath); err != nil {
			// The requested file doesn't exist.
			if strings.ContainsAny(rest, "/.") {
				// This looks like a real file path, so return Not Found.
				http.NotFound(w, r)
				return
			}
			// It looks like a virtual route path (for pages within the app).
			// For these, we must serve index.html to initialize the app.
			filePath = path.Join(*webDir2, "index.html")
		}
		http.ServeFile(w, r, filePath)
	})

	var realtimeStats *realtimeStats
	if *enableRealtimeStats {
		var err error
		realtimeStats, err = newRealtimeStats(ts)
		if err != nil {
			log.Errorf("Failed to instantiate RealtimeStats at startup: %v", err)
		}
	}

	// Serve the REST API for the vtctld web app.
	initAPI(context.Background(), ts, actionRepo, realtimeStats)

	// Init redirects for explorers
	initExplorer(ts)

	// Init workflow manager.
	initWorkflowManager(ts)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

var (
	tokenFile   = flag.String("token_file", "", "File containing Auth Token.")
	owner       = flag.String("owner", "istio", "Github Owner or org.")
	repo        = flag.String("repo", "", "Github repo within the org.")
	base        = flag.String("base", "stable", "The base branch used for PR.")
	head        = flag.String("head", "master", "The head branch used for PR.")
	pullRequest = flag.Int("pr", 0, "The Pull request to use.")
	fastForward = flag.Bool("fast_forward", false, "Creates a PR updating Base to Head.")
	verify      = flag.Bool("verify", false, "Verifies PR on Base and push them if success.")
	comment     = flag.String("comment", "", "The comment to send to the Pull Request.")
	GH          = newGhConst()
)

type ghConst struct {
	success string
	failure string
	pending string
	closed  string
	all     string
	commit  string
}

// Simple Github Helper
type helper struct {
	Owner  string
	Repo   string
	Base   string
	Head   string
	Pr     int
	Client *github.Client
}

// Get token from tokenFile is set, otherwise is anonymous.
func getToken() (*http.Client, error) {
	if *tokenFile == "" {
		return nil, nil
	}
	b, err := ioutil.ReadFile(*tokenFile)
	if err != nil {
		return nil, err
	}
	token := strings.TrimSpace(string(b[:]))
	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: string(token[:])})
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	return tc, nil
}

// Creates a new ghConst
func newGhConst() *ghConst {
	return &ghConst{
		success: "success",
		failure: "failure",
		pending: "pending",
		closed:  "closed",
		all:     "all",
		commit:  "commit",
	}
}

// Creates a new Github Helper from provided
func newHelper() (*helper, error) {
	if tc, err := getToken(); err == nil {
		if *repo == "" {
			return nil, errors.New("repo flag must be set!")
		}
		client := github.NewClient(tc)
		return &helper{
			Owner:  *owner,
			Repo:   *repo,
			Base:   *base,
			Head:   *head,
			Pr:     *pullRequest,
			Client: client,
		}, nil
	} else {
		return nil, err
	}
}

// Creates a pull request from Base branch
func (h helper) createPullRequestToBase(commit *string) error {
	if commit == nil {
		return errors.New("commit cannot be nil.")
	}
	title := fmt.Sprintf(
		"DO NOT MERGE! Fast Forward %s to %s.", h.Base, *commit)
	body := "This PR will be merged automatically once checks are successful."
	req := github.NewPullRequest{
		Head:  &h.Head,
		Base:  &h.Base,
		Title: &title,
		Body:  &body,
	}
	log.Printf("Creating a PR with Title: \"%s\"", title)
	if pr, _, err := h.Client.PullRequests.Create(h.Owner, h.Repo, &req); err == nil {
		log.Printf("Created new PR at %s", *pr.HTMLURL)
	} else {
		return err
	}
	return nil
}

// Gets the last commit from Head branch.
func (h helper) getLastCommitFromHead() (*string, error) {
	comp, _, err := h.Client.Repositories.CompareCommits(h.Owner, h.Repo, h.Head, h.Base)
	if err == nil {
		if *comp.BehindBy > 0 {
			commit := comp.BaseCommit.SHA
			log.Printf(
				"%s is %d commits ahead from %s, and HEAD commit is %s",
				h.Head, *comp.BehindBy, h.Base, *commit)
			return commit, nil
		}
	}
	return nil, err
}

// Fast forward Base branch to the last commit of Head branch.
func (h helper) fastForwardBase() error {
	commit, err := h.getLastCommitFromHead()
	if err != nil {
		return err
	}
	if commit != nil {
		options := github.PullRequestListOptions{
			Head:  h.Head,
			Base:  h.Base,
			State: GH.all,
		}

		prs, _, err := h.Client.PullRequests.List(h.Owner, h.Repo, &options)
		if err == nil {
			for _, pr := range prs {
				if strings.Contains(*pr.Title, *commit) {
					log.Printf("A PR already exist for %s", *commit)
					return nil
				}
			}
		}
		return h.createPullRequestToBase(commit)
	}
	log.Printf("Branches %s and %s are in sync.", h.Base, h.Head)
	return nil
}

// Close an existing PR
func (h helper) closePullRequest(pr *github.PullRequest) error {
	log.Printf("Closing PR %d", *pr.Number)
	*pr.State = GH.closed
	_, _, err := h.Client.PullRequests.Edit(h.Owner, h.Repo, *pr.Number, pr)
	return err
}

// Create an annotated stable tag from the given commit.
func (h helper) createStableTag(commit *string) error {
	if commit == nil {
		return errors.New("commit cannot be nil.")
	}
	sha := *commit
	tag := fmt.Sprintf("stable-%s", sha[0:7])
	message := "Stable build"
	log.Printf("Creating tag %s on %s for commit %s", tag, h.Base, *commit)
	gho := github.GitObject{
		SHA:  commit,
		Type: &GH.commit,
	}
	gt := github.Tag{
		Object:  &gho,
		Message: &message,
		Tag:     &tag,
	}
	t, resp, err := h.Client.Git.CreateTag(h.Owner, h.Repo, &gt)
	if err != nil {
		return err
	}
	log.Printf("Creating ref tag %s on %s for commit %s", tag, h.Base, *commit)
	ref := fmt.Sprintf("refs/tags/%s", tag)
	r := github.Reference{
		Ref:    &ref,
		Object: t.Object,
	}
	_, resp, err = h.Client.Git.CreateRef(h.Owner, h.Repo, &r)
	// Already exists
	if resp.StatusCode != 422 {
		return err
	}
	return nil
}

// Update the Base branch reference to a given commit.
func (h helper) updateBaseReference(commit *string) error {
	if commit == nil {
		return errors.New("commit cannot be nil")
	}
	ref := fmt.Sprintf("refs/heads/%s", h.Base)
	log.Printf("Updating ref %s to commit %s", ref, *commit)
	gho := github.GitObject{
		SHA:  commit,
		Type: &GH.commit,
	}
	r := github.Reference{
		Ref:    &ref,
		Object: &gho,
	}
	r.Ref = new(string)
	*r.Ref = ref

	_, _, err := h.Client.Git.UpdateRef(h.Owner, h.Repo, &r, false)
	return err
}

// Checks if a PR is ready to be pushed. Create a stable tag and
// fast forward Base to the PR's head commit.
func (h helper) updatePullRequest(pr *github.PullRequest, s *github.CombinedStatus) error {
	state := *s.State
	if *s.TotalCount == 0 {
		// There seems to be a bug in the Github API when no check are set.
		// The status stays in pending state forever
		state = GH.success
	}
	switch state {
	case GH.success:
		if err := h.createStableTag(s.SHA); err == nil {
			if err := h.updateBaseReference(s.SHA); err != nil {
				log.Printf("Could not update %s reference to %s.\n%v", h.Base, *s.SHA, err)
				return nil
			}
			// Note there is no need to close the PR here.
			// It will be done automatically once Base ref is updated
		} else {
			return err
		}
	case GH.failure:
		return h.closePullRequest(pr)
	case GH.pending:
		log.Printf("Pull Request %d is still being checked", *pr.Number)
	}
	return nil
}

// Checks all the PR on Base and calls updatePullRequest on each.
func (h helper) verifyPullRequestStatus() error {
	options := github.PullRequestListOptions{
		Head:  h.Head,
		Base:  h.Base,
		State: "open",
	}
	prs, _, err := h.Client.PullRequests.List(h.Owner, h.Repo, &options)
	if err != nil {
		return err
	}
	for _, pr := range prs {
		statuses, _, err := h.Client.Repositories.GetCombinedStatus(
			h.Owner, h.Repo, *pr.Head.SHA, new(github.ListOptions))
		if err == nil {
			err = h.updatePullRequest(pr, statuses)
		}
		if err != nil {
			log.Fatalf("Could not update PR %d. \n%v", *pr.Number, err)
		}
	}
	log.Printf("No more PR to verify for branch %s.", h.Base)
	return nil
}

// Creates a comment on a Pull Request
func (h helper) createComment(comment *string) error {
	if h.Pr <= 0 {
		return errors.New("PR number needs to be greather than 0")
	}
	c := github.IssueComment{
		Body: comment,
	}
	log.Printf("Commenting \"%s\" on PR %d for %s/%s", *comment, h.Pr, h.Owner, h.Repo)
	_, _, err := h.Client.Issues.CreateComment(h.Owner, h.Repo, h.Pr, &c)
	return err
}

func main() {
	flag.Parse()
	h, err := newHelper()
	if err != nil {
		log.Fatalf("Could not instantiate a github client %v", err)
	}
	if *verify {
		if err = h.verifyPullRequestStatus(); err != nil {
			log.Fatalf("Unable to verify PR from %s.\n%v", h.Base, err)
		}
	}
	if *fastForward {
		if err = h.fastForwardBase(); err != nil {
			log.Fatalf("Unable to fast forward %s.\n%v", h.Base, err)
		}
	}
	if *comment != "" {
		if err := h.createComment(comment); err != nil {
			log.Fatalf("Unable to create a comment on PR %d.\n%v", h.Pr, err)
		}
	}
}
<commit_msg>Pass a context object into github client calls (#74)<commit_after>package main

import (
	"errors"
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

var (
	tokenFile   = flag.String("token_file", "", "File containing Auth Token.")
	owner       = flag.String("owner", "istio", "Github Owner or org.")
	repo        = flag.String("repo", "", "Github repo within the org.")
	base        = flag.String("base", "stable", "The base branch used for PR.")
	head        = flag.String("head", "master", "The head branch used for PR.")
	pullRequest = flag.Int("pr", 0, "The Pull request to use.")
	fastForward = flag.Bool("fast_forward", false, "Creates a PR updating Base to Head.")
	verify      = flag.Bool("verify", false, "Verifies PR on Base and push them if success.")
	comment     = flag.String("comment", "", "The comment to send to the Pull Request.")
	GH          = newGhConst()
)

type ghConst struct {
	success string
	failure string
	pending string
	closed  string
	all     string
	commit  string
}

// Simple Github Helper
type helper struct {
	Owner  string
	Repo   string
	Base   string
	Head   string
	Pr     int
	Client *github.Client
}

// Get token from tokenFile is set, otherwise is anonymous.
func getToken() (*http.Client, error) {
	if *tokenFile == "" {
		return nil, nil
	}
	b, err := ioutil.ReadFile(*tokenFile)
	if err != nil {
		return nil, err
	}
	token := strings.TrimSpace(string(b[:]))
	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: string(token[:])})
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	return tc, nil
}

// Creates a new ghConst
func newGhConst() *ghConst {
	return &ghConst{
		success: "success",
		failure: "failure",
		pending: "pending",
		closed:  "closed",
		all:     "all",
		commit:  "commit",
	}
}

// Creates a new Github Helper from provided
func newHelper() (*helper, error) {
	if tc, err := getToken(); err == nil {
		if *repo == "" {
			return nil, errors.New("repo flag must be set!")
		}
		client := github.NewClient(tc)
		return &helper{
			Owner:  *owner,
			Repo:   *repo,
			Base:   *base,
			Head:   *head,
			Pr:     *pullRequest,
			Client: client,
		}, nil
	} else {
		return nil, err
	}
}

// Creates a pull request from Base branch
func (h helper) createPullRequestToBase(commit *string) error {
	if commit == nil {
		return errors.New("commit cannot be nil.")
	}
	title := fmt.Sprintf(
		"DO NOT MERGE! Fast Forward %s to %s.", h.Base, *commit)
	body := "This PR will be merged automatically once checks are successful."
	req := github.NewPullRequest{
		Head:  &h.Head,
		Base:  &h.Base,
		Title: &title,
		Body:  &body,
	}
	log.Printf("Creating a PR with Title: \"%s\"", title)
	if pr, _, err := h.Client.PullRequests.Create(context.TODO(), h.Owner, h.Repo, &req); err == nil {
		log.Printf("Created new PR at %s", *pr.HTMLURL)
	} else {
		return err
	}
	return nil
}

// Gets the last commit from Head branch.
func (h helper) getLastCommitFromHead() (*string, error) {
	comp, _, err := h.Client.Repositories.CompareCommits(context.TODO(), h.Owner, h.Repo, h.Head, h.Base)
	if err == nil {
		if *comp.BehindBy > 0 {
			commit := comp.BaseCommit.SHA
			log.Printf(
				"%s is %d commits ahead from %s, and HEAD commit is %s",
				h.Head, *comp.BehindBy, h.Base, *commit)
			return commit, nil
		}
	}
	return nil, err
}

// Fast forward Base branch to the last commit of Head branch.
func (h helper) fastForwardBase() error {
	commit, err := h.getLastCommitFromHead()
	if err != nil {
		return err
	}
	if commit != nil {
		options := github.PullRequestListOptions{
			Head:  h.Head,
			Base:  h.Base,
			State: GH.all,
		}

		prs, _, err := h.Client.PullRequests.List(context.TODO(), h.Owner, h.Repo, &options)
		if err == nil {
			for _, pr := range prs {
				if strings.Contains(*pr.Title, *commit) {
					log.Printf("A PR already exist for %s", *commit)
					return nil
				}
			}
		}
		return h.createPullRequestToBase(commit)
	}
	log.Printf("Branches %s and %s are in sync.", h.Base, h.Head)
	return nil
}

// Close an existing PR
func (h helper) closePullRequest(pr *github.PullRequest) error {
	log.Printf("Closing PR %d", *pr.Number)
	*pr.State = GH.closed
	_, _, err := h.Client.PullRequests.Edit(context.TODO(), h.Owner, h.Repo, *pr.Number, pr)
	return err
}

// Create an annotated stable tag from the given commit.
func (h helper) createStableTag(commit *string) error {
	if commit == nil {
		return errors.New("commit cannot be nil.")
	}
	sha := *commit
	tag := fmt.Sprintf("stable-%s", sha[0:7])
	message := "Stable build"
	log.Printf("Creating tag %s on %s for commit %s", tag, h.Base, *commit)
	gho := github.GitObject{
		SHA:  commit,
		Type: &GH.commit,
	}
	gt := github.Tag{
		Object:  &gho,
		Message: &message,
		Tag:     &tag,
	}
	t, resp, err := h.Client.Git.CreateTag(context.TODO(), h.Owner, h.Repo, &gt)
	if err != nil {
		return err
	}
	log.Printf("Creating ref tag %s on %s for commit %s", tag, h.Base, *commit)
	ref := fmt.Sprintf("refs/tags/%s", tag)
	r := github.Reference{
		Ref:    &ref,
		Object: t.Object,
	}
	_, resp, err = h.Client.Git.CreateRef(context.TODO(), h.Owner, h.Repo, &r)
	// Already exists
	if resp.StatusCode != 422 {
		return err
	}
	return nil
}

// Update the Base branch reference to a given commit.
func (h helper) updateBaseReference(commit *string) error {
	if commit == nil {
		return errors.New("commit cannot be nil")
	}
	ref := fmt.Sprintf("refs/heads/%s", h.Base)
	log.Printf("Updating ref %s to commit %s", ref, *commit)
	gho := github.GitObject{
		SHA:  commit,
		Type: &GH.commit,
	}
	r := github.Reference{
		Ref:    &ref,
		Object: &gho,
	}
	r.Ref = new(string)
	*r.Ref = ref

	_, _, err := h.Client.Git.UpdateRef(context.TODO(), h.Owner, h.Repo, &r, false)
	return err
}

// Checks if a PR is ready to be pushed. Create a stable tag and
// fast forward Base to the PR's head commit.
func (h helper) updatePullRequest(pr *github.PullRequest, s *github.CombinedStatus) error {
	state := *s.State
	if *s.TotalCount == 0 {
		// There seems to be a bug in the Github API when no check are set.
		// The status stays in pending state forever
		state = GH.success
	}
	switch state {
	case GH.success:
		if err := h.createStableTag(s.SHA); err == nil {
			if err := h.updateBaseReference(s.SHA); err != nil {
				log.Printf("Could not update %s reference to %s.\n%v", h.Base, *s.SHA, err)
				return nil
			}
			// Note there is no need to close the PR here.
			// It will be done automatically once Base ref is updated
		} else {
			return err
		}
	case GH.failure:
		return h.closePullRequest(pr)
	case GH.pending:
		log.Printf("Pull Request %d is still being checked", *pr.Number)
	}
	return nil
}

// Checks all the PR on Base and calls updatePullRequest on each.
func (h helper) verifyPullRequestStatus() error {
	options := github.PullRequestListOptions{
		Head:  h.Head,
		Base:  h.Base,
		State: "open",
	}
	prs, _, err := h.Client.PullRequests.List(context.TODO(), h.Owner, h.Repo, &options)
	if err != nil {
		return err
	}
	for _, pr := range prs {
		statuses, _, err := h.Client.Repositories.GetCombinedStatus(
			context.TODO(), h.Owner, h.Repo, *pr.Head.SHA, new(github.ListOptions))
		if err == nil {
			err = h.updatePullRequest(pr, statuses)
		}
		if err != nil {
			log.Fatalf("Could not update PR %d. \n%v", *pr.Number, err)
		}
	}
	log.Printf("No more PR to verify for branch %s.", h.Base)
	return nil
}

// Creates a comment on a Pull Request
func (h helper) createComment(comment *string) error {
	if h.Pr <= 0 {
		return errors.New("PR number needs to be greather than 0")
	}
	c := github.IssueComment{
		Body: comment,
	}
	log.Printf("Commenting \"%s\" on PR %d for %s/%s", *comment, h.Pr, h.Owner, h.Repo)
	_, _, err := h.Client.Issues.CreateComment(context.TODO(), h.Owner, h.Repo, h.Pr, &c)
	return err
}

func main() {
	flag.Parse()
	h, err := newHelper()
	if err != nil {
		log.Fatalf("Could not instantiate a github client %v", err)
	}
	if *verify {
		if err = h.verifyPullRequestStatus(); err != nil {
			log.Fatalf("Unable to verify PR from %s.\n%v", h.Base, err)
		}
	}
	if *fastForward {
		if err = h.fastForwardBase(); err != nil {
			log.Fatalf("Unable to fast forward %s.\n%v", h.Base, err)
		}
	}
	if *comment != "" {
		if err := h.createComment(comment); err != nil {
			log.Fatalf("Unable to create a comment on PR %d.\n%v", h.Pr, err)
		}
	}
}
<|endoftext|>


<commit_before>package resources

import (
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
)

type IAMRole struct {
	svc  *iam.IAM
	name string
	path string
}

func init() {
	register("IAMRole", ListIAMRoles)
}

func ListIAMRoles(sess *session.Session) ([]Resource, error) {
	svc := iam.New(sess)
	// REMOVE: Baesd on https://github.com/rebuy-de/aws-nuke/blob/master/resources/ecr-repository.go
	params := &iam.ListRolesInput{}
	resources := make([]Resource, 0)

	for {
		resp, err := svc.ListRoles(params)
		if err != nil {
			return nil, err
		}

		for _, out := range resp.Roles {
			resources = append(resources, &IAMRole{
				svc:  svc,
				name: *out.RoleName,
				path: *out.Path,
			})
		}

		if *resp.IsTruncated == false {
			break
		}

		params.Marker = resp.Marker
	}

	return resources, nil
}

func (e *IAMRole) Filter() error {
	if strings.HasPrefix(e.path, "/aws-service-role/") {
		return fmt.Errorf("cannot delete service roles")
	}
	return nil
}

func (e *IAMRole) Remove() error {
	_, err := e.svc.DeleteRole(&iam.DeleteRoleInput{
		RoleName: &e.name,
	})
	if err != nil {
		return err
	}

	return nil
}

func (e *IAMRole) String() string {
	return e.name
}
<commit_msg>Remove comment.<commit_after>package resources

import (
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
)

type IAMRole struct {
	svc  *iam.IAM
	name string
	path string
}

func init() {
	register("IAMRole", ListIAMRoles)
}

func ListIAMRoles(sess *session.Session) ([]Resource, error) {
	svc := iam.New(sess)
	params := &iam.ListRolesInput{}
	resources := make([]Resource, 0)

	for {
		resp, err := svc.ListRoles(params)
		if err != nil {
			return nil, err
		}

		for _, out := range resp.Roles {
			resources = append(resources, &IAMRole{
				svc:  svc,
				name: *out.RoleName,
				path: *out.Path,
			})
		}

		if *resp.IsTruncated == false {
			break
		}

		params.Marker = resp.Marker
	}

	return resources, nil
}

func (e *IAMRole) Filter() error {
	if strings.HasPrefix(e.path, "/aws-service-role/") {
		return fmt.Errorf("cannot delete service roles")
	}
	return nil
}

func (e *IAMRole) Remove() error {
	_, err := e.svc.DeleteRole(&iam.DeleteRoleInput{
		RoleName: &e.name,
	})
	if err != nil {
		return err
	}

	return nil
}

func (e *IAMRole) String() string {
	return e.name
}
<|endoftext|>


<commit_before>package gocd

import (
	"context"
	"fmt"
	"strings"
)

// PipelineConfigsService describes the HAL _link resource for the api response object for a pipelineconfig
type PipelineConfigsService service

// PipelineConfig describes the configuration for a pipeline
type PipelineConfig struct{}

// PipelineConfigRequest describes a request object for creating or updating pipelines
type PipelineConfigRequest struct {
	Group    string    `json:"group,omitempty"`
	Pipeline *Pipeline `json:"pipeline"`
}

// Get a single PipelineTemplate object in the GoCD API.
func (pcs *PipelineConfigsService) Get(ctx context.Context, name string) (*Pipeline, *APIResponse, error) {
	p := Pipeline{}
	_, resp, err := pcs.client.getAction(ctx, &APIClientRequest{
		Path:         "admin/pipelines/" + name,
		APIVersion:   apiV4,
		ResponseBody: &p,
	})

	if err == nil {
		p.Version = strings.Replace(resp.HTTP.Header.Get("Etag"), "\"", "", -1)
	}

	return &p, resp, err
}

// Update a pipeline configuration
func (pcs *PipelineConfigsService) Update(ctx context.Context, name string, p *Pipeline) (*Pipeline, *APIResponse, error) {

	pt := &PipelineConfigRequest{
		Pipeline: p,
	}
	pr := Pipeline{}

	_, resp, err := pcs.client.putAction(ctx, &APIClientRequest{
		Path:         "admin/pipelines/" + name,
		APIVersion:   apiV4,
		RequestBody:  pt,
		ResponseBody: &pr,
		Headers: map[string]string{
			"If-Match": fmt.Sprintf("\"%s\"", p.Version),
		},
	})

	return &pr, resp, err

}

// Create a pipeline configuration
func (pcs *PipelineConfigsService) Create(ctx context.Context, group string, p *Pipeline) (*Pipeline, *APIResponse, error) {
	pt := &PipelineConfigRequest{
		Group:    group,
		Pipeline: p,
	}
	pc := Pipeline{}

	_, resp, err := pcs.client.postAction(ctx, &APIClientRequest{
		Path:         "admin/pipelines",
		APIVersion:   apiV4,
		RequestBody:  pt,
		ResponseBody: &pc,
	})
	if err == nil {
		pc.Version = strings.Replace(resp.HTTP.Header.Get("Etag"), "\"", "", -1)
	}

	return &pc, resp, err
}

// Delete a pipeline configuration
func (pcs *PipelineConfigsService) Delete(ctx context.Context, name string) (string, *APIResponse, error) {
	return pcs.client.deleteAction(ctx, "admin/pipelines/"+name, apiV4)
}
<commit_msg>Removed unused struct<commit_after>package gocd

import (
	"context"
	"fmt"
	"strings"
)

// PipelineConfigsService describes the HAL _link resource for the api response object for a pipelineconfig
type PipelineConfigsService service

// PipelineConfigRequest describes a request object for creating or updating pipelines
type PipelineConfigRequest struct {
	Group    string    `json:"group,omitempty"`
	Pipeline *Pipeline `json:"pipeline"`
}

// Get a single PipelineTemplate object in the GoCD API.
func (pcs *PipelineConfigsService) Get(ctx context.Context, name string) (*Pipeline, *APIResponse, error) {
	p := Pipeline{}
	_, resp, err := pcs.client.getAction(ctx, &APIClientRequest{
		Path:         "admin/pipelines/" + name,
		APIVersion:   apiV4,
		ResponseBody: &p,
	})
	if err == nil {
		p.Version = strings.Replace(resp.HTTP.Header.Get("Etag"), "\"", "", -1)
	}

	return &p, resp, err
}

// Update a pipeline configuration
func (pcs *PipelineConfigsService) Update(ctx context.Context, name string, p *Pipeline) (*Pipeline, *APIResponse, error) {

	pt := &PipelineConfigRequest{
		Pipeline: p,
	}
	pr := Pipeline{}

	_, resp, err := pcs.client.putAction(ctx, &APIClientRequest{
		Path:         "admin/pipelines/" + name,
		APIVersion:   apiV4,
		RequestBody:  pt,
		ResponseBody: &pr,
		Headers: map[string]string{
			"If-Match": fmt.Sprintf("\"%s\"", p.Version),
		},
	})

	return &pr, resp, err

}

// Create a pipeline configuration
func (pcs *PipelineConfigsService) Create(ctx context.Context, group string, p *Pipeline) (*Pipeline, *APIResponse, error) {
	pt := &PipelineConfigRequest{
		Group:    group,
		Pipeline: p,
	}
	pc := Pipeline{}

	_, resp, err := pcs.client.postAction(ctx, &APIClientRequest{
		Path:         "admin/pipelines",
		APIVersion:   apiV4,
		RequestBody:  pt,
		ResponseBody: &pc,
	})
	if err == nil {
		pc.Version = strings.Replace(resp.HTTP.Header.Get("Etag"), "\"", "", -1)
	}

	return &pc, resp, err
}

// Delete a pipeline configuration
func (pcs *PipelineConfigsService) Delete(ctx context.Context, name string) (string, *APIResponse, error) {
	return pcs.client.deleteAction(ctx, "admin/pipelines/"+name, apiV4)
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"reflect"
	"time"
)

// MessageFormat indicates how the connector should display the content of
// the message. One of Variable, Fixed or Raw (To Be Implemented)
type MessageFormat int

// Outgoing message format, Variable or Fixed
const (
	Variable MessageFormat = iota // variable font width
	Fixed
	Raw
)

// Robot is passed to the plugin to enable convenience functions Say and Reply
type Robot struct {
	User     string        // The user who sent the message; this can be modified for replying to an arbitrary user
	Channel  string        // The channel where the message was received, or "" for a direct message. This can be modified to send a message to an arbitrary channel.
	Format   MessageFormat // The outgoing message format, one of Fixed or Variable
	pluginID string        // Pass the ID in for later identificaton of the plugin
}

/* robot.go defines some convenience functions on struct Robot to
   simplify use by plugins. */

// CheckAdmin returns true if the user is a configured administrator of the
// robot. Should be used sparingly, when a single plugin has multiple commands,
// some which require admin. Otherwise the plugin should just configure
// RequireAdmin: true
func (r *Robot) CheckAdmin() bool {
	b.lock.RLock()
	defer b.lock.RUnlock()
	for _, adminUser := range b.adminUsers {
		if r.User == adminUser {
			return true
		}
	}
	return false
}

// Elevate lets a plugin request elevation on the fly. When immediate = true,
// the elevator should always prompt for 2fa; otherwise a configured timeout
// should apply.
func (r *Robot) Elevate(immediate bool) bool {
	b.lock.RLock()
	e := b.elevator
	p := b.elevatorProvider
	b.lock.RUnlock()
	r.pluginID = "elevator-" + p
	return e(r, immediate)
}

// Fixed is a convenience function for sending a message with fixed width
// font. e.g. r.Reply(xxx) replies in variable width font, but
// r.Fixed().Reply(xxx) replies in a fixed-width font.
func (r *Robot) Fixed() *Robot {
	nr := *r
	nr.Format = Fixed
	return &nr
}

// Direct is a convenience function for initiating a DM conversation with a
// user. Created initially so a plugin could prompt for a password in a DM.
func (r *Robot) Direct() *Robot {
	nr := *r
	nr.Channel = ""
	return &nr
}

// Pause is a convenience function to pause some fractional number of seconds.
func (r *Robot) Pause(s float64) {
	ms := time.Duration(s * float64(1000))
	time.Sleep(ms * time.Millisecond)
}

// RandomString is a convenience function for returning a random string
// from a slice of strings, so that replies can vary.
func (r *Robot) RandomString(s []string) string {
	l := len(s)
	if l == 0 {
		return ""
	}
	return s[random.Intn(l)]
}

// RandomInt uses the robot's seeded random to return a random int 0 <= retval < n
func (r *Robot) RandomInt(n int) int {
	return random.Intn(n)
}

// GetBotAttribute returns an attribute of the robot or "" if unknown.
// Current attributes:
// name, alias, fullName, contact
func (r *Robot) GetBotAttribute(a string) *BotAttrRet {
	b.lock.RLock()
	defer b.lock.RUnlock()
	ret := Ok
	var attr string
	switch a {
	case "name":
		attr = b.name
	case "fullName", "realName":
		attr = b.fullName
	case "alias":
		attr = string(b.alias)
	case "email":
		attr = b.email
	case "contact", "admin", "adminContact":
		attr = b.adminContact
	default:
		ret = AttributeNotFound
	}
	return &BotAttrRet{attr, ret}
}

// GetUserAttribute returns a BotAttrRet with
// - The string Attribute of a user, or "" if unknown/error
// - A RetVal which is one of Ok, UserNotFound, AttributeNotFound
// Current attributes:
// name(handle), fullName, email, firstName, lastName, phone
// TODO: supplement data with gopherbot.json user's table
func (r *Robot) GetUserAttribute(u, a string) *BotAttrRet {
	attr, ret := b.GetProtocolUserAttribute(u, a)
	return &BotAttrRet{attr, ret}
}

// GetSenderAttribute returns a BotAttrRet with
// - The string Attribute of the sender, or "" if unknown/error
// - A RetVal which is one of Ok, UserNotFound, AttributeNotFound
// Current attributes:
// name(handle), fullName, email, firstName, lastName, phone
// TODO: supplement data with gopherbot.json user's table
func (r *Robot) GetSenderAttribute(a string) *BotAttrRet {
	attr, ret := b.GetProtocolUserAttribute(r.User, a)
	return &BotAttrRet{attr, ret}
}

/*

GetPluginConfig sets a struct pointer to point to a config struct populated
from configuration when plugins were loaded. To use, a plugin should define
a struct for it's configuration data, e.g.:

	type pConf struct {
		Username, Password string
	}

In conf/plugins/<pluginname>.yaml, you would add a Config: stanza, e.g.:

	Config:
	  Username: foo
	  Password: bar

When registering the plugin, you pass a pointer to an empty config template, which the
robot will use to populate a struct when configuration is loaded:

	func init() {
		bot.RegisterPlugin("memes", bot.PluginHandler{
			DefaultConfig: defaultConfig, // yaml string providing default configuration
			Handler:       plugfunc, // callback function
			Config:        &pConf{}, // pointer to empty config struct
		})
	}

Then, to get a current copy of configuration when the plugin runs, define a struct pointer
and call GetPluginConfig with a double-pointer:

	var c *pConf
	r.GetPluginConfig(&c)

... And voila! *pConf is populated with the contents from the configured Config: stanza
*/
func (r *Robot) GetPluginConfig(dptr interface{}) RetVal {
	b.lock.RLock()
	defer b.lock.RUnlock()
	plugin := plugins[plugIDmap[r.pluginID]]
	if plugin.config == nil {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig, but no config was found.", plugin.name))
		return NoConfigFound
	}
	tp := reflect.ValueOf(dptr)
	if tp.Kind() != reflect.Ptr {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig, but didn't pass a double-pointer to a struct", plugin.name))
		return InvalidDblPtr
	}
	p := reflect.Indirect(tp)
	if p.Kind() != reflect.Ptr {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig, but didn't pass a double-pointer to a struct", plugin.name))
		return InvalidDblPtr
	}
	if p.Type() != reflect.ValueOf(plugin.config).Type() {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig with an invalid double-pointer", plugin.name))
		return InvalidCfgStruct
	}
	p.Set(reflect.ValueOf(plugin.config))
	return Ok
}

// Log logs a message to the robot's log file (or stderr) if the level
// is lower than or equal to the robot's current log level
func (r *Robot) Log(l LogLevel, v ...interface{}) {
	Log(l, v...)
}

// SendChannelMessage lets a plugin easily send a message to an arbitrary
// channel. Use Robot.Fixed().SencChannelMessage(...) for fixed-width
// font.
func (r *Robot) SendChannelMessage(channel, msg string) RetVal {
	return b.SendProtocolChannelMessage(channel, msg, r.Format)
}

// SendUserChannelMessage lets a plugin easily send a message directed to
// a specific user in a specific channel without fiddling with the robot
// object. Use Robot.Fixed().SencChannelMessage(...) for fixed-width
// font.
func (r *Robot) SendUserChannelMessage(user, channel, msg string) RetVal {
	return b.SendProtocolUserChannelMessage(user, channel, msg, r.Format)
}

// SendUserMessage lets a plugin easily send a DM to a user. If a DM
// isn't possible, the connector should message the user in a channel.
func (r *Robot) SendUserMessage(user, msg string) RetVal {
	return b.SendProtocolUserMessage(user, msg, r.Format)
}

// Reply directs a message to the user
func (r *Robot) Reply(msg string) RetVal {
	if r.Channel == "" {
		return b.SendProtocolUserMessage(r.User, msg, r.Format)
	}
	return b.SendProtocolUserChannelMessage(r.User, r.Channel, msg, r.Format)
}

// Say just sends a message to the user or channel
func (r *Robot) Say(msg string) RetVal {
	if r.Channel == "" {
		return b.SendProtocolUserMessage(r.User, msg, r.Format)
	}
	return b.SendProtocolChannelMessage(r.Channel, msg, r.Format)
}
<commit_msg>Check elevator before using it<commit_after>package bot

import (
	"fmt"
	"reflect"
	"time"
)

// MessageFormat indicates how the connector should display the content of
// the message. One of Variable, Fixed or Raw (To Be Implemented)
type MessageFormat int

// Outgoing message format, Variable or Fixed
const (
	Variable MessageFormat = iota // variable font width
	Fixed
	Raw
)

// Robot is passed to the plugin to enable convenience functions Say and Reply
type Robot struct {
	User     string        // The user who sent the message; this can be modified for replying to an arbitrary user
	Channel  string        // The channel where the message was received, or "" for a direct message. This can be modified to send a message to an arbitrary channel.
	Format   MessageFormat // The outgoing message format, one of Fixed or Variable
	pluginID string        // Pass the ID in for later identificaton of the plugin
}

/* robot.go defines some convenience functions on struct Robot to
   simplify use by plugins. */

// CheckAdmin returns true if the user is a configured administrator of the
// robot. Should be used sparingly, when a single plugin has multiple commands,
// some which require admin. Otherwise the plugin should just configure
// RequireAdmin: true
func (r *Robot) CheckAdmin() bool {
	b.lock.RLock()
	defer b.lock.RUnlock()
	for _, adminUser := range b.adminUsers {
		if r.User == adminUser {
			return true
		}
	}
	return false
}

// Elevate lets a plugin request elevation on the fly. When immediate = true,
// the elevator should always prompt for 2fa; otherwise a configured timeout
// should apply.
func (r *Robot) Elevate(immediate bool) bool {
	b.lock.RLock()
	e := b.elevator
	p := b.elevatorProvider
	b.lock.RUnlock()
	if e == nil {
		return false
	}
	r.pluginID = "elevator-" + p
	return e(r, immediate)
}

// Fixed is a convenience function for sending a message with fixed width
// font. e.g. r.Reply(xxx) replies in variable width font, but
// r.Fixed().Reply(xxx) replies in a fixed-width font.
func (r *Robot) Fixed() *Robot {
	nr := *r
	nr.Format = Fixed
	return &nr
}

// Direct is a convenience function for initiating a DM conversation with a
// user. Created initially so a plugin could prompt for a password in a DM.
func (r *Robot) Direct() *Robot {
	nr := *r
	nr.Channel = ""
	return &nr
}

// Pause is a convenience function to pause some fractional number of seconds.
func (r *Robot) Pause(s float64) {
	ms := time.Duration(s * float64(1000))
	time.Sleep(ms * time.Millisecond)
}

// RandomString is a convenience function for returning a random string
// from a slice of strings, so that replies can vary.
func (r *Robot) RandomString(s []string) string {
	l := len(s)
	if l == 0 {
		return ""
	}
	return s[random.Intn(l)]
}

// RandomInt uses the robot's seeded random to return a random int 0 <= retval < n
func (r *Robot) RandomInt(n int) int {
	return random.Intn(n)
}

// GetBotAttribute returns an attribute of the robot or "" if unknown.
// Current attributes:
// name, alias, fullName, contact
func (r *Robot) GetBotAttribute(a string) *BotAttrRet {
	b.lock.RLock()
	defer b.lock.RUnlock()
	ret := Ok
	var attr string
	switch a {
	case "name":
		attr = b.name
	case "fullName", "realName":
		attr = b.fullName
	case "alias":
		attr = string(b.alias)
	case "email":
		attr = b.email
	case "contact", "admin", "adminContact":
		attr = b.adminContact
	default:
		ret = AttributeNotFound
	}
	return &BotAttrRet{attr, ret}
}

// GetUserAttribute returns a BotAttrRet with
// - The string Attribute of a user, or "" if unknown/error
// - A RetVal which is one of Ok, UserNotFound, AttributeNotFound
// Current attributes:
// name(handle), fullName, email, firstName, lastName, phone
// TODO: supplement data with gopherbot.json user's table
func (r *Robot) GetUserAttribute(u, a string) *BotAttrRet {
	attr, ret := b.GetProtocolUserAttribute(u, a)
	return &BotAttrRet{attr, ret}
}

// GetSenderAttribute returns a BotAttrRet with
// - The string Attribute of the sender, or "" if unknown/error
// - A RetVal which is one of Ok, UserNotFound, AttributeNotFound
// Current attributes:
// name(handle), fullName, email, firstName, lastName, phone
// TODO: supplement data with gopherbot.json user's table
func (r *Robot) GetSenderAttribute(a string) *BotAttrRet {
	attr, ret := b.GetProtocolUserAttribute(r.User, a)
	return &BotAttrRet{attr, ret}
}

/*

GetPluginConfig sets a struct pointer to point to a config struct populated
from configuration when plugins were loaded. To use, a plugin should define
a struct for it's configuration data, e.g.:

	type pConf struct {
		Username, Password string
	}

In conf/plugins/<pluginname>.yaml, you would add a Config: stanza, e.g.:

	Config:
	  Username: foo
	  Password: bar

When registering the plugin, you pass a pointer to an empty config template, which the
robot will use to populate a struct when configuration is loaded:

	func init() {
		bot.RegisterPlugin("memes", bot.PluginHandler{
			DefaultConfig: defaultConfig, // yaml string providing default configuration
			Handler:       plugfunc, // callback function
			Config:        &pConf{}, // pointer to empty config struct
		})
	}

Then, to get a current copy of configuration when the plugin runs, define a struct pointer
and call GetPluginConfig with a double-pointer:

	var c *pConf
	r.GetPluginConfig(&c)

... And voila! *pConf is populated with the contents from the configured Config: stanza
*/
func (r *Robot) GetPluginConfig(dptr interface{}) RetVal {
	b.lock.RLock()
	defer b.lock.RUnlock()
	plugin := plugins[plugIDmap[r.pluginID]]
	if plugin.config == nil {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig, but no config was found.", plugin.name))
		return NoConfigFound
	}
	tp := reflect.ValueOf(dptr)
	if tp.Kind() != reflect.Ptr {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig, but didn't pass a double-pointer to a struct", plugin.name))
		return InvalidDblPtr
	}
	p := reflect.Indirect(tp)
	if p.Kind() != reflect.Ptr {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig, but didn't pass a double-pointer to a struct", plugin.name))
		return InvalidDblPtr
	}
	if p.Type() != reflect.ValueOf(plugin.config).Type() {
		Log(Debug, fmt.Sprintf("Plugin \"%s\" called GetPluginConfig with an invalid double-pointer", plugin.name))
		return InvalidCfgStruct
	}
	p.Set(reflect.ValueOf(plugin.config))
	return Ok
}

// Log logs a message to the robot's log file (or stderr) if the level
// is lower than or equal to the robot's current log level
func (r *Robot) Log(l LogLevel, v ...interface{}) {
	Log(l, v...)
}

// SendChannelMessage lets a plugin easily send a message to an arbitrary
// channel. Use Robot.Fixed().SencChannelMessage(...) for fixed-width
// font.
func (r *Robot) SendChannelMessage(channel, msg string) RetVal {
	return b.SendProtocolChannelMessage(channel, msg, r.Format)
}

// SendUserChannelMessage lets a plugin easily send a message directed to
// a specific user in a specific channel without fiddling with the robot
// object. Use Robot.Fixed().SencChannelMessage(...) for fixed-width
// font.
func (r *Robot) SendUserChannelMessage(user, channel, msg string) RetVal {
	return b.SendProtocolUserChannelMessage(user, channel, msg, r.Format)
}

// SendUserMessage lets a plugin easily send a DM to a user. If a DM
// isn't possible, the connector should message the user in a channel.
func (r *Robot) SendUserMessage(user, msg string) RetVal {
	return b.SendProtocolUserMessage(user, msg, r.Format)
}

// Reply directs a message to the user
func (r *Robot) Reply(msg string) RetVal {
	if r.Channel == "" {
		return b.SendProtocolUserMessage(r.User, msg, r.Format)
	}
	return b.SendProtocolUserChannelMessage(r.User, r.Channel, msg, r.Format)
}

// Say just sends a message to the user or channel
func (r *Robot) Say(msg string) RetVal {
	if r.Channel == "" {
		return b.SendProtocolUserMessage(r.User, msg, r.Format)
	}
	return b.SendProtocolChannelMessage(r.Channel, msg, r.Format)
}
<|endoftext|>


<commit_before>package bufferapi

import (
	"bytes"
	"code.google.com/p/goauth2/oauth"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"regexp"
)

type Client struct {
	AccessToken string
	transport   *oauth.Transport
}

type Profile struct {
	Id                string         `json:"id"`
	UserId            string         `json:"user_id"`
	Avatar            string         `json:"avatar"`
	CreatedAt         int            `json:"created_at"`
	Default           bool           `json:"default"`
	FormattedUsername string         `json:"formatted_username"`
	Schedules         []Schedule     `json:"schedules"`
	Service           string         `json:service"`
	ServiceId         string         `json:"service_id"`
	ServiceUsername   string         `json:"service_username"`
	Statistics        map[string]int `json:"statistics"`
	TeamMembers       []string       `json:"team_members"`
	Timezone          string         `json:"timezone"`
}

type Schedule struct {
	Days  []string `json:"days"`
	Times []string `json:"times"`
}

type NewUpdate struct {
	Text       string            `json:"text"`
	ProfileIds []string          `json:"profile_ids"`
	Shorten    bool              `json:"shorten"`
	Now        bool              `json:"now"`
	Media      map[string]string `json:"media"`
}

type Update struct {
	Id             string            `json:"id"`
	CreatedAt      int               `json:"created_at"`
	Day            string            `json:"day"`
	DueAt          int               `json:"due_at"`
	DueTime        string            `json:"due_time"`
	media          map[string]string `json:"media"`
	ProfileId      string            `json:"profile_id"`
	ProfileService string            `json:"profile_service"`
	Status         string            `json:"status"`
	Text           string            `json:"text"`
	TextFormatted  string            `json:"text_formatted"`
	UserId         string            `json:"user_id"`
	Via            string            `json:"via"`
}

type UpdateResponse struct {
	Success          bool     `json:"success"`
	BufferCount      int      `json:"buffer_count"`
	BufferPercentage int      `json:"buffer_percentage"`
	Updates          []Update `json:"updates"`
}

func ClientFactory(token string) *Client {
	config := &oauth.Config{}
	transport := &oauth.Transport{Config: config}
	t := &oauth.Token{AccessToken: token}
	transport.Token = t
	c := Client{AccessToken: token, transport: transport}
	return &c
}

func (c *Client) API(method, url string, data interface{}) (respBody []byte, err error) {
	jsonPattern, _ := regexp.Compile(`\.json$`)
	if !jsonPattern.Match([]byte(url)) {
		url += ".json"
	}
	url = "https://api.bufferapp.com/1/" + url

	jsonBody, _ := json.Marshal(data)
	b := bytes.NewBuffer(jsonBody)

	var resp *http.Response
	switch method {
	case "get":
		resp, err = c.transport.Client().Get(url)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()

	case "post":
		resp, err := c.transport.Client().Post(url, "application/json", b)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()

	default:
		return nil, errors.New("Not a valid request type")
	}

	if resp.StatusCode >= 400 {
		return nil, errors.New(resp.Status)
	}

	respBody, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if len(respBody) < 2 {
		return nil, errors.New("Malformed JSON response")
	}
	return respBody, nil
}

func (c *Client) Get(url string) (resp []byte, err error) {
	return c.API("get", url, nil)
}

func (c *Client) Post(url string, params interface{}) (resp []byte, err error) {
	return c.API("post", url, params)
}

func (c *Client) Profiles() (profiles *[]Profile, err error) {
	body, err := c.Get("profiles.json")
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(body, profiles)
	if err != nil {
		return nil, err
	}
	return
}

func (c *Client) Update(update *NewUpdate) (resp *UpdateResponse, err error) {
	respBody, err := c.Post("updates/create.json", update)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(respBody, resp)
	if err != nil {
		return nil, err
	} else if !resp.Success {
		return nil, errors.New(respBody)
	}
	return resp, nil
}
<commit_msg>pass config into the client factory<commit_after>package bufferapi

import (
	"bytes"
	"code.google.com/p/goauth2/oauth"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"regexp"
)

type Client struct {
	AccessToken string
	transport   *oauth.Transport
}

type Profile struct {
	Id                string         `json:"id"`
	UserId            string         `json:"user_id"`
	Avatar            string         `json:"avatar"`
	CreatedAt         int            `json:"created_at"`
	Default           bool           `json:"default"`
	FormattedUsername string         `json:"formatted_username"`
	Schedules         []Schedule     `json:"schedules"`
	Service           string         `json:service"`
	ServiceId         string         `json:"service_id"`
	ServiceUsername   string         `json:"service_username"`
	Statistics        map[string]int `json:"statistics"`
	TeamMembers       []string       `json:"team_members"`
	Timezone          string         `json:"timezone"`
}

type Schedule struct {
	Days  []string `json:"days"`
	Times []string `json:"times"`
}

type NewUpdate struct {
	Text       string            `json:"text"`
	ProfileIds []string          `json:"profile_ids"`
	Shorten    bool              `json:"shorten"`
	Now        bool              `json:"now"`
	Media      map[string]string `json:"media"`
}

type Update struct {
	Id             string            `json:"id"`
	CreatedAt      int               `json:"created_at"`
	Day            string            `json:"day"`
	DueAt          int               `json:"due_at"`
	DueTime        string            `json:"due_time"`
	media          map[string]string `json:"media"`
	ProfileId      string            `json:"profile_id"`
	ProfileService string            `json:"profile_service"`
	Status         string            `json:"status"`
	Text           string            `json:"text"`
	TextFormatted  string            `json:"text_formatted"`
	UserId         string            `json:"user_id"`
	Via            string            `json:"via"`
}

type UpdateResponse struct {
	Success          bool     `json:"success"`
	BufferCount      int      `json:"buffer_count"`
	BufferPercentage int      `json:"buffer_percentage"`
	Updates          []Update `json:"updates"`
}

func ClientFactory(token, clientId, clientSecret, scope, authUrl, tokenUrl, cacheFile string) *Client {
	config := &oauth.Config{
		ClientId:     clientId,
		ClientSecret: clientSecret,
		Scope:        scope,
		AuthURL:      authUrl,
		TokenURL:     tokenUrl,
		TokenCache:   oauth.CacheFile(cacheFile),
	}
	transport := &oauth.Transport{Config: config}
	t := &oauth.Token{AccessToken: token}
	transport.Token = t
	c := Client{AccessToken: token, transport: transport}
	return &c
}

func (c *Client) API(method, url string, data interface{}) (respBody []byte, err error) {
	jsonPattern, _ := regexp.Compile(`\.json$`)
	if !jsonPattern.Match([]byte(url)) {
		url += ".json"
	}
	url = "https://api.bufferapp.com/1/" + url

	jsonBody, _ := json.Marshal(data)
	b := bytes.NewBuffer(jsonBody)

	var resp *http.Response
	switch method {
	case "get":
		resp, err = c.transport.Client().Get(url)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()

	case "post":
		resp, err := c.transport.Client().Post(url, "application/json", b)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()

	default:
		return nil, errors.New("Not a valid request type")
	}

	if resp.StatusCode >= 400 {
		return nil, errors.New(resp.Status)
	}

	respBody, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if len(respBody) < 2 {
		return nil, errors.New("Malformed JSON response")
	}
	return respBody, nil
}

func (c *Client) Get(url string) (resp []byte, err error) {
	return c.API("get", url, nil)
}

func (c *Client) Post(url string, params interface{}) (resp []byte, err error) {
	return c.API("post", url, params)
}

func (c *Client) Profiles() (profiles *[]Profile, err error) {
	body, err := c.Get("profiles.json")
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(body, profiles)
	if err != nil {
		return nil, err
	}
	return
}

func (c *Client) Update(update *NewUpdate) (resp *UpdateResponse, err error) {
	respBody, err := c.Post("updates/create.json", update)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(respBody, resp)
	if err != nil {
		return nil, err
	} else if !resp.Success {
		return nil, errors.New(respBody)
	}
	return resp, nil
}
<|endoftext|>


<commit_before>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: Sun Jul  1 07:08:42 PDT 2018

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "1c2fafa43b85e3513bcb055cce4802252c3c17ca"
<commit_msg>Tagging Sparta commit<commit_after>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: Sun Jul  1 07:09:27 PDT 2018

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "b76b71d7b0f9a79f9359f85b2740050eb743fcc3"
<|endoftext|>


<commit_before>package cache

import (
	"container/list"
    "sync"
)

type Service struct {
	//RecordType string
	Value      string
	TTL        int
	Aliases    string
}

type LRUCache struct {
	mu sync.Mutex

	// list & table of *entry objects
	list  *list.List
	table map[string]*list.Element

	// Our current size, in bytes. Obviously a gross simplification and low-grade
	// approximation.
	size uint64

	// How many bytes we are limiting the cache to.
	capacity uint64
}

type Record struct {
	size int
	list *list.List
}
type RecordCache struct {
	mu sync.Mutex

	list *list.List
	table map[string]*Record
}

//name to type
func NewLRUCache(capacity uint64) *LRUCache {
	return &LRUCache{
		list:     list.New(),
		table:    make(map[string]*list.Element),
		capacity: capacity,
	}
}

//type to record
func NewRecordCache() *RecordCache{
	return &RecordCache{
		list:	list.New(),
		table:  make(map[string]&Record(size:0,list:[])),
	}
}<commit_msg>update<commit_after>package cache

import (
	"container/list"
    "sync"
)

type Service struct {
	//RecordType string
	Value      string
	TTL        int
	Aliases    string
}

type LRUCache struct {
	mu sync.Mutex

	// list & table of *entry objects
	list  *list.List
	table map[string]*list.Element

	// Our current size, in bytes. Obviously a gross simplification and low-grade
	// approximation.
	size uint64

	// How many bytes we are limiting the cache to.
	capacity uint64
}

type Record struct {
	size int
	list *list.List
}
type RecordCache struct {
	mu sync.Mutex

	list *list.List
	table map[string]*Record
}

//name to type
func NewLRUCache(capacity uint64) *LRUCache {
	return &LRUCache{
		list:     list.New(),
		table:    make(map[string]*list.Element),
		capacity: capacity,
	}
}

//type to record
func NewRecordCache() *RecordCache{
	return &RecordCache{
		list:	list.New(),
		table:  make(map[string]*Record(size:0,list:[])),
	}
}<|endoftext|>


<commit_before>package sandra

import (
	"errors"
	"fmt"
	"time"

	"github.com/gocql/gocql"
	"github.com/mailgun/log"
)

type Cassandra interface {
	ExecuteQuery(string, ...interface{}) error
	ExecuteBatch(gocql.BatchType, []string, [][]interface{}) error
	ExecuteUnloggedBatch([]string, [][]interface{}) error
	ScanQuery(string, []interface{}, ...interface{}) error
	ScanCASQuery(string, []interface{}, ...interface{}) (bool, error)
	IterQuery(string, []interface{}, ...interface{}) func() (int, bool, error)
	Close() error
}

type cassandra struct {
	session *gocql.Session
	config  CassandraConfig
}

// CassandraConfig is a json and yaml friendly configuration struct
type CassandraConfig struct {
	Nodes       []string // addresses for the initial connections
	Timeout     string   // connection timeout (default: 600ms)
	Keyspace    string   // initial keyspace (optional)
	NumConns    int      // number of connections per host (default: 2)
	NumStreams  int      // number of streams per connection (default: 128)
	Consistency string   // consistency to use, default quorum
	Port        int      // port to connect to, default: 9042
	KeepAlive   string   // The keepalive period to use default: 0

	// TestMode affects whether a keyspace creation will be attempted on Cassandra initialization.
	TestMode bool `config:"optional"`
}

func (c CassandraConfig) String() string {
	return fmt.Sprintf("CassandraConfig(Nodes=%v, Keyspace=%v, Consistency=%v, TestMode=%v)",
		c.Nodes, c.Keyspace, c.Consistency, c.TestMode)
}

var NotFound = errors.New("Not found")

func NewCassandra(config CassandraConfig) (Cassandra, error) {
	log.Infof("Connecting to Cassandra with config: %v", config)

	cluster, err := setDefaults(config)
	if err != nil {
		return nil, err
	}

	// in test mode, create a keyspace if necessary
	if config.TestMode == true {
		session, err := cluster.CreateSession()
		if err != nil {
			return nil, err
		}

		query := session.Query(
			fmt.Sprintf(
				`create keyspace if not exists %v
                   with replication = {'class': 'SimpleStrategy', 'replication_factor': 1}`,
				config.Keyspace))

		if err := query.Exec(); err != nil {
			log.Errorf("Error creating keyspace: %v", err)
		}

		session.Close()
	}

	// switch the keyspace
	cluster.Keyspace = config.Keyspace

	session, err := cluster.CreateSession()
	if err != nil {
		return nil, err
	}

	return &cassandra{session: session}, nil
}

func (c *cassandra) Close() error {
	c.session.Close()
	return nil
}

func (c *cassandra) ExecuteQuery(queryString string, queryParams ...interface{}) error {
	return c.session.Query(queryString, queryParams...).Exec()
}

func (c *cassandra) ExecuteBatch(batchType gocql.BatchType, queries []string, params [][]interface{}) error {
	count := len(queries)

	// quick sanity check
	if count != len(params) {
		return errors.New("Amount of queries and params does not match")
	}

	batch := gocql.NewBatch(batchType)
	for idx := 0; idx < count; idx++ {
		batch.Query(queries[idx], params[idx]...)
	}

	return c.session.ExecuteBatch(batch)
}

func (c *cassandra) ExecuteUnloggedBatch(queries []string, params [][]interface{}) error {
	return c.ExecuteBatch(gocql.UnloggedBatch, queries, params)
}

func (c *cassandra) ScanQuery(queryString string, queryParams []interface{}, outParams ...interface{}) error {
	if err := c.session.Query(queryString, queryParams...).Scan(outParams...); err != nil {
		if err == gocql.ErrNotFound {
			return NotFound
		}
		return err
	}
	return nil
}

// Execute a lightweight transaction (an UPDATE or INSERT statement containing an IF clause)
func (c *cassandra) ScanCASQuery(queryString string, queryParams []interface{}, outParams ...interface{}) (bool, error) {
	return c.session.Query(queryString, queryParams...).ScanCAS(outParams...)
}

func (c *cassandra) IterQuery(queryString string, queryParams []interface{}, outParams ...interface{}) func() (int, bool, error) {
	iter := c.session.Query(queryString, queryParams...).Iter()
	idx := -1
	return func() (int, bool, error) {
		idx++
		if iter.Scan(outParams...) {
			return idx, true, nil
		}
		if err := iter.Close(); err != nil {
			return idx, true, err
		}
		return idx, false, nil
	}
}

// Return appropriate gocql.Consistency based on the provided consistency level name.
func translateConsistency(consistencyName string) (gocql.Consistency, error) {
	for index, name := range gocql.ConsistencyNames {
		if name == consistencyName {
			return gocql.Consistency(index), nil
		}
	}
	return gocql.One, fmt.Errorf("unknown consistency: %v", consistencyName)
}

func translateDuration(k string, df time.Duration) (time.Duration, error) {
	if k == "" {
		return df, nil
	}
	return time.ParseDuration(k)
}

func setDefaults(cfg CassandraConfig) (*gocql.ClusterConfig, error) {
	// convert consistency name into appropriate gocql.Consistency value
	consistency, err := translateConsistency(cfg.Consistency)
	if err != nil {
		return nil, err
	}

	keepAlive, err := translateDuration(cfg.KeepAlive, 0)
	if err != nil {
		return nil, err
	}
	timeout, err := translateDuration(cfg.Timeout, 600*time.Millisecond)
	if err != nil {
		return nil, err
	}
	if cfg.Port == 0 {
		cfg.Port = 9042
	}
	if cfg.NumConns == 0 {
		cfg.NumConns = 2
	}
	if cfg.NumStreams == 0 {
		cfg.NumStreams = 128
	}
	cluster := gocql.NewCluster(cfg.Nodes...)
	cluster.ProtoVersion = 2
	cluster.CQLVersion = "3.0.0"
	cluster.Timeout = timeout
	cluster.NumConns = cfg.NumConns
	cluster.Consistency = consistency
	cluster.SocketKeepalive = keepAlive
	return cluster, nil
}
<commit_msg>Add num streams<commit_after>package sandra

import (
	"errors"
	"fmt"
	"time"

	"github.com/gocql/gocql"
	"github.com/mailgun/log"
)

type Cassandra interface {
	ExecuteQuery(string, ...interface{}) error
	ExecuteBatch(gocql.BatchType, []string, [][]interface{}) error
	ExecuteUnloggedBatch([]string, [][]interface{}) error
	ScanQuery(string, []interface{}, ...interface{}) error
	ScanCASQuery(string, []interface{}, ...interface{}) (bool, error)
	IterQuery(string, []interface{}, ...interface{}) func() (int, bool, error)
	Close() error
}

type cassandra struct {
	session *gocql.Session
	config  CassandraConfig
}

// CassandraConfig is a json and yaml friendly configuration struct
type CassandraConfig struct {
	Nodes       []string // addresses for the initial connections
	Timeout     string   // connection timeout (default: 600ms)
	Keyspace    string   // initial keyspace (optional)
	NumConns    int      // number of connections per host (default: 2)
	NumStreams  int      // number of streams per connection (default: 128)
	Consistency string   // consistency to use, default quorum
	Port        int      // port to connect to, default: 9042
	KeepAlive   string   // The keepalive period to use default: 0

	// TestMode affects whether a keyspace creation will be attempted on Cassandra initialization.
	TestMode bool `config:"optional"`
}

func (c CassandraConfig) String() string {
	return fmt.Sprintf("CassandraConfig(Nodes=%v, Keyspace=%v, Consistency=%v, TestMode=%v)",
		c.Nodes, c.Keyspace, c.Consistency, c.TestMode)
}

var NotFound = errors.New("Not found")

func NewCassandra(config CassandraConfig) (Cassandra, error) {
	log.Infof("Connecting to Cassandra with config: %v", config)

	cluster, err := setDefaults(config)
	if err != nil {
		return nil, err
	}

	// in test mode, create a keyspace if necessary
	if config.TestMode == true {
		session, err := cluster.CreateSession()
		if err != nil {
			return nil, err
		}

		query := session.Query(
			fmt.Sprintf(
				`create keyspace if not exists %v
                   with replication = {'class': 'SimpleStrategy', 'replication_factor': 1}`,
				config.Keyspace))

		if err := query.Exec(); err != nil {
			log.Errorf("Error creating keyspace: %v", err)
		}

		session.Close()
	}

	// switch the keyspace
	cluster.Keyspace = config.Keyspace

	session, err := cluster.CreateSession()
	if err != nil {
		return nil, err
	}

	return &cassandra{session: session}, nil
}

func (c *cassandra) Close() error {
	c.session.Close()
	return nil
}

func (c *cassandra) ExecuteQuery(queryString string, queryParams ...interface{}) error {
	return c.session.Query(queryString, queryParams...).Exec()
}

func (c *cassandra) ExecuteBatch(batchType gocql.BatchType, queries []string, params [][]interface{}) error {
	count := len(queries)

	// quick sanity check
	if count != len(params) {
		return errors.New("Amount of queries and params does not match")
	}

	batch := gocql.NewBatch(batchType)
	for idx := 0; idx < count; idx++ {
		batch.Query(queries[idx], params[idx]...)
	}

	return c.session.ExecuteBatch(batch)
}

func (c *cassandra) ExecuteUnloggedBatch(queries []string, params [][]interface{}) error {
	return c.ExecuteBatch(gocql.UnloggedBatch, queries, params)
}

func (c *cassandra) ScanQuery(queryString string, queryParams []interface{}, outParams ...interface{}) error {
	if err := c.session.Query(queryString, queryParams...).Scan(outParams...); err != nil {
		if err == gocql.ErrNotFound {
			return NotFound
		}
		return err
	}
	return nil
}

// Execute a lightweight transaction (an UPDATE or INSERT statement containing an IF clause)
func (c *cassandra) ScanCASQuery(queryString string, queryParams []interface{}, outParams ...interface{}) (bool, error) {
	return c.session.Query(queryString, queryParams...).ScanCAS(outParams...)
}

func (c *cassandra) IterQuery(queryString string, queryParams []interface{}, outParams ...interface{}) func() (int, bool, error) {
	iter := c.session.Query(queryString, queryParams...).Iter()
	idx := -1
	return func() (int, bool, error) {
		idx++
		if iter.Scan(outParams...) {
			return idx, true, nil
		}
		if err := iter.Close(); err != nil {
			return idx, true, err
		}
		return idx, false, nil
	}
}

// Return appropriate gocql.Consistency based on the provided consistency level name.
func translateConsistency(consistencyName string) (gocql.Consistency, error) {
	for index, name := range gocql.ConsistencyNames {
		if name == consistencyName {
			return gocql.Consistency(index), nil
		}
	}
	return gocql.One, fmt.Errorf("unknown consistency: %v", consistencyName)
}

func translateDuration(k string, df time.Duration) (time.Duration, error) {
	if k == "" {
		return df, nil
	}
	return time.ParseDuration(k)
}

func setDefaults(cfg CassandraConfig) (*gocql.ClusterConfig, error) {
	// convert consistency name into appropriate gocql.Consistency value
	consistency, err := translateConsistency(cfg.Consistency)
	if err != nil {
		return nil, err
	}

	keepAlive, err := translateDuration(cfg.KeepAlive, 0)
	if err != nil {
		return nil, err
	}
	timeout, err := translateDuration(cfg.Timeout, 600*time.Millisecond)
	if err != nil {
		return nil, err
	}
	if cfg.Port == 0 {
		cfg.Port = 9042
	}
	if cfg.NumConns == 0 {
		cfg.NumConns = 2
	}
	if cfg.NumStreams == 0 {
		cfg.NumStreams = 128
	}
	cluster := gocql.NewCluster(cfg.Nodes...)
	cluster.ProtoVersion = 2
	cluster.CQLVersion = "3.0.0"
	cluster.Timeout = timeout
	cluster.NumConns = cfg.NumConns
	cluster.NumStreams = cfg.NumStreams
	cluster.Consistency = consistency
	cluster.SocketKeepalive = keepAlive
	return cluster, nil
}
<|endoftext|>