<commit_before>package handlers

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"

	"github.com/docker/distribution"
	ctxu "github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/registry/api/v2"
	"github.com/gorilla/handlers"
)

// layerUploadDispatcher constructs and returns the layer upload handler for
// the given request context.
func layerUploadDispatcher(ctx *Context, r *http.Request) http.Handler {
	luh := &layerUploadHandler{
		Context: ctx,
		UUID:    getUploadUUID(ctx),
	}

	handler := http.Handler(handlers.MethodHandler{
		"POST": http.HandlerFunc(luh.StartLayerUpload),
		"GET":  http.HandlerFunc(luh.GetUploadStatus),
		"HEAD": http.HandlerFunc(luh.GetUploadStatus),
		// TODO(stevvooe): Must implement patch support.
		// "PATCH":    http.HandlerFunc(luh.PutLayerChunk),
		"PUT":    http.HandlerFunc(luh.PutLayerUploadComplete),
		"DELETE": http.HandlerFunc(luh.CancelLayerUpload),
	})

	if luh.UUID != "" {
		state, err := hmacKey(ctx.Config.HTTP.Secret).unpackUploadState(r.FormValue("_state"))
		if err != nil {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("error resolving upload: %v", err)
				w.WriteHeader(http.StatusBadRequest)
				luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
			})
		}
		luh.State = state

		if state.Name != ctx.Repository.Name() {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched repository name in upload state: %q != %q", state.Name, luh.Repository.Name())
				w.WriteHeader(http.StatusBadRequest)
				luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
			})
		}

		if state.UUID != luh.UUID {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched uuid in upload state: %q != %q", state.UUID, luh.UUID)
				w.WriteHeader(http.StatusBadRequest)
				luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
			})
		}

		layers := ctx.Repository.Layers()
		upload, err := layers.Resume(luh.UUID)
		if err != nil {
			ctxu.GetLogger(ctx).Errorf("error resolving upload: %v", err)
			if err == distribution.ErrLayerUploadUnknown {
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.WriteHeader(http.StatusNotFound)
					luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown, err)
				})
			}

			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusInternalServerError)
				luh.Errors.Push(v2.ErrorCodeUnknown, err)
			})
		}
		luh.Upload = upload

		if state.Offset > 0 {
			// Seek the layer upload to the correct spot if it's non-zero.
			// These error conditions should be rare and demonstrate really
			// problems. We basically cancel the upload and tell the client to
			// start over.
			if nn, err := upload.Seek(luh.State.Offset, os.SEEK_SET); err != nil {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("error seeking layer upload: %v", err)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.WriteHeader(http.StatusBadRequest)
					luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
					upload.Cancel()
				})
			} else if nn != luh.State.Offset {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("seek to wrong offest: %d != %d", nn, luh.State.Offset)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.WriteHeader(http.StatusBadRequest)
					luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
					upload.Cancel()
				})
			}
		}

		handler = closeResources(handler, luh.Upload)
	}

	return handler
}

// layerUploadHandler handles the http layer upload process.
type layerUploadHandler struct {
	*Context

	// UUID identifies the upload instance for the current request.
	UUID string

	Upload distribution.LayerUpload

	State layerUploadState
}

// StartLayerUpload begins the layer upload process and allocates a server-
// side upload session.
func (luh *layerUploadHandler) StartLayerUpload(w http.ResponseWriter, r *http.Request) {
	layers := luh.Repository.Layers()
	upload, err := layers.Upload()
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError) // Error conditions here?
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	luh.Upload = upload
	defer luh.Upload.Close()

	if err := luh.layerUploadResponse(w, r); err != nil {
		w.WriteHeader(http.StatusInternalServerError) // Error conditions here?
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	w.Header().Set("Docker-Upload-UUID", luh.Upload.UUID())
	w.WriteHeader(http.StatusAccepted)
}

// GetUploadStatus returns the status of a given upload, identified by uuid.
func (luh *layerUploadHandler) GetUploadStatus(w http.ResponseWriter, r *http.Request) {
	if luh.Upload == nil {
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown)
		return
	}

	if err := luh.layerUploadResponse(w, r); err != nil {
		w.WriteHeader(http.StatusInternalServerError) // Error conditions here?
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	w.Header().Set("Docker-Upload-UUID", luh.UUID)
	w.WriteHeader(http.StatusNoContent)
}

// PutLayerUploadComplete takes the final request of a layer upload. The final
// chunk may include all the layer data, the final chunk of layer data or no
// layer data. Any data provided is received and verified. If successful, the
// layer is linked into the blob store and 201 Created is returned with the
// canonical url of the layer.
func (luh *layerUploadHandler) PutLayerUploadComplete(w http.ResponseWriter, r *http.Request) {
	if luh.Upload == nil {
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown)
		return
	}

	dgstStr := r.FormValue("digest") // TODO(stevvooe): Support multiple digest parameters!

	if dgstStr == "" {
		// no digest? return error, but allow retry.
		w.WriteHeader(http.StatusBadRequest)
		luh.Errors.Push(v2.ErrorCodeDigestInvalid, "digest missing")
		return
	}

	dgst, err := digest.ParseDigest(dgstStr)
	if err != nil {
		// no digest? return error, but allow retry.
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeDigestInvalid, "digest parsing failed")
		return
	}

	// TODO(stevvooe): Check the incoming range header here, per the
	// specification. LayerUpload should be seeked (sought?) to that position.

	// TODO(stevvooe): Consider checking the error on this copy.
	// Theoretically, problems should be detected during verification but we
	// may miss a root cause.

	// Read in the final chunk, if any.
	if _, err := io.Copy(luh.Upload, r.Body); err != nil {
		ctxu.GetLogger(luh).Errorf("unknown error copying into upload: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
	}

	layer, err := luh.Upload.Finish(dgst)
	if err != nil {
		switch err := err.(type) {
		case distribution.ErrLayerInvalidDigest:
			w.WriteHeader(http.StatusBadRequest)
			luh.Errors.Push(v2.ErrorCodeDigestInvalid, err)
		default:
			ctxu.GetLogger(luh).Errorf("unknown error completing upload: %#v", err)
			w.WriteHeader(http.StatusInternalServerError)
			luh.Errors.Push(v2.ErrorCodeUnknown, err)
		}

		// Clean up the backend layer data if there was an error.
		if err := luh.Upload.Cancel(); err != nil {
			// If the cleanup fails, all we can do is observe and report.
			ctxu.GetLogger(luh).Errorf("error canceling upload after error: %v", err)
		}

		return
	}

	// Build our canonical layer url
	layerURL, err := luh.urlBuilder.BuildBlobURL(luh.Repository.Name(), layer.Digest())
	if err != nil {
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Header().Set("Location", layerURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Docker-Content-Digest", layer.Digest().String())
	w.WriteHeader(http.StatusCreated)
}

// CancelLayerUpload cancels an in-progress upload of a layer.
func (luh *layerUploadHandler) CancelLayerUpload(w http.ResponseWriter, r *http.Request) {
	if luh.Upload == nil {
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown)
		return
	}

	w.Header().Set("Docker-Upload-UUID", luh.UUID)
	if err := luh.Upload.Cancel(); err != nil {
		ctxu.GetLogger(luh).Errorf("error encountered canceling upload: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		luh.Errors.PushErr(err)
	}

	w.WriteHeader(http.StatusNoContent)
}

// layerUploadResponse provides a standard request for uploading layers and
// chunk responses. This sets the correct headers but the response status is
// left to the caller.
func (luh *layerUploadHandler) layerUploadResponse(w http.ResponseWriter, r *http.Request) error {

	offset, err := luh.Upload.Seek(0, os.SEEK_CUR)
	if err != nil {
		ctxu.GetLogger(luh).Errorf("unable get current offset of layer upload: %v", err)
		return err
	}

	// TODO(stevvooe): Need a better way to manage the upload state automatically.
	luh.State.Name = luh.Repository.Name()
	luh.State.UUID = luh.Upload.UUID()
	luh.State.Offset = offset
	luh.State.StartedAt = luh.Upload.StartedAt()

	token, err := hmacKey(luh.Config.HTTP.Secret).packUploadState(luh.State)
	if err != nil {
		ctxu.GetLogger(luh).Infof("error building upload state token: %s", err)
		return err
	}

	uploadURL, err := luh.urlBuilder.BuildBlobUploadChunkURL(
		luh.Repository.Name(), luh.Upload.UUID(),
		url.Values{
			"_state": []string{token},
		})
	if err != nil {
		ctxu.GetLogger(luh).Infof("error building upload url: %s", err)
		return err
	}

	w.Header().Set("Docker-Upload-UUID", luh.UUID)
	w.Header().Set("Location", uploadURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Range", fmt.Sprintf("0-%d", luh.State.Offset))

	return nil
}
<commit_msg>Return after error in handler<commit_after>package handlers

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"

	"github.com/docker/distribution"
	ctxu "github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/registry/api/v2"
	"github.com/gorilla/handlers"
)

// layerUploadDispatcher constructs and returns the layer upload handler for
// the given request context.
func layerUploadDispatcher(ctx *Context, r *http.Request) http.Handler {
	luh := &layerUploadHandler{
		Context: ctx,
		UUID:    getUploadUUID(ctx),
	}

	handler := http.Handler(handlers.MethodHandler{
		"POST": http.HandlerFunc(luh.StartLayerUpload),
		"GET":  http.HandlerFunc(luh.GetUploadStatus),
		"HEAD": http.HandlerFunc(luh.GetUploadStatus),
		// TODO(stevvooe): Must implement patch support.
		// "PATCH":    http.HandlerFunc(luh.PutLayerChunk),
		"PUT":    http.HandlerFunc(luh.PutLayerUploadComplete),
		"DELETE": http.HandlerFunc(luh.CancelLayerUpload),
	})

	if luh.UUID != "" {
		state, err := hmacKey(ctx.Config.HTTP.Secret).unpackUploadState(r.FormValue("_state"))
		if err != nil {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("error resolving upload: %v", err)
				w.WriteHeader(http.StatusBadRequest)
				luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
			})
		}
		luh.State = state

		if state.Name != ctx.Repository.Name() {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched repository name in upload state: %q != %q", state.Name, luh.Repository.Name())
				w.WriteHeader(http.StatusBadRequest)
				luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
			})
		}

		if state.UUID != luh.UUID {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched uuid in upload state: %q != %q", state.UUID, luh.UUID)
				w.WriteHeader(http.StatusBadRequest)
				luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
			})
		}

		layers := ctx.Repository.Layers()
		upload, err := layers.Resume(luh.UUID)
		if err != nil {
			ctxu.GetLogger(ctx).Errorf("error resolving upload: %v", err)
			if err == distribution.ErrLayerUploadUnknown {
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.WriteHeader(http.StatusNotFound)
					luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown, err)
				})
			}

			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusInternalServerError)
				luh.Errors.Push(v2.ErrorCodeUnknown, err)
			})
		}
		luh.Upload = upload

		if state.Offset > 0 {
			// Seek the layer upload to the correct spot if it's non-zero.
			// These error conditions should be rare and demonstrate really
			// problems. We basically cancel the upload and tell the client to
			// start over.
			if nn, err := upload.Seek(luh.State.Offset, os.SEEK_SET); err != nil {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("error seeking layer upload: %v", err)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.WriteHeader(http.StatusBadRequest)
					luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
					upload.Cancel()
				})
			} else if nn != luh.State.Offset {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("seek to wrong offest: %d != %d", nn, luh.State.Offset)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.WriteHeader(http.StatusBadRequest)
					luh.Errors.Push(v2.ErrorCodeBlobUploadInvalid, err)
					upload.Cancel()
				})
			}
		}

		handler = closeResources(handler, luh.Upload)
	}

	return handler
}

