make it possible to wait for a query lambda

errors/wait.go

@@ -0,0 +1,25 @@
+package errors
+
+import (
+	"errors"
+	"fmt"
+)
+
+var ErrBadWaitState = errors.New("encountered bad state while waiting for resource")
+
+type BadWaitState struct {
+	State string
+}
+
+func (e BadWaitState) Error() string {
+	return fmt.Sprintf("%s: %s", ErrBadWaitState.Error(), e.State)
+}
+func (e BadWaitState) Unwrap() error {
+	return ErrBadWaitState
+}
+
+func NewBadWaitStateError(state string) BadWaitState {
+	return BadWaitState{
+		State: state,
+	}
+}

option/collection.go

@@ -7,10 +7,14 @@ import (
 	"github.com/rockset/rockset-go-client/openapi"
 )
 
+type CollectionStatus string
+
+func (c CollectionStatus) String() string { return string(c) }
+
 const (
-	CollectionStatusCreated     = "CREATED"
-	CollectionStatusInitialized = "INITIALIZED"
-	CollectionStatusReady       = "READY"
+	CollectionStatusCreated     CollectionStatus = "CREATED"
+	CollectionStatusInitialized CollectionStatus = "INITIALIZED"
+	CollectionStatusReady       CollectionStatus = "READY"
 )
 
 type ListCollectionOptions struct {

option/query.go

@@ -4,6 +4,8 @@ import "github.com/rockset/rockset-go-client/openapi"
 
 type QueryState string
 
+func (q QueryState) String() string { return string(q) }
+
 const (
 	QueryQueued    QueryState = "QUEUED"
 	QueryRunning   QueryState = "RUNNING"

option/query_lambda.go

@@ -2,6 +2,15 @@ package option
 
 import "github.com/rockset/rockset-go-client/openapi"
 
+type QueryLambdaState string
+
+func (q QueryLambdaState) String() string { return string(q) }
+
+const (
+	QueryLambdaActive  QueryLambdaState = "ACTIVE"
+	QueryLambdaInvalid QueryLambdaState = "INVALID"
+)
+
 type ExecuteQueryLambdaRequest struct {
 	openapi.ExecuteQueryLambdaRequest
 	Tag     string

option/virtual_instance.go

@@ -2,15 +2,19 @@ package option
 
 import "time"
 
+type VirtualInstanceState string
+
+func (v VirtualInstanceState) String() string { return string(v) }
+
 const (
-	VirtualInstanceInitializing           = "INITIALIZING"
-	VirtualInstanceProvisioningResources  = "PROVISIONING_RESOURCES"
-	VirtualInstanceRebalancingCollections = "REBALANCING_COLLECTIONS"
-	VirtualInstanceActive                 = "ACTIVE"
-	VirtualInstanceSuspending             = "SUSPENDING"
-	VirtualInstanceSuspended              = "SUSPENDED"
-	VirtualInstanceResuming               = "RESUMING"
-	VirtualInstanceDeleted                = "DELETED"
+	VirtualInstanceInitializing           VirtualInstanceState = "INITIALIZING"
+	VirtualInstanceProvisioningResources  VirtualInstanceState = "PROVISIONING_RESOURCES"
+	VirtualInstanceRebalancingCollections VirtualInstanceState = "REBALANCING_COLLECTIONS"
+	VirtualInstanceActive                 VirtualInstanceState = "ACTIVE"
+	VirtualInstanceSuspending             VirtualInstanceState = "SUSPENDING"
+	VirtualInstanceSuspended              VirtualInstanceState = "SUSPENDED"
+	VirtualInstanceResuming               VirtualInstanceState = "RESUMING"
+	VirtualInstanceDeleted                VirtualInstanceState = "DELETED"
 
 	MountCreating             = "CREATING"
 	MountActive               = "ACTIVE"

wait/collections.go

@@ -11,10 +11,10 @@ import (
 
 // UntilCollectionReady waits until the collection is ready.
 func (w *Waiter) UntilCollectionReady(ctx context.Context, workspace, name string) error {
-	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []string{option.CollectionStatusReady},
-		func(ctx context.Context) (string, error) {
+	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []option.CollectionStatus{option.CollectionStatusReady}, nil,
+		func(ctx context.Context) (option.CollectionStatus, error) {
 			c, err := w.rc.GetCollection(ctx, workspace, name)
-			return c.GetStatus(), err
+			return option.CollectionStatus(c.GetStatus()), err
 		}))
 }
 

wait/collections_test.go

@@ -15,9 +15,9 @@ func TestWait_untilCollectionReady(t *testing.T) {
 	ctx := context.TODO()
 
 	rs := fakeRocksetClient()
-	rs.GetCollectionReturnsOnCall(0, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusInitialized)}, nil)
-	rs.GetCollectionReturnsOnCall(1, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusCreated)}, nil)
-	rs.GetCollectionReturnsOnCall(2, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusReady)}, nil)
+	rs.GetCollectionReturnsOnCall(0, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusInitialized.String())}, nil)
+	rs.GetCollectionReturnsOnCall(1, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusCreated.String())}, nil)
+	rs.GetCollectionReturnsOnCall(2, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusReady.String())}, nil)
 
 	err := wait.New(&rs).UntilCollectionReady(ctx, "workspace", "collection")
 	assert.NoError(t, err)
