How to write a URL Shortener in Go

Exercise 2 of Gophercises was to create a URL shortener. The framework for this code was given and the student needed to fill in the missing pieces. The solution of the main problem is given but no solutions are given for the bonus challenges. This is my solution to the main problem along with the bonus challenges in which a solution was not given.

The main goal of this exercise was to get familiar working with different markup languages in Go. It also introduced working with command line arguments and reading data from a database.

The URL shortener

The first step of the problem was to create a map handler. The map handler needs to check a map for a key with the specified URL and redirect the user to the value stored in that key. If no key is found then the fallback handler should be used instead.

return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
    if pathsToUrls[r.URL.Path] != "" {
        http.Redirect(rw, r, pathsToUrls[r.URL.Path], http.StatusFound)
    }

    fallback.ServeHTTP(rw, r)
})

The next step was creating a handler for YAML input. This function will parse YAML input and build a map from the results. The map is then passed to the map handler built in the previous step.

var redirects []struct {
    Path string `yaml:"path"`
    URL  string `yaml:"url"`
}

err := yaml.Unmarshal([]byte(yml), &redirects)
if err != nil {
    return nil, err
}

paths := make(map[string]string)
for _, path := range redirects {
    paths[path.Path] = path.URL
}

return MapHandler(paths, fallback), nil

Bonus #1: Read the YAML from a file

The given code creates an example YAML string inside the program. This bonus challenge was to instead read the YAML from a file. The file path would be passed in through the command line and have a default value.

The first step was to read the file path from the command line. This is done using the Golang flag package. The flag package makes it very easy to setup command line arguments.

yamlPath := flag.String("YAML", "paths.yaml", "A YAML file")
flag.Parse()

Here the library is setup so it knows it will be receiving a String argument via the command line. That argument will be referred to as “YAML” with a default value of “paths.yaml” and a description of “A YAML file”. The value of this argument will be stored in yamlPath variable. The flag.Parse() function is then called which will process the command line arguments and store them in the associated variables.

One of the great things about using the flag library is that with this one line the program will now have an automatically generated help menu. Passing the --help flag into your program will print a menu for someone trying to figure out the available flags.

Usage of /tmp/go-build244514782/b001/exe/main:
  -YAML string
        A YAML file containing redirects (default "paths.yaml")

The file path now needs to be passed to the YAML handler and have it open and read from the file. There are a few different options available when it comes to reading input from a file.

reader, err := os.Open(*yamlPath)
if err != nil {
    panic(err)
}

err = yaml.NewDecoder(reader).Decode(&redirects)
if err != nil {
    return nil, err
}

Creating an io.reader allows this reader to be passed to the YAML library. This allows the YAML library to manage the buffering of reading from the file. This is then passed to a new YAML decoder which decodes the file into our redirects variable.

Other options include reading the entire file into memory or reading in the file one line at a time. Reading the entire file into memory would increase the memory usage of the program, especially if there are a large number of paths. Reading in the file one line at a time prevents this issue. Luckily, this is similar to how a decoder is buffering the file without having to manually manage the file’s location pointer.

That’s it for handling the YAML input being stored in a file.

Bonus #2: Build a JSON Handler

The next bonus question was very similar to the main challenge but slightly different. This time data will be passed in using JSON instead of the original YAML.

reader, err := os.Open(*jsonPath)
if err != nil {
    panic(err)
}

err = json.NewDecoder(reader).Decode(&redirects)
if err != nil {
    return nil, err
}

Not much changes here other than using the JSON library instead of the YAML library. These libraries are setup to use the same interface so switching between the two is very convenient.

Bonus 3: Build a database Handler

The final bonus challenge was to store paths in a database and access them from the program. BoltDB is recommended in the exercise but I chose to use MongoDB instead. MongoDB is a name I had heard of in production systems and I already had it set up on my system. Installing and setting up MongoDb instructions can be found on the official site.

var db *mongo.Client

collection := db.Database("test").Collection("Paths")
cur, err := collection.Find(context.TODO(), bson.M{})
if err != nil {
    log.Panic(err)
}

db is used as a global variable which has package scope. This assumes db has been initialized in another function. The code first gets access to the “Paths” collection on the “test” database. A search is then ran on that collection using an empty interface. The empty interface causes all rows to be retrieved from the collection. If there were values in the interface, then the results would be filtered by those values. Context.TODO() is an empty context that is used when a context is not available.

Converting the other Handler’s data into a map was pretty straightforward. This process is a little more involved when working with a database.

paths := make(map[string]string)

for cur.Next(context.TODO()) {
    var redirect Redirect
    err := cur.Decode(&redirect)
    if err != nil {
        log.Panic(err)
    }

    paths[redirect.Path] = redirect.URL
}

Collection.Find() in the previous code returned a cursor to the row in the collection. This cursor can now be used to iterator over the rows in the collection and add their value to a map. A struct is created on each iteration to be added to the map. The path is then used as a key in the map with the URL as the value. The map can then be passed into the MapHandler that was first created.

Wrapping up

This Gophercise was a great way to get used to working with different data storage formats. The bonus challenges help to slowly introduce additional libraries that are nice to know about and have some experience working with.