add support for suspend/resume engines

examples/suspend_engine/main.go

@@ -0,0 +1,64 @@
+// Copyright 2022 RelationalAI, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"log"
+	"os"
+
+	"github.com/jessevdk/go-flags"
+	"github.com/relationalai/rai-sdk-go/rai"
+)
+
+type Options struct {
+	Engine  string `short:"e" long:"engine" required:"true" description:"engine name"`
+	Profile string `long:"profile" default:"default" description:"config profile"`
+}
+
+func run(opts *Options) error {
+	client, err := rai.NewClientFromConfig(opts.Profile)
+	if err != nil {
+		return err
+	}
+	rsp, err := client.GetEngine(opts.Engine)
+	if err != nil {
+		return err
+	}
+
+	switch rsp.State {
+	case "PROVISIONED":
+		err := client.SuspendEngine(opts.Engine)
+		if err != nil {
+			return err
+		}
+	case "SUSPENDED":
+		err := client.ResumeEngine(opts.Engine)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func main() {
+	var opts Options
+	if _, err := flags.ParseArgs(&opts, os.Args); err != nil {
+		os.Exit(1)
+	}
+	if err := run(&opts); err != nil {
+		log.Fatal(err)
+	}
+}

rai/client.go

@@ -601,6 +601,28 @@ func (c *Client) ListEngines(filters ...interface{}) ([]Engine, error) {
 	return result.Engines, nil
 }
 
+func (c *Client) ResumeEngine(engineName string) error {
+	var result interface{}
+	data := &SuspendEngineRequest{Suspend: false}
+	uri := makePath(PathEngine, engineName)
+	err := c.Patch(uri, nil, data, &result)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (c *Client) SuspendEngine(engineName string) error {
+	var result interface{}
+	data := &SuspendEngineRequest{Suspend: true}
+	uri := makePath(PathEngine, engineName)
+	err := c.Patch(uri, nil, data, &result)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 //
 // OAuth Clients
 //

rai/client_test.go

@@ -217,6 +217,38 @@ func TestEngine(t *testing.T) {
 	assert.Nil(t, err)
 	engine = findEngine(engines, test.engineName)
 	assert.Nil(t, engine)
+
+	// suspend the test engine
+	err = client.SuspendEngine(test.engineName)
+	assert.Nil(t, err)
+
+	waitOrFail := func(state string) {
+		const maxWaitTime = 600 // seconds
+		const duration = 5      // seconds
+		eng := &Engine{}
+		waitTime := 0
+		for !isTerminalState(eng.State, state) {
+			time.Sleep(duration * time.Second)
+			waitTime += duration
+			eng, err = client.GetEngine(test.engineName)
+			assert.Nil(t, err)
+			assert.Less(t, waitTime, maxWaitTime, "Failed waiting for engine state change: %s", eng.State)
+		}
+	}
+	waitOrFail("SUSPENDED")
+
+	// check status
+	engine, err = client.GetEngine(test.engineName)
+	assert.Nil(t, err)
+	assert.NotNil(t, engine)
+	assert.Equal(t, "SUSPENDED", engine.State)
+
+	// suspend the test engine
+	err = client.ResumeEngine(test.engineName)
+	assert.Nil(t, err)
+
+	waitOrFail("PROVISIONED")
+
 }
 
 // Test transaction execution.

rai/models.go

@@ -228,6 +228,10 @@ type createEngineResponse struct {
 	Engine Engine `json:"compute"`
 }
 
+type SuspendEngineRequest struct {
+	Suspend bool `json:"suspend"`
+}
+
 type createOAuthClientRequest struct {
 	Name        string   `json:"name"`
 	Permissions []string `json:"permissions"`