Skip to content
Turtle
Instagram

Key Generator Service with DynamoDB

— DynamoDB, Key Generation Service, Golang — 3 min read

In this post, we'll explore how DynamoDB works by building a key generation service (KGS) using Go.

Background

A key generation service is useful for decoupling unique key creation from any other business logic. This service comes up in System Design questions such as URL shortner and Pastebin. It works by generating random base64 keys and storing them in an available_keys table. Once a key is requested to be used, that key is then moved to the used_keys table.

Setup

For our storage, we will use Amazon's DynamoDB, a NoSQL storage. To run it locally, type:

docker run -p 8000:8000 amazon/dynamodb-local

Generating the keys

To generate the keys, we'll use the uuid package to generate the random string and base64 encode it.

import (
    "encoding/base64"


    "github.com/google/uuid"
)


hash := uuid.New()
encodedHash := base64.StdEncoding.EncodeToString(hash[:])

Saving keys into DynamoDB

DynamoDB consists of tables, items and indices. Each table does not need a schema like in relational DBs but does require a primary key. This can be any unique attribute of the item. If two attributes are used, one will be the partition key and the other the sort key.

For our example, we will use the key as the primary key.

type KeyItem struct {
        Key       string `dynamodbav:"key"`
        CreatedAt string
        UpdatedAt string
    }


now := time.Now().UTC().String()


item := KeyItem{
    Key:       encodedHash,
    CreatedAt: now,
    UpdatedAt: now,
}

To create the tables and save the key in the available_keys table we will use the aws-sdk-go-V2 package. This package is the successor to the original AWS sdk. We assume that the client is already initialized and the used_keys table is created identically to the available_keys table.

// Create the table
client.CreateTable(ctx, &dynamodb.CreateTableInput{
        TableName: aws.String(TableAvailableKeys),
        AttributeDefinitions: []types.AttributeDefinition{
            {
                AttributeName: aws.String("key"),
                AttributeType: types.ScalarAttributeTypeS,
            },
        },
        KeySchema: []types.KeySchemaElement{
            {
                AttributeName: aws.String("key"),
                KeyType:       types.KeyTypeHash,
            },
        },
        ProvisionedThroughput: &types.ProvisionedThroughput{
            ReadCapacityUnits:  aws.Int64(10),
            WriteCapacityUnits: aws.Int64(10),
        },
    }
)
// Save the item
marshalledItem, err := attributevalue.MarshalMap(item)
if err != nil {
    panic(err)
}


name := "available_keys"
_, err = client.PutItem(ctx, &dynamodb.PutItemInput{
    Item:      marshalledItem,
    TableName: &name,
})
if err != nil {
    panic(err)
}

As you can see, creating tables and adding items is pretty straight forward in DynamoDB.

Accessing the keys

When a key is requested from the service, we retrieve the first available key and move it to the used_keys table. This allows us to reuse keys when they are no longer needed.

First, scan the available_keys table for a key:

var limit int32 = 1
res, err := k.client.Scan(ctx, &dynamodb.ScanInput{
    TableName: aws.String("available_keys"),
    Limit:     &limit,
})
if err != nil {
    panic(err)
}

Next, save that key in the used_keys table:

item := tables.KeyItem{}
err = attributevalue.UnmarshalMap(res.Items[0], &item)
if err != nil {
    panic(err)
}


item.CreatedAt = time.Now().UTC().String()
item.UpdatedAt = item.CreatedAt


marshalledItem, err := attributevalue.MarshalMap(item)
if err != nil {
    panic(err)
}
_, err = k.client.PutItem(ctx, &dynamodb.PutItemInput{
    Item:      marshalledItem,
    TableName: aws.String("used_keys"),
})
if err != nil {
    panic(err)
}

Finally, delete the key from the available_keys table:

key, err := attributevalue.Marshal(item.Key)
if err != nil {
    panic(err)
}


_, err = k.client.DeleteItem(ctx, &dynamodb.DeleteItemInput{
    Key: map[string]types.AttributeValue{
        "key": key,
    },
    TableName: aws.String("available_keys"),
})
if err != nil {
    panic(err)
}

Note: deleting the key involves doing the inverse approach. We get the key from the used_keys table, delete it from there and add it back to the available_keys table.

It's easy to play any musical instrument: all you have to do is touch the right key at the right time and the instrument will play itself. - Johann Sebastian Bach