@@ -28,7 +28,7 @@ func TestWait_untilCollectionGone(t *testing.T) {
 	ctx := context.TODO()
 
 	rs := fakeRocksetClient()
-	rs.GetCollectionReturnsOnCall(0, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusReady)}, nil)
+	rs.GetCollectionReturnsOnCall(0, openapi.Collection{Status: openapi.PtrString(option.CollectionStatusReady.String())}, nil)
 	rs.GetCollectionReturnsOnCall(1, openapi.Collection{}, NotFoundErr)
 
 	err := wait.New(&rs).UntilCollectionGone(ctx, "workspace", "collection")

wait/query.go

@@ -6,12 +6,14 @@ import (
 	"github.com/rockset/rockset-go-client/option"
 )
 
-// UntilQueryDone waits until queryID has either completed, errored, or been cancelled.
+// UntilQueryDone waits until queryID has either completed.
+// Returns ErrBadWaitState if the query failed or was cancelled.
 func (w *Waiter) UntilQueryDone(ctx context.Context, queryID string) error {
-	// TODO should this only wait for COMPLETED and return an error for ERROR and CANCELLED?
-	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []option.QueryState{option.QueryCompleted, option.QueryError, option.QueryCancelled},
-		func(ctx context.Context) (option.QueryState, error) {
-			q, err := w.rc.GetQueryInfo(ctx, queryID)
-			return option.QueryState(q.GetStatus()), err
-		}))
+	return w.rc.RetryWithCheck(ctx,
+		ResourceHasState(ctx, []option.QueryState{option.QueryCompleted},
+			[]option.QueryState{option.QueryError, option.QueryCancelled},
+			func(ctx context.Context) (option.QueryState, error) {
+				q, err := w.rc.GetQueryInfo(ctx, queryID)
+				return option.QueryState(q.GetStatus()), err
+			}))
 }

wait/query_lambda.go

@@ -0,0 +1,25 @@
+package wait
+
+import (
+	"context"
+
+	"github.com/rockset/rockset-go-client/option"
+)
+
+// UntilQueryLambdaVersionGone waits until a query lambda is deleted, i.e. GetQueryLambda() returns "not found".
+func (w *Waiter) UntilQueryLambdaVersionGone(ctx context.Context, workspace, name, version string) error {
+	return w.rc.RetryWithCheck(ctx, ResourceIsGone(ctx, func(ctx context.Context) error {
+		_, err := w.rc.GetQueryLambdaVersion(ctx, workspace, name, version)
+		return err
+	}))
+}
+
+// UntilQueryLambdaVersionActive waits until the Virtual Instance is active.
+func (w *Waiter) UntilQueryLambdaVersionActive(ctx context.Context, workspace, name, version string) error {
+	return w.rc.RetryWithCheck(ctx,
+		ResourceHasState(ctx, []option.QueryLambdaState{option.QueryLambdaActive}, []option.QueryLambdaState{option.QueryLambdaInvalid},
+			func(ctx context.Context) (option.QueryLambdaState, error) {
+				ql, err := w.rc.GetQueryLambdaVersion(ctx, workspace, name, version)
+				return option.QueryLambdaState(ql.GetState()), err
+			}))
+}

wait/query_lambda_test.go

