#46 - RunMiddleware for GetRaw and SetRaw

command.go

@@ -5,6 +5,9 @@ type Command struct {
  // Key information
  bucket string
  bucketType string
+
+ // Middleware
+ runMiddleware []RunMiddleware
 }
 
 // Result contains your query result data from Run()
@@ -21,3 +24,8 @@ func Bucket(bucket, bucketType string) *Command {
  bucketType: bucketType,
  }
 }
+
+func (c *Command) RegisterRunMiddleware(middleware RunMiddleware) *Command {
+ c.runMiddleware = append(c.runMiddleware, middleware)
+ return c
+}

middleware.go

@@ -1,17 +1,12 @@
 package goriak
 
-type ExecMiddlewarer interface {
+type RunMiddlewarer interface {
  Key() string
  Bucket() string
  BucketType() string
 }
 
-type ExecuteMiddleware func(cmd ExecMiddlewarer, next func() (*Result, error)) (*Result, error)
-
-func (c *SetRawCommand) AddMiddleware(mid ExecuteMiddleware) *SetRawCommand {
- c.execMiddleware = append(c.execMiddleware, mid)
- return c
-}
+type RunMiddleware func(cmd RunMiddlewarer, next func() (*Result, error)) (*Result, error)
 
 type setRawMiddlewarer struct {
  cmd *SetRawCommand
@@ -28,3 +23,19 @@ func (c setRawMiddlewarer) Bucket() string {
 func (c setRawMiddlewarer) BucketType() string {
  return c.cmd.c.bucketType
 }
+
+type getRawMiddlewarer struct {
+ cmd *GetRawCommand
+}
+
+func (c getRawMiddlewarer) Key() string {
+ return c.cmd.key
+}
+
+func (c getRawMiddlewarer) Bucket() string {
+ return c.cmd.c.bucket
+}
+
+func (c getRawMiddlewarer) BucketType() string {
+ return c.cmd.c.bucketType
+}

middleware_test.go

@@ -6,13 +6,13 @@ import (
 )
 
 func TestMiddlewareSetAbort(t *testing.T) {
- m := func(cmd ExecMiddlewarer, next func() (*Result, error)) (*Result, error) {
+ m := func(cmd RunMiddlewarer, next func() (*Result, error)) (*Result, error) {
  return nil, errors.New("aborted middleware")
  }
 
  _, err := Bucket("middleware", "tests").
+ RegisterRunMiddleware(m).
  SetRaw([]byte{1, 2, 3, 4, 5}).
- AddMiddleware(m).
  Run(con())
  if err == nil {
  t.Error("no error")
@@ -27,7 +27,7 @@ func TestMiddlewareSetKeyTest(t *testing.T) {
  beforeEmpty := 0
  afterWithKey := 0
 
- m := func(cmd ExecMiddlewarer, next func() (*Result, error)) (*Result, error) {
+ m := func(cmd RunMiddlewarer, next func() (*Result, error)) (*Result, error) {
  if cmd.Key() == "" {
  beforeEmpty++
  }
@@ -42,9 +42,9 @@ func TestMiddlewareSetKeyTest(t *testing.T) {
  }
 
  _, err := Bucket("middleware", "tests").
+ RegisterRunMiddleware(m).
+ RegisterRunMiddleware(m).
  SetRaw([]byte{1, 2, 3, 4, 5}).
- AddMiddleware(m).
- AddMiddleware(m).
  Run(con())
  if err != nil {
  t.Error(err)
@@ -60,7 +60,7 @@ func TestMiddlewareSetKeyTest(t *testing.T) {
 }
 
 func TestMiddlewareSetBucketType(t *testing.T) {
- m := func(cmd ExecMiddlewarer, next func() (*Result, error)) (*Result, error) {
+ m := func(cmd RunMiddlewarer, next func() (*Result, error)) (*Result, error) {
  if cmd.Bucket() != "middleware" {
  t.Error("unexpected before bucket")
  }
@@ -81,10 +81,48 @@ func TestMiddlewareSetBucketType(t *testing.T) {
  }
 
  _, err := Bucket("middleware", "tests").
+ RegisterRunMiddleware(m).
  SetRaw([]byte{1, 2, 3, 4, 5}).
- AddMiddleware(m).
  Run(con())
  if err != nil {
  t.Error(err)
  }
 }
+
+func TestMiddlewareGetRaw(t *testing.T) {
+ m := func(cmd RunMiddlewarer, next func() (*Result, error)) (*Result, error) {
+ if cmd.Bucket() != "middleware" {
+ t.Error("unexpected before bucket")
+ }
+ if cmd.BucketType() != "tests" {
+ t.Error("unexpected before bucket type")
+ }
+ if cmd.Key() != "hello123" {
+ t.Error("unexpected before key")
+ }
+
+ res, err := next()
+
+ if cmd.Bucket() != "middleware" {
+ t.Error("unexpected after bucket")
+ }
+ if cmd.BucketType() != "tests" {
+ t.Error("unexpected after bucket type")
+ }
+ if cmd.Key() != "hello123" {
+ t.Error("unexpected after key")
+ }
+
+ if err == nil {
+ t.Error("did not get error")
+ }
+
+ return res, err
+ }
+
+ var out []byte
+ Bucket("middleware", "tests").
+ RegisterRunMiddleware(m).
+ GetRaw("hello123", &out).
+ Run(con())
+}

query_get_value.go

@@ -86,7 +86,37 @@ func (c *GetRawCommand) WithR(r uint32) *GetRawCommand {
  return c
 }
 
+func runMiddleware(middlewarer RunMiddlewarer, middlewareList []RunMiddleware, execFunc func(*Session) (*Result, error), session *Session) (*Result, error) {
+ // Keep track of whick middleware that we should execute next
+ middlewareI := 0
+
+ // Needed so that next can call itself
+ var next2 func() (*Result, error)
+
+ next := func() (*Result, error) {
+ if middlewareI == len(middlewareList) {
+ return execFunc(session)
+ }
+
+ middlewareI++
+
+ return middlewareList[middlewareI-1](middlewarer, next2)
+ }
+
+ next2 = next
+
+ return next()
+}
+
 func (c *GetRawCommand) Run(session *Session) (*Result, error) {
+ middlewarer := &getRawMiddlewarer{
+ cmd: c,
+ }
+
+ return runMiddleware(middlewarer, c.c.runMiddleware, c.runExec, session)
+}
+
+func (c *GetRawCommand) runExec(session *Session) (*Result, error) {
  cmd, err := c.builder.Build()
  if err != nil {
  return nil, err

query_store_value.go

@@ -14,8 +14,6 @@ type SetRawCommand struct {
  storeValueCommandBuilder *riak.StoreValueCommandBuilder
  storeValueObject *riak.Object
 
- execMiddleware []ExecuteMiddleware
-
  key string
 
  err error
@@ -64,29 +62,11 @@ func (c *SetRawCommand) Run(session *Session) (*Result, error) {
  // Set object
  c.storeValueCommandBuilder.WithContent(c.storeValueObject)
 
- // Keep track of whick middleware that we should execute next
- middlewareI := 0
-
- // Needed so that next can call itself
- var next2 func() (*Result, error)
-
  middlewarer := setRawMiddlewarer{
  cmd: c,
  }
 
- next := func() (*Result, error) {
- if middlewareI == len(c.execMiddleware) {
- return c.riakExecute(session)
- }
-
- middlewareI++
-
- return c.execMiddleware[middlewareI-1](middlewarer, next2)
- }
-
- next2 = next
-
- return next()
+ return runMiddleware(middlewarer, c.c.runMiddleware, c.riakExecute, session)
 }
 
 func (c *SetRawCommand) riakExecute(session *Session) (*Result, error) {

raw.go

@@ -27,6 +27,7 @@ func (c *Command) SetRaw(value []byte) *SetRawCommand {
 // The output will be written to output by Run().
 func (c *Command) GetRaw(key string, output *[]byte) *GetRawCommand {
  cmd := &GetRawCommand{
+ c: c,
  key: key,
  outputBytes: output,
  isRawOutput: true,