Source

go-appengine-demo / shortie / shortie.go

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

import (
	"appengine"
	"appengine/datastore"
	"appengine/memcache"
	"appengine/user"
	"appengine/taskqueue"
	"fmt"
	"http"
	"os"
	"regexp"
	"strings"
	"template"
	"time"
)

const (
	base62Alphabet = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	counterKeyName = "counter-key-name"
	counterKind    = "Counter"
	urlKind        = "Url"
	hitWorkerPath  = "/_worker/hit"
)

// base62Encode encodes a number to a base62 string representation.
func base62Encode(num uint64) string {
	if num == 0 {
		return "0"
	}

	arr := []uint8{}
	base := uint64(len(base62Alphabet))

	for num > 0 {
		rem := num % base
		num = num / base
		arr = append(arr, base62Alphabet[rem])
	}

	for i, j := 0, len(arr)-1; i < j; i, j = i+1, j-1 {
		arr[i], arr[j] = arr[j], arr[i]
	}

	return string(arr)
}

type Counter struct {
	Count int64
}

// getCount return the current count of urls.
func getCount(ctx appengine.Context) (int64, os.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, os.Error) {
	var count int64

	err := datastore.RunInTransaction(ctx, func(ctx appengine.Context) os.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
}

type URL struct {
	Short   string
	Long    string
	User    string
	Created int64
	Hits    int64
}

// urlKey returns a datastore key for short url.
func urlKey(ctx appengine.Context, short string) *datastore.Key {
	return datastore.NewKey(ctx, urlKind, short, 0, nil)
}

// getURL fetches a URL from the datastore by short url.
func getURL(ctx appengine.Context, short string) (*URL, os.Error) {
	key := urlKey(ctx, short)
	url := new(URL)
	err := datastore.Get(ctx, key, url)

	return url, err
}

// incHits increments hit count on url (this is done when short url is resolved).
func incHits(ctx appengine.Context, short string) os.Error {
	return datastore.RunInTransaction(ctx, func(ctx appengine.Context) os.Error {
		url, err := getURL(ctx, short)
		if err != nil {
			return err
		}

		url.Hits++
		key := urlKey(ctx, short)
		if _, err := datastore.Put(ctx, key, url); err != nil {
			return err
		}
		return nil
	}, nil)
}

// userURLs return a list of user urls, ordered by time (max 1000).
func userURLs(ctx appengine.Context, user string) ([]*URL, os.Error) {
	query := datastore.NewQuery(urlKind).
		Filter("User =", user).
		Order("-Created").
		Limit(1000)
	buff := []*URL{}
	it := query.Run(ctx)

	for {
		var url URL
		_, err := it.Next(&url)
		if err == datastore.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		buff = append(buff, &url)
	}

	return buff, nil
}

var homeTemplate *template.Template

func init() {
	var fmap = template.FuncMap{
		"strftime": timeFormatter,
	}
	// homeHTML is in html.go
	homeTemplate = template.Must(template.New("home").Funcs(fmap).Parse(homeHTML))
	http.HandleFunc(hitWorkerPath, hitHandler)
	http.HandleFunc("/", rootHandler)
}

// Parameters for homeTemplate.
type homeParams struct {
	User       string
	LoginTitle string
	LoginURL   string
	Count      int64
	Error      string
	ShortURL   string
	URLs       []*URL
}

/* 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)
}

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

// rootHandler handles the main page.
func rootHandler(w http.ResponseWriter, r *http.Request) {
	// If it has something, we assume it's a short url
	if r.URL.Path != "/" {
		redirectHandler(w, r)
		return
	}

	ctx := appengine.NewContext(r)
	var err os.Error
	params := &homeParams{
		User: "stranger",
	}

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

	u := user.Current(ctx)

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

		params.URLs, err = userURLs(ctx, params.User)
		if err != nil {
			return
		}
	} else {
		params.LoginTitle = "Login"
		params.LoginURL, err = user.LoginURL(ctx, r.URL.String())
	}

	if r.Method == "POST" {
		longURL := strings.TrimSpace(r.FormValue("url"))
		if len(longURL) == 0 {
			err = os.NewError("Empty URL")
			return
		}

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

		if !strings.Contains(longURL, ".") {
			err = os.NewError(fmt.Sprintf("Bad URL - %s", longURL))
			return
		}

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

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

	params.Count, err = getCount(ctx)
}

/* redirectHandler handles redirects.
All urls that are not / and worker are assumed to be redirects (short).
*/
func redirectHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	path := r.URL.Path[1:]

	var longURL string

	// Try memcache first and if not get URL from datastore and updatge memcache
	if item, err := memcache.Get(ctx, path); err == memcache.ErrCacheMiss {
		url, err1 := getURL(ctx, path)
		if err1 != nil {
			http.NotFound(w, r)
			return
		}
		longURL = url.Long

		item1 := &memcache.Item{
			Key:   path,
			Value: []byte(longURL),
		}

		if err1 := memcache.Set(ctx, item1); err1 != nil {
			ctx.Errorf("memcache setting error: %v", err1)
		}
	} else if err != nil {
		ctx.Errorf("memcache error - %s", err)
		http.Error(w, err.String(), http.StatusInternalServerError)
		return
	} else { // Found in memcache
		longURL = string(item.Value)
	}

	// Offload hit counter update to a task
	task := taskqueue.NewPOSTTask(hitWorkerPath, map[string][]string{"path": {path}})
	if _, err := taskqueue.Add(ctx, task, ""); err != nil {
		return
	}

	http.Redirect(w, r, longURL, http.StatusTemporaryRedirect)
}

// timeFormatter formats time. It is used by homeTemplate.
func timeFormatter(args ...interface{}) string {
	nsec, ok := args[0].(int64)
	if !ok {
		return "ERR"
	}
	t := time.NanosecondsToLocalTime(nsec)
	return t.Format("2006-01-02 15:04")
}

// hitHandler is handler for updating hits task.
func hitHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	path := r.FormValue("path")
	if len(path) == 0 {
		ctx.Errorf("Empty value for hit handler")
		return
	}
	incHits(ctx, path)
}