// layerUploadHandler handles the http layer upload process.
type layerUploadHandler struct {
	*Context

	// UUID identifies the upload instance for the current request.
	UUID string

	Upload distribution.LayerUpload

	State layerUploadState
}

// StartLayerUpload begins the layer upload process and allocates a server-
// side upload session.
func (luh *layerUploadHandler) StartLayerUpload(w http.ResponseWriter, r *http.Request) {
	layers := luh.Repository.Layers()
	upload, err := layers.Upload()
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError) // Error conditions here?
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	luh.Upload = upload
	defer luh.Upload.Close()

	if err := luh.layerUploadResponse(w, r); err != nil {
		w.WriteHeader(http.StatusInternalServerError) // Error conditions here?
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	w.Header().Set("Docker-Upload-UUID", luh.Upload.UUID())
	w.WriteHeader(http.StatusAccepted)
}

// GetUploadStatus returns the status of a given upload, identified by uuid.
func (luh *layerUploadHandler) GetUploadStatus(w http.ResponseWriter, r *http.Request) {
	if luh.Upload == nil {
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown)
		return
	}

	if err := luh.layerUploadResponse(w, r); err != nil {
		w.WriteHeader(http.StatusInternalServerError) // Error conditions here?
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	w.Header().Set("Docker-Upload-UUID", luh.UUID)
	w.WriteHeader(http.StatusNoContent)
}

// PutLayerUploadComplete takes the final request of a layer upload. The final
// chunk may include all the layer data, the final chunk of layer data or no
// layer data. Any data provided is received and verified. If successful, the
// layer is linked into the blob store and 201 Created is returned with the
// canonical url of the layer.
func (luh *layerUploadHandler) PutLayerUploadComplete(w http.ResponseWriter, r *http.Request) {
	if luh.Upload == nil {
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown)
		return
	}

	dgstStr := r.FormValue("digest") // TODO(stevvooe): Support multiple digest parameters!

	if dgstStr == "" {
		// no digest? return error, but allow retry.
		w.WriteHeader(http.StatusBadRequest)
		luh.Errors.Push(v2.ErrorCodeDigestInvalid, "digest missing")
		return
	}

	dgst, err := digest.ParseDigest(dgstStr)
	if err != nil {
		// no digest? return error, but allow retry.
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeDigestInvalid, "digest parsing failed")
		return
	}

	// TODO(stevvooe): Check the incoming range header here, per the
	// specification. LayerUpload should be seeked (sought?) to that position.

	// TODO(stevvooe): Consider checking the error on this copy.
	// Theoretically, problems should be detected during verification but we
	// may miss a root cause.

	// Read in the final chunk, if any.
	if _, err := io.Copy(luh.Upload, r.Body); err != nil {
		ctxu.GetLogger(luh).Errorf("unknown error copying into upload: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		return
	}

	layer, err := luh.Upload.Finish(dgst)
	if err != nil {
		switch err := err.(type) {
		case distribution.ErrLayerInvalidDigest:
			w.WriteHeader(http.StatusBadRequest)
			luh.Errors.Push(v2.ErrorCodeDigestInvalid, err)
		default:
			ctxu.GetLogger(luh).Errorf("unknown error completing upload: %#v", err)
			w.WriteHeader(http.StatusInternalServerError)
			luh.Errors.Push(v2.ErrorCodeUnknown, err)
		}

		// Clean up the backend layer data if there was an error.
		if err := luh.Upload.Cancel(); err != nil {
			// If the cleanup fails, all we can do is observe and report.
			ctxu.GetLogger(luh).Errorf("error canceling upload after error: %v", err)
		}

		return
	}

	// Build our canonical layer url
	layerURL, err := luh.urlBuilder.BuildBlobURL(luh.Repository.Name(), layer.Digest())
	if err != nil {
		luh.Errors.Push(v2.ErrorCodeUnknown, err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Header().Set("Location", layerURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Docker-Content-Digest", layer.Digest().String())
	w.WriteHeader(http.StatusCreated)
}

// CancelLayerUpload cancels an in-progress upload of a layer.
func (luh *layerUploadHandler) CancelLayerUpload(w http.ResponseWriter, r *http.Request) {
	if luh.Upload == nil {
		w.WriteHeader(http.StatusNotFound)
		luh.Errors.Push(v2.ErrorCodeBlobUploadUnknown)
		return
	}

	w.Header().Set("Docker-Upload-UUID", luh.UUID)
	if err := luh.Upload.Cancel(); err != nil {
		ctxu.GetLogger(luh).Errorf("error encountered canceling upload: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		luh.Errors.PushErr(err)
	}

	w.WriteHeader(http.StatusNoContent)
}

// layerUploadResponse provides a standard request for uploading layers and
// chunk responses. This sets the correct headers but the response status is
// left to the caller.
func (luh *layerUploadHandler) layerUploadResponse(w http.ResponseWriter, r *http.Request) error {

	offset, err := luh.Upload.Seek(0, os.SEEK_CUR)
	if err != nil {
		ctxu.GetLogger(luh).Errorf("unable get current offset of layer upload: %v", err)
		return err
	}

	// TODO(stevvooe): Need a better way to manage the upload state automatically.
	luh.State.Name = luh.Repository.Name()
	luh.State.UUID = luh.Upload.UUID()
	luh.State.Offset = offset
	luh.State.StartedAt = luh.Upload.StartedAt()

	token, err := hmacKey(luh.Config.HTTP.Secret).packUploadState(luh.State)
	if err != nil {
		ctxu.GetLogger(luh).Infof("error building upload state token: %s", err)
		return err
	}

	uploadURL, err := luh.urlBuilder.BuildBlobUploadChunkURL(
		luh.Repository.Name(), luh.Upload.UUID(),
		url.Values{
			"_state": []string{token},
		})
	if err != nil {
		ctxu.GetLogger(luh).Infof("error building upload url: %s", err)
		return err
	}

	w.Header().Set("Docker-Upload-UUID", luh.UUID)
	w.Header().Set("Location", uploadURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Range", fmt.Sprintf("0-%d", luh.State.Offset))

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"net/http"
)

func main() {
	addr := flag.String("addr", "localhost:8081", "host:port")
	flag.Parse()
	log.Printf("Serving at %s\n", *addr)
	log.Fatal(http.ListenAndServe(*addr, http.FileServer(http.Dir("."))))
}
<commit_msg>Make it easier to launch browser from terminal<commit_after>package main

import (
	"flag"
	"log"
	"net/http"
)

func main() {
	addr := flag.String("addr", "localhost:8081", "host:port")
	flag.Parse()
	log.Printf("Serving at http://%s\n", *addr)
	log.Fatal(http.ListenAndServe(*addr, http.FileServer(http.Dir("."))))
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/transfer"

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

func resourceAwsTransferUser() *schema.Resource {

	return &schema.Resource{
		Create: resourceAwsTransferUserCreate,
		Read:   resourceAwsTransferUserRead,
		Update: resourceAwsTransferUserUpdate,
		Delete: resourceAwsTransferUserDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},
		Schema: map[string]*schema.Schema{

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

			"home_directory": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validation.StringLenBetween(0, 1024),
			},

			"policy": {
				Type:             schema.TypeString,
				Optional:         true,
				ValidateFunc:     validateIAMPolicyJson,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},

			"role": {
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validateArn,
			},

			"server_id": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateTransferServerID,
			},

			"tags": tagsSchema(),

			"user_name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateTransferUserName,
			},
		},
	}
}

func resourceAwsTransferUserCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	userName := d.Get("user_name").(string)
	serverID := d.Get("server_id").(string)

	createOpts := &transfer.CreateUserInput{
		ServerId: aws.String(serverID),
		UserName: aws.String(userName),
		Role:     aws.String(d.Get("role").(string)),
	}

	if attr, ok := d.GetOk("home_directory"); ok {
		createOpts.HomeDirectory = aws.String(attr.(string))
	}

	if attr, ok := d.GetOk("policy"); ok {
		createOpts.Policy = aws.String(attr.(string))
	}

	if attr, ok := d.GetOk("tags"); ok {
		createOpts.Tags = tagsFromMapTransfer(attr.(map[string]interface{}))
	}

	log.Printf("[DEBUG] Create Transfer User Option: %#v", createOpts)

	_, err := conn.CreateUser(createOpts)
	if err != nil {
		return fmt.Errorf("error creating Transfer User: %s", err)
	}

	d.SetId(fmt.Sprintf("%s/%s", serverID, userName))

	return resourceAwsTransferUserRead(d, meta)
}

func resourceAwsTransferUserRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	serverID, userName, err := decodeTransferUserId(d.Id())
	if err != nil {
		return fmt.Errorf("error parsing Transfer User ID: %s", err)
	}

	descOpts := &transfer.DescribeUserInput{
		UserName: aws.String(userName),
		ServerId: aws.String(serverID),
	}

	log.Printf("[DEBUG] Describe Transfer User Option: %#v", descOpts)

	resp, err := conn.DescribeUser(descOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			log.Printf("[WARN] Transfer User (%s) for Server (%s) not found, removing from state", userName, serverID)
			d.SetId("")
			return nil
		}
		return fmt.Errorf("error reading Transfer User (%s): %s", d.Id(), err)
	}

	d.Set("server_id", resp.ServerId)
	d.Set("user_name", resp.User.UserName)
	d.Set("arn", resp.User.Arn)
	d.Set("home_directory", resp.User.HomeDirectory)
	d.Set("policy", resp.User.Policy)
	d.Set("role", resp.User.Role)

	if err := d.Set("tags", tagsToMapTransfer(resp.User.Tags)); err != nil {
		return fmt.Errorf("Error setting tags: %s", err)
	}
	return nil
}

func resourceAwsTransferUserUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	updateFlag := false
	serverID, userName, err := decodeTransferUserId(d.Id())
	if err != nil {
		return fmt.Errorf("error parsing Transfer User ID: %s", err)
	}

	updateOpts := &transfer.UpdateUserInput{
		UserName: aws.String(userName),
		ServerId: aws.String(serverID),
	}

	if d.HasChange("home_directory") {
		updateOpts.HomeDirectory = aws.String(d.Get("home_directory").(string))
		updateFlag = true
	}

	if d.HasChange("policy") {
		updateOpts.Policy = aws.String(d.Get("policy").(string))
		updateFlag = true
	}

	if d.HasChange("role") {
		updateOpts.Role = aws.String(d.Get("role").(string))
		updateFlag = true
	}

	if updateFlag {
		_, err := conn.UpdateUser(updateOpts)
		if err != nil {
			if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
				log.Printf("[WARN] Transfer User (%s) for Server (%s) not found, removing from state", userName, serverID)
				d.SetId("")
				return nil
			}
			return fmt.Errorf("error updating Transfer User (%s): %s", d.Id(), err)
		}
	}

	if err := setTagsTransfer(conn, d); err != nil {
		return fmt.Errorf("Error update tags: %s", err)
	}

	return resourceAwsTransferUserRead(d, meta)
}

func resourceAwsTransferUserDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	serverID, userName, err := decodeTransferUserId(d.Id())
	if err != nil {
		return fmt.Errorf("error parsing Transfer User ID: %s", err)
	}

	delOpts := &transfer.DeleteUserInput{
		UserName: aws.String(userName),
		ServerId: aws.String(serverID),
	}

	log.Printf("[DEBUG] Delete Transfer User Option: %#v", delOpts)

	_, err = conn.DeleteUser(delOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			return nil
		}
		return fmt.Errorf("error deleting Transfer User (%s) for Server(%s): %s", userName, serverID, err)
	}

	if err := waitForTransferUserDeletion(conn, serverID, userName); err != nil {
		return fmt.Errorf("error waiting for Transfer User (%s) for Server (%s): %s", userName, serverID, err)
	}

	return nil
}

func decodeTransferUserId(id string) (string, string, error) {
	idParts := strings.SplitN(id, "/", 2)
	if len(idParts) != 2 || idParts[0] == "" || idParts[1] == "" {
		return "", "", fmt.Errorf("unexpected format of ID (%s), expected SERVERID/USERNAME", id)
	}
	return idParts[0], idParts[1], nil
}

func waitForTransferUserDeletion(conn *transfer.Transfer, serverID, userName string) error {
	params := &transfer.DescribeUserInput{
		ServerId: aws.String(serverID),
		UserName: aws.String(userName),
	}

	return resource.Retry(10*time.Minute, func() *resource.RetryError {
		_, err := conn.DescribeUser(params)

		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			return nil
		}

		if err != nil {
			return resource.NonRetryableError(err)
		}

		return resource.RetryableError(fmt.Errorf("Transfer User (%s) for Server (%s) still exists", userName, serverID))
	})
}
<commit_msg>final retry when waiting for transfer user deletion<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/transfer"

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

func resourceAwsTransferUser() *schema.Resource {

	return &schema.Resource{
		Create: resourceAwsTransferUserCreate,
		Read:   resourceAwsTransferUserRead,
		Update: resourceAwsTransferUserUpdate,
		Delete: resourceAwsTransferUserDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},
		Schema: map[string]*schema.Schema{

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

			"home_directory": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validation.StringLenBetween(0, 1024),
			},

			"policy": {
				Type:             schema.TypeString,
				Optional:         true,
				ValidateFunc:     validateIAMPolicyJson,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},

			"role": {
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validateArn,
			},

			"server_id": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateTransferServerID,
			},

			"tags": tagsSchema(),

			"user_name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateTransferUserName,
			},
		},
	}
}

func resourceAwsTransferUserCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	userName := d.Get("user_name").(string)
	serverID := d.Get("server_id").(string)

	createOpts := &transfer.CreateUserInput{
		ServerId: aws.String(serverID),
		UserName: aws.String(userName),
		Role:     aws.String(d.Get("role").(string)),
	}

	if attr, ok := d.GetOk("home_directory"); ok {
		createOpts.HomeDirectory = aws.String(attr.(string))
	}

	if attr, ok := d.GetOk("policy"); ok {
		createOpts.Policy = aws.String(attr.(string))
	}

	if attr, ok := d.GetOk("tags"); ok {
		createOpts.Tags = tagsFromMapTransfer(attr.(map[string]interface{}))
	}

	log.Printf("[DEBUG] Create Transfer User Option: %#v", createOpts)

	_, err := conn.CreateUser(createOpts)
	if err != nil {
		return fmt.Errorf("error creating Transfer User: %s", err)
	}

	d.SetId(fmt.Sprintf("%s/%s", serverID, userName))

	return resourceAwsTransferUserRead(d, meta)
}

func resourceAwsTransferUserRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	serverID, userName, err := decodeTransferUserId(d.Id())
	if err != nil {
		return fmt.Errorf("error parsing Transfer User ID: %s", err)
	}

	descOpts := &transfer.DescribeUserInput{
		UserName: aws.String(userName),
		ServerId: aws.String(serverID),
	}

	log.Printf("[DEBUG] Describe Transfer User Option: %#v", descOpts)

	resp, err := conn.DescribeUser(descOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			log.Printf("[WARN] Transfer User (%s) for Server (%s) not found, removing from state", userName, serverID)
			d.SetId("")
			return nil
		}
		return fmt.Errorf("error reading Transfer User (%s): %s", d.Id(), err)
	}

	d.Set("server_id", resp.ServerId)
	d.Set("user_name", resp.User.UserName)
	d.Set("arn", resp.User.Arn)
	d.Set("home_directory", resp.User.HomeDirectory)
	d.Set("policy", resp.User.Policy)
	d.Set("role", resp.User.Role)

	if err := d.Set("tags", tagsToMapTransfer(resp.User.Tags)); err != nil {
		return fmt.Errorf("Error setting tags: %s", err)
	}
	return nil
}

func resourceAwsTransferUserUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	updateFlag := false
	serverID, userName, err := decodeTransferUserId(d.Id())
	if err != nil {
		return fmt.Errorf("error parsing Transfer User ID: %s", err)
	}

	updateOpts := &transfer.UpdateUserInput{
		UserName: aws.String(userName),
		ServerId: aws.String(serverID),
	}

	if d.HasChange("home_directory") {
		updateOpts.HomeDirectory = aws.String(d.Get("home_directory").(string))
		updateFlag = true
	}

	if d.HasChange("policy") {
		updateOpts.Policy = aws.String(d.Get("policy").(string))
		updateFlag = true
	}

	if d.HasChange("role") {
		updateOpts.Role = aws.String(d.Get("role").(string))
		updateFlag = true
	}

	if updateFlag {
		_, err := conn.UpdateUser(updateOpts)
		if err != nil {
			if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
				log.Printf("[WARN] Transfer User (%s) for Server (%s) not found, removing from state", userName, serverID)
				d.SetId("")
				return nil
			}
			return fmt.Errorf("error updating Transfer User (%s): %s", d.Id(), err)
		}
	}

	if err := setTagsTransfer(conn, d); err != nil {
		return fmt.Errorf("Error update tags: %s", err)
	}

	return resourceAwsTransferUserRead(d, meta)
}

func resourceAwsTransferUserDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	serverID, userName, err := decodeTransferUserId(d.Id())
	if err != nil {
		return fmt.Errorf("error parsing Transfer User ID: %s", err)
	}

	delOpts := &transfer.DeleteUserInput{
		UserName: aws.String(userName),
		ServerId: aws.String(serverID),
	}

	log.Printf("[DEBUG] Delete Transfer User Option: %#v", delOpts)

	_, err = conn.DeleteUser(delOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			return nil
		}
		return fmt.Errorf("error deleting Transfer User (%s) for Server(%s): %s", userName, serverID, err)
	}

	if err := waitForTransferUserDeletion(conn, serverID, userName); err != nil {
		return fmt.Errorf("error waiting for Transfer User (%s) for Server (%s): %s", userName, serverID, err)
	}

	return nil
}