@@ -0,0 +1,36 @@
+package wait_test
+
+import (
+	"context"
+	rockerr "github.com/rockset/rockset-go-client/errors"
+	"testing"
+
+	"github.com/rockset/rockset-go-client/openapi"
+	"github.com/rockset/rockset-go-client/option"
+	"github.com/rockset/rockset-go-client/wait"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWait_untilQueryLambdaActive(t *testing.T) {
+	ctx := context.TODO()
+
+	rs := fakeRocksetClient()
+	rs.GetQueryLambdaVersionReturnsOnCall(0, openapi.QueryLambdaVersion{State: openapi.PtrString("")}, nil)
+	rs.GetQueryLambdaVersionReturnsOnCall(1, openapi.QueryLambdaVersion{State: openapi.PtrString(option.QueryLambdaActive.String())}, nil)
+
+	err := wait.New(&rs).UntilQueryLambdaVersionActive(ctx, "ws", "ql", "v")
+	assert.NoError(t, err)
+	assert.Equal(t, 2, rs.GetQueryLambdaVersionCallCount())
+}
+
+func TestWait_untilQueryLambdaActive_invalid(t *testing.T) {
+	ctx := context.TODO()
+
+	rs := fakeRocksetClient()
+	rs.GetQueryLambdaVersionReturnsOnCall(0, openapi.QueryLambdaVersion{State: openapi.PtrString("")}, nil)
+	rs.GetQueryLambdaVersionReturnsOnCall(1, openapi.QueryLambdaVersion{State: openapi.PtrString(option.QueryLambdaInvalid.String())}, nil)
+
+	err := wait.New(&rs).UntilQueryLambdaVersionActive(ctx, "ws", "ql", "v")
+	assert.ErrorIs(t, err, rockerr.ErrBadWaitState)
+	assert.Equal(t, 2, rs.GetQueryLambdaVersionCallCount())
+}

wait/virtual_instance.go

@@ -12,10 +12,10 @@ import (
 
 // UntilVirtualInstanceActive waits until the Virtual Instance is active.
 func (w *Waiter) UntilVirtualInstanceActive(ctx context.Context, id string) error {
-	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []string{option.VirtualInstanceActive},
-		func(ctx context.Context) (string, error) {
+	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []option.VirtualInstanceState{option.VirtualInstanceActive}, nil,
+		func(ctx context.Context) (option.VirtualInstanceState, error) {
 			vi, err := w.rc.GetVirtualInstance(ctx, id)
-			return vi.GetState(), err
+			return option.VirtualInstanceState(vi.GetState()), err
 		}))
 }
 
@@ -29,20 +29,20 @@ func (w *Waiter) UntilVirtualInstanceGone(ctx context.Context, id string) error
 
 // UntilVirtualInstanceSuspended waits until the Virtual Instance is suspended.
 func (w *Waiter) UntilVirtualInstanceSuspended(ctx context.Context, id string) error {
-	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []string{option.VirtualInstanceSuspended},
-		func(ctx context.Context) (string, error) {
+	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []option.VirtualInstanceState{option.VirtualInstanceSuspended}, nil,
+		func(ctx context.Context) (option.VirtualInstanceState, error) {
 			vi, err := w.rc.GetVirtualInstance(ctx, id)
-			return vi.GetState(), err
+			return option.VirtualInstanceState(vi.GetState()), err
 		}))
 }
 
 // UntilMountActive waits until the collection mount is active, and queries can be issued to it on the
 // virtual instance.
 func (w *Waiter) UntilMountActive(ctx context.Context, vID, workspace, collection string) error {
-	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []string{option.MountActive},
-		func(ctx context.Context) (string, error) {
+	return w.rc.RetryWithCheck(ctx, ResourceHasState(ctx, []option.VirtualInstanceState{option.MountActive}, nil,
+		func(ctx context.Context) (option.VirtualInstanceState, error) {
 			cm, err := w.rc.GetCollectionMount(ctx, vID, workspace+"."+collection)
-			return cm.GetState(), err
+			return option.VirtualInstanceState(cm.GetState()), err
 		}))
 }
 

wait/virtual_instance_test.go

@@ -17,8 +17,8 @@ func TestWait_untilVirtualInstanceActive(t *testing.T) {
 	ctx := context.TODO()
 
 	rs := fakeRocksetClient()
-	rs.GetVirtualInstanceReturnsOnCall(0, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceInitializing)}, nil)
-	rs.GetVirtualInstanceReturnsOnCall(1, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceActive)}, nil)
+	rs.GetVirtualInstanceReturnsOnCall(0, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceInitializing.String())}, nil)
+	rs.GetVirtualInstanceReturnsOnCall(1, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceActive.String())}, nil)
 
 	err := wait.New(&rs).UntilVirtualInstanceActive(ctx, "id")
 	assert.NoError(t, err)
@@ -29,9 +29,9 @@ func TestWait_untilVirtualInstanceSuspended(t *testing.T) {
 	ctx := context.TODO()
 
 	rs := fakeRocksetClient()
-	rs.GetVirtualInstanceReturnsOnCall(0, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceActive)}, nil)
-	rs.GetVirtualInstanceReturnsOnCall(1, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceSuspending)}, nil)
-	rs.GetVirtualInstanceReturnsOnCall(2, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceSuspended)}, nil)
+	rs.GetVirtualInstanceReturnsOnCall(0, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceActive.String())}, nil)
+	rs.GetVirtualInstanceReturnsOnCall(1, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceSuspending.String())}, nil)
+	rs.GetVirtualInstanceReturnsOnCall(2, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceSuspended.String())}, nil)
 
 	err := wait.New(&rs).UntilVirtualInstanceSuspended(ctx, "id")
 	assert.NoError(t, err)
@@ -42,7 +42,7 @@ func TestWait_untilVirtualInstanceGone(t *testing.T) {
 	ctx := context.TODO()
 
 	rs := fakeRocksetClient()
-	rs.GetVirtualInstanceReturnsOnCall(0, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceActive)}, nil)
+	rs.GetVirtualInstanceReturnsOnCall(0, openapi.VirtualInstance{State: openapi.PtrString(option.VirtualInstanceActive.String())}, nil)
 	rs.GetVirtualInstanceReturnsOnCall(1, openapi.VirtualInstance{}, NotFoundErr)
 
 	err := wait.New(&rs).UntilVirtualInstanceGone(ctx, "id")