Source

go-meetup / 5-short / shortie / shortie.go

/* AppEngine demo - a URL shortener */
package shortie

import (
	"fmt"
	"html/template"
	"net/http"
	"regexp"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/user"
)

const (
	counterKeyName = "counter-key-name"
	counterKind    = "Counter"
	urlKind        = "Url"
)

// Global counter of urls
type Counter struct {
	Count int64
}

// URL stored in database
type URL struct {
	Short   string
	Long    string
	User    string
	Created time.Time
	Hits    int64
}

// Parameters for homeTemplate.
type homeParams struct {
	User       string
	LoginTitle string
	LoginURL   string
	Error      string
	Count      int64
	ShortURL   string
}

var homeTemplate *template.Template

func init() {
	homeTemplate = template.Must(template.New("home").Parse(homeHTML))
	http.HandleFunc("/", rootHandler)
}

// rootHandler handles the main page.
func rootHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	var err error
	params := new(homeParams)

	// Run at end. We check "err" and update params if needed. Then serve homeTemplate.
	defer func() {
		if err != nil {
			params.Error = err.Error()
			ctx.Errorf("%v", err) // Log error
		}
		homeTemplate.Execute(w, params)
	}()

	err = fillUser(r, ctx, params)
	if err != nil {
		return
	}
	params.Count, err = urlCount(ctx)

	if r.Method == "POST" {
		longURL := strings.TrimSpace(r.FormValue("url"))

		if !isValidURL(longURL) {
			err = fmt.Errorf("Bad URL - %s", longURL)
			return
		}

		if !hasSchema(longURL) {
			longURL = fmt.Sprintf("http://%s", longURL)
		}

		var id string
		id, err = newShortURL(ctx, id, longURL)
		if err != nil {
			return
		}

		params.ShortURL = fullURL(r, id)
	}
}

// fillUser fills user details in template parameters
func fillUser(r *http.Request, ctx appengine.Context, params *homeParams) error {
	var err error
	u := user.Current(ctx)

	if u != nil {
		params.User = u.String()
		params.LoginTitle = "Logout"
		params.LoginURL, err = user.LogoutURL(ctx, r.URL.String())
	} else {
		params.User = "Stranger"
		params.LoginTitle = "Login"
		params.LoginURL, err = user.LoginURL(ctx, r.URL.String())
	}

	return err
}

// urlCount return the current count of urls.
func urlCount(ctx appengine.Context) (int64, error) {
	key := datastore.NewKey(ctx, counterKind, counterKeyName, 0, nil)
	counter := new(Counter)
	if err := datastore.Get(ctx, key, counter); err != nil && err != datastore.ErrNoSuchEntity {
		return 0, err
	}

	return counter.Count, nil
}

/* nextId returns the next short url. 
We use the global counter and then encode the last count in base62.
*/
func nextId(ctx appengine.Context) (string, error) {
	var count int64

	err := datastore.RunInTransaction(ctx, func(ctx appengine.Context) error {
		key := datastore.NewKey(ctx, counterKind, counterKeyName, 0, nil)
		counter := new(Counter)
		if err := datastore.Get(ctx, key, counter); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}

		counter.Count++
		if _, err := datastore.Put(ctx, key, counter); err != nil {
			return err
		}
		count = counter.Count
		return nil
	}, nil)

	return base62Encode(uint64(count)), err
}

// isValidURL check that URL is valid
func isValidURL(url string) bool {
	return (len(url) > 0) && strings.Contains(url, ".")
}

// hasSchema check if url has schema prefix.
func hasSchema(url string) bool {
	match, _ := regexp.MatchString("^[a-zA-Z]+://", url)
	return match
}

/* fullURL adds http://<host> suffix to short url. 
This works both locally and on AppEngine.
*/
func fullURL(r *http.Request, id string) string {
	return fmt.Sprintf("http://%s/%s", r.Host, id)
}

func newShortURL(ctx appengine.Context, longURL, user string) (string, error) {
	var id string
	id, err := nextId(ctx)

	if err != nil {
		return "", err
	}

	url := &URL{
		Short:   id,
		Long:    longURL,
		User:    user,
		Created: time.Now(),
		Hits:    0,
	}
	key := datastore.NewKey(ctx, urlKind, id, 0, nil)
	_, err = datastore.Put(ctx, key, url)
	if err != nil {
		return "", err
	}

	return id, nil
}