func decodeTransferUserId(id string) (string, string, error) {
	idParts := strings.SplitN(id, "/", 2)
	if len(idParts) != 2 || idParts[0] == "" || idParts[1] == "" {
		return "", "", fmt.Errorf("unexpected format of ID (%s), expected SERVERID/USERNAME", id)
	}
	return idParts[0], idParts[1], nil
}

func waitForTransferUserDeletion(conn *transfer.Transfer, serverID, userName string) error {
	params := &transfer.DescribeUserInput{
		ServerId: aws.String(serverID),
		UserName: aws.String(userName),
	}

	err := resource.Retry(10*time.Minute, func() *resource.RetryError {
		_, err := conn.DescribeUser(params)

		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			return nil
		}

		if err != nil {
			return resource.NonRetryableError(err)
		}

		return resource.RetryableError(fmt.Errorf("Transfer User (%s) for Server (%s) still exists", userName, serverID))
	})

	if isResourceTimeoutError(err) {
		_, err = conn.DescribeUser(params)
	}
	if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
		return nil
	}
	if err != nil {
		return fmt.Errorf("Error decoding transfer user ID: %s", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

	conn := meta.(*AWSClient).ec2conn

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

	d.SetPartial("tags")

	return nil
}

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

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

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

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

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

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

		return resource.NonRetryableError(err)
	})
}

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

		ig := resp.InternetGateways[0]

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

		return ig, *ig.Attachments[0].State, nil
	}
}
<commit_msg>Issue #2174 Check that InternetGateway exists before returning from creation (#6105)<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

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

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

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

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

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

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

	resource.Retry(5*time.Minute, func() *resource.RetryError {
		igRaw, _, err := IGStateRefreshFunc(conn, d.Id())()
		if igRaw != nil {
			return nil
		}
		if err == nil {
			return resource.RetryableError(err)
		} else {
			return resource.NonRetryableError(err)
		}
	})

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

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

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

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

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

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

	return nil
}

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

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

	conn := meta.(*AWSClient).ec2conn

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

	d.SetPartial("tags")

	return nil
}

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

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

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

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

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

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

		return resource.NonRetryableError(err)
	})
}

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

		ig := resp.InternetGateways[0]

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

		return ig, *ig.Attachments[0].State, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"os/exec"
	"regexp"
)

type session struct {
	name        string
	ipVersion   int
	established int
}

var regex *regexp.Regexp

func main() {
	r, _ := regexp.Compile("^([^\\s]+)[\\s]+BGP[\\s]+([^\\s]+)[\\s]+([^\\s]+)[\\s]+([\\d]+)[\\s]+(.*?)[\\s]*$")
	regex = r

	fmt.Println("Starting bgp exporter")
	http.HandleFunc("/metrics", handleMetricsRequest)

	log.Fatal(http.ListenAndServe(":9200", nil))
}

func handleMetricsRequest(w http.ResponseWriter, r *http.Request) {
	sessions := getSessions()

	for _, s := range sessions {
		writeLineForSession(s, w)
	}
}

func writeLineForSession(s session, w io.Writer) {
	fmt.Fprintf(w, "bgp%d_session_up{name=\"%s\"} %d\n", s.ipVersion, s.name, s.established)
}

func getSessions() []session {
	birdSessions := getSessionsFromBird(4)
	bird6Sessions := getSessionsFromBird(6)

	return append(birdSessions, bird6Sessions...)
}

func getSessionsFromBird(ipVersion int) []session {
	client := "birdc"

	if ipVersion == 6 {
		client += "6"
	}

	output := getBirdOutput(client)
	return parseOutput(output, ipVersion)
}

func getBirdOutput(birdClient string) []byte {
	b, err := exec.Command(birdClient, "show", "protocols").Output()

	if err != nil {
		b = make([]byte, 0)
		log.Print(err)
	}

	return b
}

func parseOutput(data []byte, ipVersion int) []session {
	sessions := make([]session, 0)

	reader := bytes.NewReader(data)
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		line := scanner.Text()
		if session, res := parseLine(line, ipVersion); res == true {
			sessions = append(sessions, *session)
		}
	}

	return sessions
}

func parseLine(line string, ipVersion int) (*session, bool) {
	match := regex.FindStringSubmatch(line)

	if match != nil {
		session := session{name: match[1], ipVersion: ipVersion, established: parseState(match[5])}
		return &session, true
	} else {
		return nil, false
	}
}

func parseState(state string) int {
	if state == "Established" {
		return 1
	} else {
		return 0
	}
}
<commit_msg>added prefix counts<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"os/exec"
	"regexp"
	"strconv"
)

type session struct {
	name        string
	ipVersion   int
	established int
	imported    int64
	exported    int64
}

var protoRegex *regexp.Regexp
var routeRegex *regexp.Regexp

func main() {
	initRegexes()

	fmt.Println("Starting bgp exporter")
	http.HandleFunc("/metrics", handleMetricsRequest)

	log.Fatal(http.ListenAndServe(":9200", nil))
}

func initRegexes() {
	protoRegex, _ = regexp.Compile("^([^\\s]+)\\s+BGP\\s+([^\\s]+)\\s+([^\\s]+)\\s+([\\d]+)\\s+(.*?)\\s*$")
	routeRegex, _ = regexp.Compile("^\\s+Routes:\\s+(\\d+) imported, \\d+ filtered, (\\d+) exported")
}

func handleMetricsRequest(w http.ResponseWriter, r *http.Request) {
	sessions := getSessions()

	for _, s := range sessions {
		writeLineForSession(s, w)
	}
}

func writeLineForSession(s *session, w io.Writer) {
	fmt.Fprintf(w, "bgp%d_session_up{name=\"%s\"} %d\n", s.ipVersion, s.name, s.established)
	fmt.Fprintf(w, "bgp%d_session_prefix_count_import{name=\"%s\"} %d\n", s.ipVersion, s.name, s.imported)
	fmt.Fprintf(w, "bgp%d_session_prefix_count_export{name=\"%s\"} %d\n", s.ipVersion, s.name, s.exported)
}

func getSessions() []*session {
	birdSessions := getSessionsFromBird(4)
	bird6Sessions := getSessionsFromBird(6)

	return append(birdSessions, bird6Sessions...)
}

func getSessionsFromBird(ipVersion int) []*session {
	client := "birdc"

	if ipVersion == 6 {
		client += "6"
	}

	output := getBirdOutput(client)
	return parseOutput(output, ipVersion)
}

func getBirdOutput(birdClient string) []byte {
	b, err := exec.Command(birdClient, "show", "protocols", "all").Output()

	if err != nil {
		b = make([]byte, 0)
		log.Print(err)
	}

	return b
}

func parseOutput(data []byte, ipVersion int) []*session {
	sessions := make([]*session, 0)

	reader := bytes.NewReader(data)
	scanner := bufio.NewScanner(reader)
	var current *session = nil

	for scanner.Scan() {
		line := scanner.Text()
		if session, res := parseLineForSession(line, ipVersion); res == true {
			current = session
			sessions = append(sessions, current)
		}

		if current != nil {
			parseLineForRoutes(line, current)
		}

		if line == "" {
			current = nil
		}
	}

	return sessions
}

func parseLineForSession(line string, ipVersion int) (*session, bool) {
	match := protoRegex.FindStringSubmatch(line)

	if match != nil {
		session := session{name: match[1], ipVersion: ipVersion, established: parseState(match[5])}
		return &session, true
	}

	return nil, false
}

func parseLineForRoutes(line string, session *session) {
	match := routeRegex.FindStringSubmatch(line)

	if match != nil {
		session.imported, _ = strconv.ParseInt(match[1], 0, 64)
		session.exported, _ = strconv.ParseInt(match[2], 0, 64)
	}
}

func parseState(state string) int {
	if state == "Established" {
		return 1
	} else {
		return 0
	}
}
<|endoftext|>


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

// The block writer example program is a write-only workload intended to insert
// a large amount of data into cockroach quickly. This example is intended to
// trigger range splits and rebalances.
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"os"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/satori/go.uuid"
	// Import postgres driver.
	_ "github.com/cockroachdb/pq"
)

const (
	insertBlockStmt = `INSERT INTO blocks (block_id, writer_id, block_num, raw_bytes) VALUES ($1, $2, $3, $4)`
)

// concurrency = number of concurrent insertion processes.
var concurrency = flag.Int("concurrency", 2*runtime.NumCPU(), "Number of concurrent writers inserting blocks")

var tolerateErrors = flag.Bool("tolerate-errors", false, "Keep running on error")

// outputInterval = interval at which information is output to console.
var outputInterval = flag.Duration("output-interval", 1*time.Second, "Interval of output")

// Minimum and maximum size of inserted blocks.
var minBlockSizeBytes = flag.Int("min-block-bytes", 256, "Minimum amount of raw data written with each insertion")
var maxBlockSizeBytes = flag.Int("max-block-bytes", 1024, "Maximum amount of raw data written with each insertion")

// numBlocks keeps a global count of successfully written blocks.
var numBlocks uint64

// A blockWriter writes blocks of random data into cockroach in an infinite
// loop.
type blockWriter struct {
	id         string
	blockCount uint64
	db         *sql.DB
	rand       *rand.Rand
}

func newBlockWriter(db *sql.DB) blockWriter {
	source := rand.NewSource(int64(time.Now().UnixNano()))
	return blockWriter{
		db:   db,
		id:   uuid.NewV4().String(),
		rand: rand.New(source),
	}
}

