Implementing pluggable backends in go

When writing a component like a storage backend you often want a way to switch between different implementations like memory, sqlite, or redis. Another common example is authentication backends for LDAP or Kerberos.

With interfaces in go you can have multiple implementations of a common API, but it takes a little more work to make them pluggable.¹

The goal is that we want to be able to have a configuration file with a section like

"store" {
    "type": "boltdb",
    "options": {
        "filename": "my.db"
    }
}

to create an arbitrary backend with the appropriate options.

Let’s write a simple frontend library for a key value store and work up to pluggable backends. The first step is to decide on the interface:

type KVStore interface {
    Get(key string) (string, error)
    Set(key, value string) error
    Delete(key string) (error)
}

var NoSuchKey = errors.New("No such key")

First interface implementation

With the interface decided on, let’s write an in-memory implementation using a map:

type MemoryKVStore struct {
    store map[string]string
}

func NewMemoryKVStore() *MemoryKVStore {
    store := make(map[string]string)
    return &MemoryKVStore{store: store}
}

func (m *MemoryKVStore) Get(key string) (string, error) {
    val, ok := m.store[key]
    if !ok {
        return "", NoSuchKey
    }
    return val, nil
}
func (m *MemoryKVStore) Set(key string, value string) error {
    m.store[key] = value
    return nil
}

func (m *MemoryKVStore) Delete(key string) error {
    delete(m.store, key)
    return nil
}

At this point we can write a main function and do an initial test:

func main() {
    kv := NewMemoryKVStore()

    fmt.Println(kv.Get("test"))
    fmt.Println(kv.Set("test", "success"))
    fmt.Println(kv.Get("test"))
    fmt.Println(kv.Delete("test"))
    fmt.Println(kv.Get("test"))
}

Which outputs the expected:

 No such key
<nil>
success <nil>
<nil>
 No such key

A real test case would be better, so let’s write one. I can plan ahead for multiple implementations by using the sub tests feature in go 1.7. Instead of testing the MemoryKVStore directly, I can test the KVStore interface:

import "testing"

func storeTestHelper(t *testing.T, store KVStore) {
    if _, err := store.Get("test"); err != NoSuchKey {
        t.Fatalf("Expected NoSuchKey, got %q", err)
    }
    if err := store.Set("test", "success"); err != nil {
        t.Fatalf("Expected nil, got %q", err)
    }
    if val, _ := store.Get("test"); val != "success" {
        t.Fatalf("Expected success, got %q", val)
    }
    if err := store.Delete("test"); err != nil {
        t.Fatalf("Expected nil, got %q", err)
    }
    _, err := store.Get("test")
    if err != NoSuchKey {
        t.Fatalf("Expected NoSuchKey, got %q", err)
    }
}

func TestKVStore(t *testing.T) {
    t.Run("backend=MemoryStore", func(t *testing.T) {
        kv := NewMemoryKVStore()
        storeTestHelper(t, kv)
    })
}

2nd interface implementation

One copy & paste, search & replace, and a few lines of code later, a 2nd Implementation is ready:

type BoltDBStore struct {
    db *bolt.DB
}

var bucketID = []byte("storage")

func NewBoltStore(filename string) (*BoltDBStore, error) {
    db, err := bolt.Open(filename, 0600, nil)
    if err != nil {
        return nil, err
    }
    err = db.Update(func(tx *bolt.Tx) error {
        _, err = tx.CreateBucketIfNotExists(bucketID)
        if err != nil {
            return fmt.Errorf("create bucket: %s", err)
        }
        return nil
    })
    if err != nil {
        return nil, err
    }

    return &BoltDBStore{db: db}, nil
}

func (b *BoltDBStore) Get(key string) (string, error) {
    var val []byte
    err := b.db.View(func(tx *bolt.Tx) error {
        bucket := tx.Bucket(bucketID)
        val = bucket.Get([]byte(key))
        return nil
    })
    if err != nil {
        return "", err
    }
    if val == nil {
        return "", NoSuchKey
    }
    return string(val), nil

}
func (b *BoltDBStore) Set(key string, value string) error {
    return b.db.Update(func(tx *bolt.Tx) error {
        bucket := tx.Bucket(bucketID)
        err := bucket.Put([]byte(key), []byte(value))
        return err
    })
}

func (b *BoltDBStore) Delete(key string) error {
    return b.db.Update(func(tx *bolt.Tx) error {
        bucket := tx.Bucket(bucketID)
        err := bucket.Delete([]byte(key))
        return err
    })
}

and hooked up to the tests:

t.Run("backend=BoltDBStore", func(t *testing.T) {
    dir, err := ioutil.TempDir("", "bolttest")
    defer os.RemoveAll(dir) // clean up
    kv, err := NewBoltStore(dir + "/db")
    if err != nil {
        t.Fatal(err)
    }
    storeTestHelper(t, kv)
})

Pluggable problems

At this point, we already start to run into trouble making these backends plugable. The new BoltDBStore implements the same interface as MemoryKVStore, but the NewBoltStore has a different signature from NewMemoryKVStore.

Changing MemoryKVStore to return a nil error is easy, but it doesn’t make sense to add an unused filename parameter.

The solution is.. more types

Option types

type MemStoreOptions struct {
}

func NewMemoryKVStore(options MemStoreOptions) (*MemoryKVStore, error) {
    store := make(map[string]string)
    return &MemoryKVStore{store: store}, nil
}


type BoltStoreOptions struct {
    filename string
}

func NewBoltStore(options BoltStoreOptions) (*BoltDBStore, error) {
    db, err := bolt.Open(options.filename, 0600, nil)
    ...

Now, these two New functions take almost the same interface, aside from MemStoreOptions and BoltStoreOptions not being the same type.

To solve this problem, we can use a type assertion.

Generic NewStore

func NewStore(backend string, options interface{}) (KVStore, error) {
    switch backend {
    case "memory":
        opts, ok := options.(MemStoreOptions)
        if !ok {
            return nil, fmt.Errorf("Invalid memory options %q", options)
        }
        return NewMemoryKVStore(opts)
    case "boltdb":
        opts, ok := options.(BoltStoreOptions)
        if !ok {
            return nil, fmt.Errorf("Invalid Bolt options %q", options)
        }
        return NewBoltStore(opts)
    default:
        return nil, fmt.Errorf("Unknown store: %s", backend)
    }
}

With this written, the tests can be updated:

func TestKVStore(t *testing.T) {
    t.Run("backend=MemoryStore", func(t *testing.T) {
        kv, _ := NewStore("memory", MemStoreOptions{})
        storeTestHelper(t, kv)
    })
    t.Run("backend=BoltDBStore", func(t *testing.T) {
        dir, err := ioutil.TempDir("", "bolttest")
        defer os.RemoveAll(dir) // clean up
        kv, err := NewStore("boltdb", BoltStoreOptions{filename: dir + "/db"})
        if err != nil {
            t.Fatal(err)
        }
        storeTestHelper(t, kv)
    })
}

This is almost everything we need to implement our config file.

Configuration file

First, define a struct to hold the configuration:

type StoreConfig struct {
    Backend string                 `json:"backend"`
    Options map[string]interface{} `json:"options"`
}

type Configuration struct {
    Store StoreConfig `json:"store"`
}

Second, define two test configuration files:

var testMemConfig = []byte(`
{
    "store": {
        "backend": "memory"
    }
}
`)

var testBoltConfig = []byte(`
{
    "store": {
        "backend": "boltdb",
        "options": {
            "filename": "config_test_tempory.db"
        }
    }
}
`)

Now, we write a function that can take an io.Reader for a configuration, and return a store:

NewStoreFromConfig(r io.Reader) (KVStore, error) {
    var cfg Configuration 
    err := json.NewDecoder(r).Decode(&cfg)
    if err != nil {
        return nil, err
    }
    return NewStore(cfg.Store.Backend, cfg.Store.Options)
}

The problem is, this does not work! cfg.Store.Options is a map, but we need an Interface{}:

--- FAIL: TestKVStore (0.00s)
    --- FAIL: TestKVStore/backend=MemoryStore (0.00s)
        kv_test.go:51: Invalid memory options map[]
    --- FAIL: TestKVStore/backend=BoltDBStore (0.00s)
        kv_test.go:60: Invalid Bolt options map["filename":"config_test_tempory.db"]

The easy way to fix this is to update the store backends to have a function that converts the options into the right interface and return a store:

    func NewMemoryStoreFromMap(options map[string]interface{}) (*MemoryKVStore, error) {
        //Nothing to do here
        opts := MemStoreOptions{}
        return NewMemoryKVStore(opts)
    }

    func NewBoltStoreFromMap(options map[string]interface{}) (*BoltDBStore, error) { 
        opts := BoltStoreOptions{}
        opts.filename = options["filename"].(string)
        return NewBoltStore(opts)
    }

A lot of fancy reflection can do this in a generic way, but this keeps things simple.

With those two functions in place, NewStoreFromConfig can be rewritten:

func NewStoreFromConfig(r io.Reader) (KVStore, error) {
    var cfg Configuration
    err := json.NewDecoder(r).Decode(&cfg)
    if err != nil {
        return nil, err
    }
    switch cfg.Store.Backend {
    case "memory":
        return NewMemoryStoreFromMap(cfg.Store.Options)
    case "boltdb":
        return NewBoltStoreFromMap(cfg.Store.Options)
    default:
        return nil, fmt.Errorf("Unknown store: %s", cfg.Store.Backend)
    }
}

And the tests can be updated to use NewStoreFromConfig

func TestKVStore(t *testing.T) {
    t.Run("backend=MemoryStore", func(t *testing.T) {
        kv, err := NewStoreFromConfig(bytes.NewReader(testMemConfig))
        if err != nil {
            t.Fatal(err)
        }
        storeTestHelper(t, kv)
    })
    t.Run("backend=BoltDBStore", func(t *testing.T) {
        os.RemoveAll("config_test_tempory.db")       // clean up
        defer os.RemoveAll("config_test_tempory.db") // clean up
        kv, err := NewStoreFromConfig(bytes.NewReader(testBoltConfig))
        if err != nil {
            t.Fatal(err)
        }
        storeTestHelper(t, kv)
    })
}

One more thing

This is fully functional but it can still be improved. NewStoreFromConfig has to have knowledge about each store type. It would be better if each store was 100% self contained. This problem can be solved by having store backends register themselves into a map.

First, to simplify things, write a NewStoreFromInterface for each backend:

func NewBoltStoreFromInterface(options interface{}) (*BoltDBStore, error) {
    opts, ok := options.(BoltStoreOptions)
    if !ok {
        return nil, fmt.Errorf("Invalid boltdb options %q", options)
    }
    return NewBoltStore(opts)
}

func NewMemoryKVStoreFromInterface(options interface{}) (*MemoryKVStore, error) {
    opts, ok := options.(MemStoreOptions)
    if !ok {
        return nil, fmt.Errorf("Invalid memory options %q", options)
    }
    return NewMemoryKVStore(opts)
}

This moves the type assertions to the backends and simplifies the NewStore function to:

func NewStore(backend string, options interface{}) (KVStore, error) {
    switch backend {
    case "memory":
        return NewMemoryKVStoreFromInterface(options)
    case "boltdb":
        return NewBoltStoreFromInterface(options)
    default:
        return nil, fmt.Errorf("Unknown store: %s", backend)
    }
}

At this point, NewStore is effectively a map from

memory -> NewMemoryKVStoreFromInterface
boltdb -> NewBoltStoreFromInterface

and NewStoreFromConfig is a map from:

memory -> NewMemoryStoreFromMap
boltdb -> NewBoltStoreFromMap

Having two functions like this is a little awkward. A grouping of similar functions is exactly what an interface is for, so let’s build another interface for creating stores:

type KVStoreDriver interface {
    NewFromInterface(interface{}) (KVStore, error)
    NewFromMap(map[string]interface{}) (KVStore, error)
}

And create a map of strings to Drivers and a function to manage them:

var drivers = make(map[string]KVStoreDriver)

func Register(backend string, driver KVStoreDriver) {
    drivers[backend] = driver
}

And then move our New functions into this interface:

type MemoryStoreDriver struct {
}

func (d MemoryStoreDriver) NewFromMap(options map[string]interface{}) (KVStore, error) {
    //Nothing to do here
    opts := MemStoreOptions{}
    return NewMemoryKVStore(opts)
}
func (d MemoryStoreDriver) NewFromInterface(options interface{}) (KVStore, error) {
    opts, ok := options.(MemStoreOptions)
    if !ok {
        return nil, fmt.Errorf("Invalid memory options %q", options)
    }
    return NewMemoryKVStore(opts)
}

And the final ’trick’ that will allow this to work: init inside each driver:

func init() {
    Register("memory", MemoryStoreDriver{})
}

With that in place, NewStore and NewStoreFromConfig can be rewritten to simply look up the driver in the map:

func NewStore(backend string, options interface{}) (KVStore, error) {
    driver, ok := drivers[backend]
    if !ok {
        return nil, fmt.Errorf("Unknown store: %s", backend)
    }
    return driver.NewFromInterface(options)
}

func NewStoreFromConfig(r io.Reader) (KVStore, error) {
    var cfg Configuration
    err := json.NewDecoder(r).Decode(&cfg)
    if err != nil {
        return nil, err
    }
    driver, ok := drivers[cfg.Store.Backend]
    if !ok {
        return nil, fmt.Errorf("Unknown store: %s", cfg.Store.Backend)
    }
    return driver.NewFromMap(cfg.Store.Options)
}

At this stage the backends are fully pluggable and the core of the kvstore library has zero knowledge of any individual backend implementation. Adding a new backend is simply a matter of importing a package that uses Register to add a new KVStoreDriver that implements the right interfaces.

Justin Azoff