// run is an infinite loop in which the blockWriter continuously attempts to
// write blocks of random data into a table in cockroach DB.
func (bw blockWriter) run(errCh chan<- error) {
	for {
		blockID := bw.rand.Int63()
		blockData := bw.randomBlock()
		bw.blockCount++
		if _, err := bw.db.Exec(insertBlockStmt, blockID, bw.id, bw.blockCount, blockData); err != nil {
			errCh <- fmt.Errorf("error running blockwriter %s: %s", bw.id, err)
		} else {
			atomic.AddUint64(&numBlocks, 1)
		}
	}
}

// randomBlock generates a slice of randomized bytes. Random data is preferred
// to prevent compression in storage.
func (bw blockWriter) randomBlock() []byte {
	blockSize := bw.rand.Intn(*maxBlockSizeBytes-*minBlockSizeBytes) + *minBlockSizeBytes
	blockData := make([]byte, blockSize)
	for i := range blockData {
		blockData[i] = byte(bw.rand.Int() & 0xff)
	}
	return blockData
}

// setupDatabase performs initial setup for the example, creating a database and
// with a single table. If the desired table already exists on the cluster, the
// existing table will be dropped.
func setupDatabase(dbURL string) (*sql.DB, error) {
	parsedURL, err := url.Parse(dbURL)
	if err != nil {
		return nil, err
	}
	parsedURL.Path = "datablocks"

	// Open connection to server and create a database.
	db, err := sql.Open("postgres", parsedURL.String())
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS datablocks"); err != nil {
		return nil, err
	}

	// Allow a maximum of concurrency+1 connections to the database.
	db.SetMaxOpenConns(*concurrency + 1)

	// Create the initial table for storing blocks.
	if _, err := db.Exec(`
	CREATE TABLE IF NOT EXISTS blocks (
	  block_id BIGINT NOT NULL,
	  writer_id STRING NOT NULL,
	  block_num BIGINT NOT NULL,
	  raw_bytes BYTES NOT NULL,
	  PRIMARY KEY (block_id, writer_id, block_num)
	)`); err != nil {
		return nil, err
	}

	return db, nil
}

var usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s <db URL>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = usage
	flag.Parse()

	dbURL := "postgresql://root@localhost:26257/photos?sslmode=disable"
	if flag.NArg() == 1 {
		dbURL = flag.Arg(0)
	}

	if *concurrency < 1 {
		log.Fatalf("Value of 'concurrency' flag (%d) must be greater than or equal to 1", *concurrency)
	}

	if max, min := *maxBlockSizeBytes, *minBlockSizeBytes; max < min {
		log.Fatalf("Value of 'max-block-bytes' (%d) must be greater than or equal to value of 'min-block-bytes' (%d)", max, min)
	}

	var db *sql.DB
	{
		var err error
		for err == nil || *tolerateErrors {
			db, err = setupDatabase(dbURL)
			if err == nil {
				break
			}
			if !*tolerateErrors {
				log.Fatal(err)
			}
		}
	}

	lastNow := time.Now()
	start := lastNow
	var lastNumDumps uint64
	writers := make([]blockWriter, *concurrency)

	errCh := make(chan error)
	for i := range writers {
		writers[i] = newBlockWriter(db)
		go writers[i].run(errCh)
	}

	var numErr int
	for range time.Tick(*outputInterval) {
		now := time.Now()
		elapsed := time.Since(lastNow)
		dumps := atomic.LoadUint64(&numBlocks)
		fmt.Printf("%6s: %6.1f/sec",
			time.Duration(time.Since(start).Seconds()+0.5)*time.Second,
			float64(dumps-lastNumDumps)/elapsed.Seconds())
		if numErr > 0 {
			fmt.Printf(" (%d total errors)\n", numErr)
		}
		fmt.Printf("\n")
		for {
			select {
			case err := <-errCh:
				numErr++
				if !*tolerateErrors {
					log.Fatal(err)
				} else {
					log.Print(err)
				}
				continue
			default:
			}
			break
		}
		lastNumDumps = dumps
		lastNow = now
	}
}
<commit_msg>Add --max-blocks and --duration flags<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Matt Tracy

// The block writer example program is a write-only workload intended to insert
// a large amount of data into cockroach quickly. This example is intended to
// trigger range splits and rebalances.
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/satori/go.uuid"
	// Import postgres driver.
	_ "github.com/cockroachdb/pq"
)

const (
	insertBlockStmt = `INSERT INTO blocks (block_id, writer_id, block_num, raw_bytes) VALUES ($1, $2, $3, $4)`
)

// concurrency = number of concurrent insertion processes.
var concurrency = flag.Int("concurrency", 2*runtime.NumCPU(), "Number of concurrent writers inserting blocks")

var tolerateErrors = flag.Bool("tolerate-errors", false, "Keep running on error")

// outputInterval = interval at which information is output to console.
var outputInterval = flag.Duration("output-interval", 1*time.Second, "Interval of output")

// Minimum and maximum size of inserted blocks.
var minBlockSizeBytes = flag.Int("min-block-bytes", 256, "Minimum amount of raw data written with each insertion")
var maxBlockSizeBytes = flag.Int("max-block-bytes", 1024, "Maximum amount of raw data written with each insertion")

var maxBlocks = flag.Uint64("max-blocks", 0, "Maximum number of blocks to write")
var duration = flag.Duration("duration", 0, "The duration to run. If 0, run forever.")

// numBlocks keeps a global count of successfully written blocks.
var numBlocks uint64

// A blockWriter writes blocks of random data into cockroach in an infinite
// loop.
type blockWriter struct {
	id         string
	blockCount uint64
	db         *sql.DB
	rand       *rand.Rand
}

func newBlockWriter(db *sql.DB) blockWriter {
	source := rand.NewSource(int64(time.Now().UnixNano()))
	return blockWriter{
		db:   db,
		id:   uuid.NewV4().String(),
		rand: rand.New(source),
	}
}

// run is an infinite loop in which the blockWriter continuously attempts to
// write blocks of random data into a table in cockroach DB.
func (bw blockWriter) run(errCh chan<- error, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		blockID := bw.rand.Int63()
		blockData := bw.randomBlock()
		bw.blockCount++
		if _, err := bw.db.Exec(insertBlockStmt, blockID, bw.id, bw.blockCount, blockData); err != nil {
			errCh <- fmt.Errorf("error running blockwriter %s: %s", bw.id, err)
		} else {
			v := atomic.AddUint64(&numBlocks, 1)
			if *maxBlocks > 0 && v >= *maxBlocks {
				return
			}
		}
	}
}

// randomBlock generates a slice of randomized bytes. Random data is preferred
// to prevent compression in storage.
func (bw blockWriter) randomBlock() []byte {
	blockSize := bw.rand.Intn(*maxBlockSizeBytes-*minBlockSizeBytes) + *minBlockSizeBytes
	blockData := make([]byte, blockSize)
	for i := range blockData {
		blockData[i] = byte(bw.rand.Int() & 0xff)
	}
	return blockData
}

// setupDatabase performs initial setup for the example, creating a database and
// with a single table. If the desired table already exists on the cluster, the
// existing table will be dropped.
func setupDatabase(dbURL string) (*sql.DB, error) {
	parsedURL, err := url.Parse(dbURL)
	if err != nil {
		return nil, err
	}
	parsedURL.Path = "datablocks"

	// Open connection to server and create a database.
	db, err := sql.Open("postgres", parsedURL.String())
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS datablocks"); err != nil {
		return nil, err
	}

	// Allow a maximum of concurrency+1 connections to the database.
	db.SetMaxOpenConns(*concurrency + 1)

	// Create the initial table for storing blocks.
	if _, err := db.Exec(`
	CREATE TABLE IF NOT EXISTS blocks (
	  block_id BIGINT NOT NULL,
	  writer_id STRING NOT NULL,
	  block_num BIGINT NOT NULL,
	  raw_bytes BYTES NOT NULL,
	  PRIMARY KEY (block_id, writer_id, block_num)
	)`); err != nil {
		return nil, err
	}

	return db, nil
}

var usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s <db URL>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = usage
	flag.Parse()

	dbURL := "postgresql://root@localhost:26257/photos?sslmode=disable"
	if flag.NArg() == 1 {
		dbURL = flag.Arg(0)
	}

	if *concurrency < 1 {
		log.Fatalf("Value of 'concurrency' flag (%d) must be greater than or equal to 1", *concurrency)
	}

	if max, min := *maxBlockSizeBytes, *minBlockSizeBytes; max < min {
		log.Fatalf("Value of 'max-block-bytes' (%d) must be greater than or equal to value of 'min-block-bytes' (%d)", max, min)
	}

	var db *sql.DB
	{
		var err error
		for err == nil || *tolerateErrors {
			db, err = setupDatabase(dbURL)
			if err == nil {
				break
			}
			if !*tolerateErrors {
				log.Fatal(err)
			}
		}
	}

	lastNow := time.Now()
	start := lastNow
	var lastNumDumps uint64
	writers := make([]blockWriter, *concurrency)

	errCh := make(chan error)
	var wg sync.WaitGroup
	for i := range writers {
		wg.Add(1)
		writers[i] = newBlockWriter(db)
		go writers[i].run(errCh, &wg)
	}

	var numErr int
	tick := time.Tick(*outputInterval)
	done := make(chan struct{}, 1)

	go func() {
		wg.Wait()
		done <- struct{}{}
	}()

	if *duration > 0 {
		go func() {
			time.Sleep(*duration)
			done <- struct{}{}
		}()
	}

	for {
		select {
		case err := <-errCh:
			numErr++
			if !*tolerateErrors {
				log.Fatal(err)
			} else {
				log.Print(err)
			}
			continue

		case <-tick:
			now := time.Now()
			elapsed := time.Since(lastNow)
			dumps := atomic.LoadUint64(&numBlocks)
			fmt.Printf("%6s: %6.1f/sec",
				time.Duration(time.Since(start).Seconds()+0.5)*time.Second,
				float64(dumps-lastNumDumps)/elapsed.Seconds())
			if numErr > 0 {
				fmt.Printf(" (%d total errors)", numErr)
			}
			fmt.Printf("\n")
			lastNumDumps = dumps
			lastNow = now

		case <-done:
			dumps := atomic.LoadUint64(&numBlocks)
			elapsed := time.Since(start).Seconds()
			fmt.Printf("%6.1fs: %d total %6.1f/sec",
				elapsed, dumps, float64(dumps)/elapsed)
			if numErr > 0 {
				fmt.Printf(" (%d total errors)\n", numErr)
			}
			fmt.Printf("\n")
			os.Exit(0)
		}
	}
}
<|endoftext|>


<commit_before>package blocks

import (
	"bytes"
	"encoding/json"
	"log"
	"net/http"
	"strings"
)

func PostValue(b *Block) {

	type KeyMapping struct {
		MsgKey   string
		QueryKey string
	}

	type postRule struct {
		Keymapping []KeyMapping
		Endpoint   string
	}

	var rule *postRule
	client := &http.Client{}

	// TODO check the endpoint for happiness
	for {
		select {
		case msg := <-b.AddChan:
			updateOutChans(msg, b)
		case <-b.QuitChan:
			quit(b)
			return
		case msg := <-b.Routes["get_rule"]:
			if rule == nil {
				marshal(msg, &postRule{Keymapping:[]KeyMapping{KeyMapping{}}})
			} else {
				marshal(msg, rule)
			}
		case msg := <-b.Routes["set_rule"]:
			if rule == nil {
				rule = &postRule{}
			}
			unmarshal(msg, rule)

		case msg := <-b.InChan:
			if rule == nil {
				break
			}

			body := make(map[string]interface{})
			for _, keymap := range rule.Keymapping {
				keys := strings.Split(keymap.MsgKey, ".")
				value, err := Get(msg, keys...)
				if err != nil {
					log.Println(err.Error())
				} else {
					Set(body, keymap.QueryKey, value)
				}
			}

			// TODO maybe check the response ?
			postBody, err := json.Marshal(body)
			if err != nil {
				log.Fatal(err.Error())
			}

			// TODO the content-type here is heavily borked but we're using a hack
			/*resp, err := http.Post(rule.Endpoint, "application/x-www-form-urlencoded", bytes.NewReader(postBody))
			if err != nil {
				log.Println(err.Error())
				break
			}
			defer resp.Body.Close()*/

			req, err := http.NewRequest("POST", rule.Endpoint, bytes.NewReader(postBody))

			req.Close()

			if err != nil {
			    log.Println(err)
			    break
			}

			resp, err := client.Do(req)
			defer resp.Body.Close()

			if err != nil {
			    log.Println(err)
			    break
			}

			bodyBuf := &bytes.Buffer{}
			_, err = bodyBuf.ReadFrom(resp.Body)
		}
	}
}
<commit_msg>added content type to header<commit_after>package blocks

import (
	"bytes"
	"encoding/json"
	"log"
	"net/http"
	"strings"
)

func PostValue(b *Block) {

	type KeyMapping struct {
		MsgKey   string
		QueryKey string
	}

	type postRule struct {
		Keymapping []KeyMapping
		Endpoint   string
	}

	var rule *postRule
	client := &http.Client{}

	// TODO check the endpoint for happiness
	for {
		select {
		case msg := <-b.AddChan:
			updateOutChans(msg, b)
		case <-b.QuitChan:
			quit(b)
			return
		case msg := <-b.Routes["get_rule"]:
			if rule == nil {
				marshal(msg, &postRule{Keymapping:[]KeyMapping{KeyMapping{}}})
			} else {
				marshal(msg, rule)
			}
		case msg := <-b.Routes["set_rule"]:
			if rule == nil {
				rule = &postRule{}
			}
			unmarshal(msg, rule)

		case msg := <-b.InChan:
			if rule == nil {
				break
			}

			body := make(map[string]interface{})
			for _, keymap := range rule.Keymapping {
				keys := strings.Split(keymap.MsgKey, ".")
				value, err := Get(msg, keys...)
				if err != nil {
					log.Println(err.Error())
				} else {
					Set(body, keymap.QueryKey, value)
				}
			}

			// TODO maybe check the response ?
			postBody, err := json.Marshal(body)
			if err != nil {
				log.Fatal(err.Error())
			}

			// TODO the content-type here is heavily borked but we're using a hack
			req, err := http.NewRequest("POST", rule.Endpoint, bytes.NewReader(postBody))
			req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
			req.Close = true

			if err != nil {
			    log.Println(err)
			    break
			}

			resp, err := client.Do(req)
			defer resp.Body.Close()

			if err != nil {
			    log.Println(err)
			    break
			}

			bodyBuf := &bytes.Buffer{}
			_, err = bodyBuf.ReadFrom(resp.Body)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package interpolate

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestParseSuccess(t *testing.T) {
	tests := []struct {
		give string
		want String
	}{
		{
			give: "foo",
			want: String{literal("foo")},
		},
		{
			give: "foo ${bar} baz",
			want: String{
				literal("foo "),
				variable{Name: "bar"},
				literal(" baz"),
			},
		},
		{
			give: "foo $ {bar} baz",
			want: String{literal("foo "), literal("$ "), literal("{bar} baz")},
		},
		{
			give: "foo ${bar:}",
			want: String{
				literal("foo "),
				variable{Name: "bar", HasDefault: true},
			},
		},
		{
			give: "${foo:bar}",
			want: String{
				variable{
					Name:       "foo",
					Default:    "bar",
					HasDefault: true,
				},
			},
		},
		{
			give: `foo \${bar:42} baz`,
			want: String{
				literal("foo "),
				literal("$"),
				literal("{bar:42} baz"),
			},
		},
		{
			give: "$foo${bar}",
			want: String{
				literal("$f"),
				literal("oo"),
				variable{Name: "bar"},
			},
		},
		{
			give: "foo${b-a-r}",
			want: String{
				literal("foo"),
				variable{Name: "b-a-r"},
			},
		},
		{
			give: "foo ${bar::baz} qux",
			want: String{
				literal("foo "),
				variable{Name: "bar", HasDefault: true, Default: ":baz"},
				literal(" qux"),
			},
		},
		{
			give: "a ${b:hello world} c",
			want: String{
				literal("a "),
				variable{Name: "b", HasDefault: true, Default: "hello world"},
				literal(" c"),
			},
		},
	}

	for _, tt := range tests {
		out, err := Parse(tt.give)
		assert.NoError(t, err)
		assert.Equal(t, tt.want, out)
	}
}

func TestParseFailures(t *testing.T) {
	tests := []string{
		"${foo",
		"${foo.}",
		"${foo-}",
		"${foo--bar}",
	}

	for _, tt := range tests {
		_, err := Parse(tt)
		assert.Error(t, err, tt)
	}
}
<commit_msg>interpolate: Add test for $$ behavior (#1268)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package interpolate

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestParseSuccess(t *testing.T) {
	tests := []struct {
		give string
		want String
	}{
		{
			give: "foo",
			want: String{literal("foo")},
		},
		{
			give: "foo ${bar} baz",
			want: String{
				literal("foo "),
				variable{Name: "bar"},
				literal(" baz"),
			},
		},
		{
			give: "foo $ {bar} baz",
			want: String{literal("foo "), literal("$ "), literal("{bar} baz")},
		},
		{
			give: "foo ${bar:}",
			want: String{
				literal("foo "),
				variable{Name: "bar", HasDefault: true},
			},
		},
		{
			give: "${foo:bar}",
			want: String{
				variable{
					Name:       "foo",
					Default:    "bar",
					HasDefault: true,
				},
			},
		},
		{
			give: `foo \${bar:42} baz`,
			want: String{
				literal("foo "),
				literal("$"),
				literal("{bar:42} baz"),
			},
		},
		{
			give: "$foo${bar}",
			want: String{
				literal("$f"),
				literal("oo"),
				variable{Name: "bar"},
			},
		},
		{
			give: "foo${b-a-r}",
			want: String{
				literal("foo"),
				variable{Name: "b-a-r"},
			},
		},
		{
			give: "foo ${bar::baz} qux",
			want: String{
				literal("foo "),
				variable{Name: "bar", HasDefault: true, Default: ":baz"},
				literal(" qux"),
			},
		},
		{
			give: "a ${b:hello world} c",
			want: String{
				literal("a "),
				variable{Name: "b", HasDefault: true, Default: "hello world"},
				literal(" c"),
			},
		},
		{
			give: "foo $${bar}",
			want: String{literal("foo "), literal("$$"), literal("{bar}")},
		},
	}

	for _, tt := range tests {
		out, err := Parse(tt.give)
		assert.NoError(t, err)
		assert.Equal(t, tt.want, out)
	}
}

func TestParseFailures(t *testing.T) {
	tests := []string{
		"${foo",
		"${foo.}",
		"${foo-}",
		"${foo--bar}",
	}

	for _, tt := range tests {
		_, err := Parse(tt)
		assert.Error(t, err, tt)
	}
}
<|endoftext|>


<commit_before>package v8js

import (
	"fmt"
	"github.com/GeertJohan/go.rice"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/runner"
	"github.com/ry/v8worker"
	"github.com/valyala/fasthttp"
	"golang.org/x/net/context"
	"math"
	"os"
	"time"
)

type libFile struct {
	Filename string
	Source   string
}

type Runner struct {
	Filename string
	Source   string
	Client   *fasthttp.Client

	stdlib []libFile
}

type VUContext struct {
	r    *Runner
	ctx  context.Context
	ch   chan runner.Result
	mods map[string]Module
}

type Module map[string]Member

type Member interface{}

func New(filename, src string) *Runner {
	r := &Runner{
		Filename: filename,
		Source:   src,
		Client: &fasthttp.Client{
			Dial:                fasthttp.Dial,
			MaxIdleConnDuration: time.Duration(0),
			MaxConnsPerHost:     math.MaxInt64,
		},
	}

	// Load the standard library as a rice box; panic if any part of this fails
	// (The only possible cause is a programming/developer error, not user error)
	box := rice.MustFindBox("lib")
	box.Walk("/", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			panic(err)
		}
		if !info.IsDir() {
			r.stdlib = append(r.stdlib, libFile{
				Filename: path,
				Source:   box.MustString(path),
			})
		}
		return nil
	})

	return r
}

func (r *Runner) Run(ctx context.Context, id int64) <-chan runner.Result {
	ch := make(chan runner.Result)

	go func() {
		defer close(ch)

		vu := VUContext{r: r, ctx: ctx, ch: ch}
		w := v8worker.New(vu.Recv, vu.RecvSync)

		w.Load("internal:constants", fmt.Sprintf(`var __id = %d;`, id))

		for _, f := range r.stdlib {
			if err := w.Load(f.Filename, f.Source); err != nil {
				log.WithError(err).WithField("file", f.Filename).Error("Couldn't load lib")
			}
		}

		if err := vu.RegisterModules(w); err != nil {
			log.WithError(err).Error("Couldn't register bridged functions")
			return
		}

		src := fmt.Sprintf(`
		function __run__() {
			try {
		%s
			} catch (e) {
				console.error(e);
			}
		}
		`, r.Source)
		if err := w.Load(r.Filename, src); err != nil {
			log.WithError(err).Error("Couldn't load JS")
			return
		}

		for {
			// log.Info("-> run")
			w.SendSync("run")
			// log.Info("<- run")

			select {
			case <-ctx.Done():
				return
			default:
			}
		}
	}()

	return ch
}
<commit_msg>Better logging<commit_after>package v8js

import (
	"fmt"
	"github.com/GeertJohan/go.rice"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/runner"
	"github.com/ry/v8worker"
	"github.com/valyala/fasthttp"
	"golang.org/x/net/context"
	"math"
	"os"
	"time"
)

type libFile struct {
	Filename string
	Source   string
}

type Runner struct {
	Filename string
	Source   string
	Client   *fasthttp.Client

	stdlib []libFile
}

type VUContext struct {
	r    *Runner
	ctx  context.Context
	ch   chan runner.Result
	mods map[string]Module
}

type Module map[string]Member

type Member interface{}

func New(filename, src string) *Runner {
	r := &Runner{
		Filename: filename,
		Source:   src,
		Client: &fasthttp.Client{
			Dial:                fasthttp.Dial,
			MaxIdleConnDuration: time.Duration(0),
			MaxConnsPerHost:     math.MaxInt64,
		},
	}

	// Load the standard library as a rice box; panic if any part of this fails
	// (The only possible cause is a programming/developer error, not user error)
	box := rice.MustFindBox("lib")
	box.Walk("/", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			panic(err)
		}
		if !info.IsDir() {
			r.stdlib = append(r.stdlib, libFile{
				Filename: path,
				Source:   box.MustString(path),
			})
		}
		return nil
	})

	return r
}

func (r *Runner) Run(ctx context.Context, id int64) <-chan runner.Result {
	ch := make(chan runner.Result)

	go func() {
		defer close(ch)

		vu := VUContext{r: r, ctx: ctx, ch: ch}
		w := v8worker.New(vu.Recv, vu.RecvSync)

		w.Load("internal:constants", fmt.Sprintf(`var __id = %d;`, id))

		for _, f := range r.stdlib {
			if err := w.Load(f.Filename, f.Source); err != nil {
				log.WithError(err).WithField("file", f.Filename).Error("Couldn't load lib")
			}
		}

		if err := vu.RegisterModules(w); err != nil {
			log.WithError(err).Error("Couldn't register bridged functions")
			return
		}

		src := fmt.Sprintf(`
		function __run__() {
			try {
		%s
			} catch (e) {
				console.error("Script Error", '' + e);
			}
		}
		`, r.Source)
		if err := w.Load(r.Filename, src); err != nil {
			log.WithError(err).Error("Couldn't load JS")
			return
		}

		for {
			// log.Info("-> run")
			w.SendSync("run")
			// log.Info("<- run")

			select {
			case <-ctx.Done():
				return
			default:
			}
		}
	}()

	return ch
}
<|endoftext|>


<commit_before>// Copyright 2017 Jonathan Pincas

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

// 	http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package core

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"

	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var isInstallDemoData, isReinstall bool

func init() {
	RootCmd.AddCommand(installCmd)
	RootCmd.AddCommand(unInstallCmd)
	installCmd.Flags().BoolVar(&isInstallDemoData, "demodata", false, "Install bundle demo data if available")
	installCmd.Flags().BoolVarP(&isReinstall, "reinstall", "r", false, "Uninstall bundle before installing")
}

// installCmd represents the install command
var installCmd = &cobra.Command{
	Use:   "install [bundle]",
	Short: "Install an EcoSystem bundle",
	Long: `Installs an EcoSystem bundle from the named folder.
	Note: does not download anything, so the bundle folder must
	exist and contain everything.  Previous to installing, either clone
	or download the bundle into the 'bundles' directory`,
	RunE: installBundle,
}

// installCmd represents the install command
var unInstallCmd = &cobra.Command{
	Use:   "uninstall [bundle]",
	Short: "Removes an EcoSystem bundle",
	Long: `Removes an EcoSystem bundle by deleting the DB schema, deleting template
	files and images,`,
	RunE: unInstallBundle,
}

//uninstallBundle is the removal function for a bundle
func unInstallBundle(cmd *cobra.Command, args []string) error {

	readConfig()

	//Check for bundle name
	if len(args) < 1 {
		return errors.New("a bundle name must be provided")
	}

	//Ask for confirmation
	c := AskForConfirmation("This will delete the bundle, causing loss of all data in the schema created by the bundle.  Are you sure you want to do this?")

	if c {
		//Establish a temporary connection as the super user
		db := SuperUserDBConfig.ReturnDBConnection("")
		defer db.Close()

		//Drop the schema
		//If it doesn't exist, it won't be dropped - no big deal
		db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))

		//Attempt to updated the bundles installed list
		newBundlesInstalled, err := Bundles(viper.GetStringSlice("bundlesInstalled")).UnInstallBundle(args[0])

		//If there is any error, log it
		if err != nil {
			Log(LogEntry{"CORE.INSTALL", false, "Error updating bundles installed list: " + err.Error()})
		}

		//Otherwise set the viper configuration to the new bundles list and overwrite the config.json
		viper.Set("bundlesInstalled", newBundlesInstalled)
		var config Config
		viper.Unmarshal(&config)
		configJSON, _ := json.MarshalIndent(config, "", "\t")
		err = ioutil.WriteFile(viper.GetString("configfile")+".json", configJSON, 0644)
		if err != nil {
			Log(LogEntry{"CORE.INSTALL", false, "Error updating config file: " + err.Error()})
		}

		Log(LogEntry{"CORE.INSTALL", true, "config.json updated"})
		Log(LogEntry{"CORE.INSTALL", true, "Uninstallation of bundle " + args[0] + " completed"})

	}

	return nil

}

//installBundle is the entire installation procedure for an EcoSystem Bundle
func installBundle(cmd *cobra.Command, args []string) error {

	readConfig()

	//Check for bundle name
	if len(args) < 1 {
		return errors.New("a bundle name must be provided")
	}

	//Check that bundle exists
	basePath := "./bundles/" + args[0]
	exists, _ := afero.IsDir(AppFs, basePath)
	if !exists {
		//Exit if doesn't exist
		Log(LogEntry{"CORE.INSTALL", false, "Bundle " + args[0] + " not found.  Please download or clone."})
	}

	//Uninstall first if requested
	if isReinstall {
		Log(LogEntry{"CORE.INSTALL", true, "Uninstalling bundle " + args[0] + " before reinstalling"})
		unInstallBundle(cmd, args)
	}

	//Establish a temporary connection as the super user
	db := SuperUserDBConfig.ReturnDBConnection("")
	defer db.Close()

	//Check for the presence of install.sql and attempt to read it
	sqlBytes, err := afero.ReadFile(AppFs, basePath+"/install.sql")
	if err != nil {
		Log(LogEntry{"CORE.INSTALL", false, "install.sql not present for this bundle, or could not be read: " + err.Error()})
	} else {
		sqlString := string(sqlBytes)

		//Install the DB setup and logic
		//Attempt to create a schema matching the bundle's name,
		_, err := db.Exec(fmt.Sprintf(SQLToCreateSchema, args[0]))

		if err != nil {

			//if the schema exists, the bundle is already installed
			//don't go any further with db part of the bundle
			Log(LogEntry{"CORE.INSTALL", false, "Failed to create DB schema: " + err.Error()})

		} else {

			//Set admin privileges for everything in this schema going forwards
			_, err = db.Exec(fmt.Sprintf(SQLToGrantBundleAdminPermissions, args[0], args[0], args[0]))

			//Set the search path to the bundle schema so that all SQL commands take
			//place within the schema
			_, err = db.Exec(fmt.Sprintf(SQLToSetSearchPathForBundle, args[0]))
			if err != nil {

				//If there is any problem with the search path, give up the db part of the bundle
				Log(LogEntry{"CORE.INSTALL", false, "search_path failed to set, aborting sql installation and cleaning up" + err.Error()})
				db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))

			} else {

				//Run the SQL
				_, err = db.Exec(sqlString)
				if err != nil {
					Log(LogEntry{"CORE.INSTALL", false, "Problem with install.sql, aborting sql installation and cleaning up" + err.Error()})
					db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))
				} else {

					//If all is good so far and the user has specified it, install the demo data (if it exists)
					if isInstallDemoData {

						//Check for the presence of demodata.sql and attempt to read it
						sqlBytes, err := afero.ReadFile(AppFs, basePath+"/demodata.sql")
						if err != nil {
							//If there is no demodata.sql
							Log(LogEntry{"CORE.INSTALL", false, "demodata.sql not present for this bundle, or could not be read: " + err.Error()})
						} else {
							sqlString := string(sqlBytes)
							_, err = db.Exec(sqlString)
							if err != nil {
								//If there is an error with demodata.sql
								Log(LogEntry{"CORE.INSTALL", false, "Error installing demo data: " + err.Error()})
							}
						}
					}
				}
			}
		}
	} //End sql installation

	//Attempt to updated the bundles installed list
	newBundlesInstalled, err := Bundles(viper.GetStringSlice("bundlesInstalled")).InstallBundle(args[0])

	//If there is any error, return it
	if err != nil {
		Log(LogEntry{"CORE.INSTALL", false, "Error updating bundles installed list: " + err.Error()})
	}

	//Otherwise set the viper configuration to the new bundles list and overwrite the config.json
	viper.Set("bundlesInstalled", newBundlesInstalled)
	var config Config
	viper.Unmarshal(&config)
	configJSON, _ := json.MarshalIndent(config, "", "\t")
	err = ioutil.WriteFile(viper.GetString("configfile")+".json", configJSON, 0644)
	if err != nil {
		Log(LogEntry{"CORE.INSTALL", false, "Error updating config file: " + err.Error()})
	}

	Log(LogEntry{"CORE.INSTALL", true, "config file updated"})

	//Bundle installation complete
	Log(LogEntry{"CORE.INSTALL", true, "Installation of bundle " + args[0] + " completed"})
	return nil

}
<commit_msg>Now use multiple install and demo data sql files<commit_after>// Copyright 2017 Jonathan Pincas

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

// 	http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package core

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"path"

	"database/sql"

	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var isInstallDemoData, isReinstall bool

func init() {
	RootCmd.AddCommand(installCmd)
	RootCmd.AddCommand(unInstallCmd)
	installCmd.Flags().BoolVar(&isInstallDemoData, "demodata", false, "Install bundle demo data if available")
	installCmd.Flags().BoolVarP(&isReinstall, "reinstall", "r", false, "Uninstall bundle before installing")
}

// installCmd represents the install command
var installCmd = &cobra.Command{
	Use:   "install [bundle]",
	Short: "Install an EcoSystem bundle",
	Long: `Installs an EcoSystem bundle from the named folder.
	Note: does not download anything, so the bundle folder must
	exist and contain everything.  Previous to installing, either clone
	or download the bundle into the 'bundles' directory`,
	RunE: installBundle,
}

// installCmd represents the install command
var unInstallCmd = &cobra.Command{
	Use:   "uninstall [bundle]",
	Short: "Removes an EcoSystem bundle",
	Long: `Removes an EcoSystem bundle by deleting the DB schema, deleting template
	files and images,`,
	RunE: unInstallBundle,
}

//uninstallBundle is the removal function for a bundle
func unInstallBundle(cmd *cobra.Command, args []string) error {

	readConfig()

	//Check for bundle name
	if len(args) < 1 {
		return errors.New("a bundle name must be provided")
	}

	//Ask for confirmation
	c := AskForConfirmation("This will delete the bundle, causing loss of all data in the schema created by the bundle.  Are you sure you want to do this?")

	if c {
		//Establish a temporary connection as the super user
		db := SuperUserDBConfig.ReturnDBConnection("")
		defer db.Close()

		//Drop the schema
		//If it doesn't exist, it won't be dropped - no big deal
		db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))

		//Attempt to updated the bundles installed list
		newBundlesInstalled, err := Bundles(viper.GetStringSlice("bundlesInstalled")).UnInstallBundle(args[0])

		//If there is any error, log it
		if err != nil {
			Log(LogEntry{"CORE.INSTALL", false, "Error updating bundles installed list: " + err.Error()})
		}

		//Otherwise set the viper configuration to the new bundles list and overwrite the config.json
		viper.Set("bundlesInstalled", newBundlesInstalled)
		var config Config
		viper.Unmarshal(&config)
		configJSON, _ := json.MarshalIndent(config, "", "\t")
		err = ioutil.WriteFile(viper.GetString("configfile")+".json", configJSON, 0644)
		if err != nil {
			Log(LogEntry{"CORE.INSTALL", false, "Error updating config file: " + err.Error()})
		}

		Log(LogEntry{"CORE.INSTALL", true, "config.json updated"})
		Log(LogEntry{"CORE.INSTALL", true, "Uninstallation of bundle " + args[0] + " completed"})

	}

	return nil

}

//installBundle is the entire installation procedure for an EcoSystem Bundle
func installBundle(cmd *cobra.Command, args []string) error {

	readConfig()

	//Check for bundle name
	if len(args) < 1 {
		return errors.New("a bundle name must be provided")
	}

	//Check that bundle installation folder exists
	basePath := "./bundles/" + args[0] + "/install"

	exists, err := afero.IsDir(AppFs, basePath)

	if !exists || err != nil {
		//Exit if doesn't exist
		LogFatal(LogEntry{"CORE.INSTALL", false, "Bundle '" + args[0] + "' install folder not found or unreadable."})
	}

	//Uninstall first if requested
	if isReinstall {
		Log(LogEntry{"CORE.INSTALL", true, "Uninstalling bundle " + args[0] + " before reinstalling"})
		unInstallBundle(cmd, args)
	}

	//Check for error reading directory or zero files
	filesInDirectory, err := afero.ReadDir(AppFs, basePath)
	if err != nil || len(filesInDirectory) == 0 {
		LogFatal(LogEntry{"CORE.INSTALL", false, "No installation files could be read for bundle"})
		return nil
	}

	Log(LogEntry{"CORE.INSTALL", true, "Installing bundle '" + args[0] + "'"})

	//Establish a temporary connection as the super user
	db := SuperUserDBConfig.ReturnDBConnection("")
	defer db.Close()

	//Set up a schema for the bundle
	err = setupDBSchema(db, args[0])
	if err != nil {
		//IF there is any type of error, drop the schema, log and exit
		db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))
		LogFatal(LogEntry{"CORE.INSTALL", false, "Schema creation failed with error " + err.Error()})
		return nil
	}

	//Iterate over the installation files
	for _, file := range filesInDirectory {
		//Ignore directories
		if !file.IsDir() {
			//Attempt to processes the sqlfile
			err := processBundleFile(db, path.Join(basePath, file.Name()))
			if err != nil {
				//IF there is any type of error, drop the schema, log and exit
				db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))
				LogFatal(LogEntry{"CORE.INSTALL", false, "Installation of '" + file.Name() + "' failed with error: " + err.Error()})
				return nil
			}
			Log(LogEntry{"CORE.INSTALL", true, file.Name() + " installed OK"})
		}
	}

	//If the user has asked for demo data
	if isInstallDemoData {

		Log(LogEntry{"CORE.INSTALL", true, "Installing demo data"})

		basePath := "./bundles/" + args[0] + "/demodata"

		//Check for error reading directory or zero files
		filesInDirectory, err := afero.ReadDir(AppFs, basePath)
		if err != nil || len(filesInDirectory) == 0 {
			//IF there is any type of error, drop the schema, log and exit
			db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))
			LogFatal(LogEntry{"CORE.INSTALL", false, "No demo data files could be read for bundle"})
			return nil
		}

		//Iterate over the demodata files
		for _, file := range filesInDirectory {
			//Ignore directories
			if !file.IsDir() {
				//Attempt to processes the sqlfile
				err := processBundleFile(db, path.Join(basePath, file.Name()))
				if err != nil {
					//IF there is any type of error, drop the schema, log and exit
					db.Exec(fmt.Sprintf(SQLToDropSchema, args[0]))
					LogFatal(LogEntry{"CORE.INSTALL", false, "Installation of '" + file.Name() + "' failed with error: " + err.Error()})
					return nil
				}

				Log(LogEntry{"CORE.INSTALL", true, file.Name() + " installed OK"})

			}
		}

	}

	//Attempt to updated the bundles installed list
	newBundlesInstalled, err := Bundles(viper.GetStringSlice("bundlesInstalled")).InstallBundle(args[0])

	//If there is any error, return it
	if err != nil {
		Log(LogEntry{"CORE.INSTALL", false, "Error updating bundles installed list: " + err.Error()})
	}

	//Otherwise set the viper configuration to the new bundles list and overwrite the config.json
	viper.Set("bundlesInstalled", newBundlesInstalled)
	var config Config
	viper.Unmarshal(&config)
	configJSON, _ := json.MarshalIndent(config, "", "\t")
	err = ioutil.WriteFile(viper.GetString("configfile")+".json", configJSON, 0644)
	if err != nil {
		Log(LogEntry{"CORE.INSTALL", false, "Error updating config file: " + err.Error()})
	}

	Log(LogEntry{"CORE.INSTALL", true, "config file updated"})

	//Bundle installation complete
	Log(LogEntry{"CORE.INSTALL", true, "Installation of bundle " + args[0] + " completed"})
	return nil

}

func processBundleFile(db *sql.DB, filename string) error {

	//Attempt to read file
	sqlBytes, err := afero.ReadFile(AppFs, filename)

	if err != nil {
		return err
	}

	//Run the SQL
	if _, err = db.Exec(string(sqlBytes)); err != nil {
		return err
	}

	return nil

}

func setupDBSchema(db *sql.DB, bundleName string) error {

	//Attempt to create a schema matching the bundle's name,
	_, err := db.Exec(fmt.Sprintf(SQLToCreateSchema, bundleName))

	if err != nil {
		return err
	}

	//Set admin privileges for everything in this schema going forwards
	_, err = db.Exec(fmt.Sprintf(SQLToGrantBundleAdminPermissions, bundleName, bundleName, bundleName))

	if err != nil {
		return err
	}

	//Set the search path to the bundle schema so that all SQL commands take
	//place within the schema
	_, err = db.Exec(fmt.Sprintf(SQLToSetSearchPathForBundle, bundleName))

	if err != nil {
		return err
	}

	return nil

}
<|endoftext|>