Skip to content

Commit 96688b0

Browse files
bradlobradlo
committed
implement support for database linking
1 parent 28f0955 commit 96688b0

File tree

8 files changed

+240
-63
lines changed

8 files changed

+240
-63
lines changed

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ go 1.18
44

55
require (
66
github.com/pkg/errors v0.9.1
7-
github.com/relationalai/rai-sdk-go v0.4.3-alpha
7+
github.com/relationalai/rai-sdk-go v0.4.4-alpha
88
github.com/spf13/cobra v1.5.0
99
)
1010

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -277,8 +277,8 @@ github.com/prometheus/procfs v0.0.0-20190117184657-bf6a532e95b1/go.mod h1:c3At6R
277277
github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
278278
github.com/prometheus/procfs v0.0.8/go.mod h1:7Qr8sr6344vo1JqZ6HhLceV9o3AJ1Ff+GxbHq6oeK9A=
279279
github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
280-
github.com/relationalai/rai-sdk-go v0.4.3-alpha h1:uz3Dkn8Qg45yM9NACYdizUGJhrU6DRDoLONNwTMy94o=
281-
github.com/relationalai/rai-sdk-go v0.4.3-alpha/go.mod h1:qfIZ7OiUpM+ZLPB8m358DFQPRPFzMjM4ghX4iKR8UHo=
280+
github.com/relationalai/rai-sdk-go v0.4.4-alpha h1:AsZsEWK8K9RPIQ0hFDFuPHqtUrhY3Yz5TeRdhkzPvNw=
281+
github.com/relationalai/rai-sdk-go v0.4.4-alpha/go.mod h1:qfIZ7OiUpM+ZLPB8m358DFQPRPFzMjM4ghX4iKR8UHo=
282282
github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
283283
github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
284284
github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=

‎rai/cmds.go‎

Lines changed: 84 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,6 @@ import (
2020
"fmt"
2121
"io"
2222
"io/ioutil"
23-
"math/rand"
24-
"net/http"
2523
"os"
2624
"path"
2725
"path/filepath"
@@ -82,16 +80,19 @@ func (a *Action) Context() context.Context {
8280
return a.cmd.Context()
8381
}
8482

83+
// Returns the bool value corresponding to the named flag.
8584
func (a *Action) getBool(name string) bool {
8685
result, _ := a.cmd.Flags().GetBool(name)
8786
return result
8887
}
8988

89+
// Returns the int value corresponding to the named flag.
9090
func (a *Action) getInt(name string) int {
9191
result, _ := a.cmd.Flags().GetInt(name)
9292
return result
9393
}
9494

95+
// Returns the rune value corresponding to the named flag.
9596
func (a *Action) getRune(name string) rune {
9697
s, _ := a.cmd.Flags().GetString(name)
9798
if s == "" {
@@ -100,11 +101,23 @@ func (a *Action) getRune(name string) rune {
100101
return []rune(s)[0]
101102
}
102103

104+
// Returns the string value corresponding to the named flag.
103105
func (a *Action) getString(name string) string {
104106
result, _ := a.cmd.Flags().GetString(name)
105107
return result
106108
}
107109

110+
// Returns the string value corresponding to the named flag, and if no value
111+
// is set, return the value corresponding to the given environment variable.
112+
func (a *Action) getStringEnv(name, key string) string {
113+
result, _ := a.cmd.Flags().GetString(name)
114+
if result == "" {
115+
result = os.Getenv(key)
116+
}
117+
return result
118+
}
119+
120+
// Returns the string array value corresponding to the named flag.
108121
func (a *Action) getStringArray(name string) []string {
109122
result, _ := a.cmd.Flags().GetStringArray(name)
110123
return result
@@ -170,7 +183,6 @@ func (a *Action) showValue(v interface{}) {
170183
}
171184
case "json":
172185
break // default
173-
174186
}
175187
showJSON(v)
176188
}
@@ -192,7 +204,7 @@ func (a *Action) Start(format string, args ...interface{}) *Action {
192204
var msg string
193205
msg = fmt.Sprintf(format, args...)
194206
msg = fmt.Sprintf("%s .. ", msg)
195-
fmt.Fprintf(os.Stderr, msg)
207+
fmt.Fprintln(os.Stderr, msg)
196208
return a
197209
}
198210

@@ -210,6 +222,7 @@ func (a *Action) Exit(result interface{}, err error) {
210222
}
211223

212224
// Pick a random PROVISIONED engine.
225+
/*
213226
func pickRandomEngine(action *Action) string {
214227
rsp, err := action.Client().ListEngines("state", "PROVISIONED")
215228
if err != nil {
@@ -224,6 +237,7 @@ func pickRandomEngine(action *Action) string {
224237
ix := rand.Intn(len(rsp))
225238
return rsp[ix].Name
226239
}
240+
*/
227241

228242
// Pick the most recently created PROVISIONED engine.
229243
func pickLatestEngine(action *Action) string {
@@ -250,25 +264,6 @@ func pickEngine(action *Action) string {
250264
return pickLatestEngine(action)
251265
}
252266

253-
// Answers if the given state represents a terminal state.
254-
func isTerminalState(state, targetState string) bool {
255-
if state == targetState {
256-
return true
257-
}
258-
if strings.Contains(state, "FAILED") {
259-
return true
260-
}
261-
return false
262-
}
263-
264-
func isStatusNotFound(err error) bool {
265-
e, ok := err.(*rai.HTTPError)
266-
if !ok {
267-
return false
268-
}
269-
return e.StatusCode == http.StatusNotFound
270-
}
271-
272267
//
273268
// Databases
274269
//
@@ -614,7 +609,7 @@ func listEdbNames(cmd *cobra.Command, args []string) {
614609
for name := range nameMap {
615610
names = append(names, name)
616611
}
617-
names = sort.StringSlice(names)
612+
sort.Strings(names)
618613
action.Exit(names, nil)
619614
}
620615

@@ -784,17 +779,17 @@ func updateUser(cmd *cobra.Command, args []string) {
784779
}
785780

786781
//
787-
// Integrations
782+
// Snowflake integrations
788783
//
789784

790785
func createSnowflakeIntegration(cmd *cobra.Command, args []string) {
791786
action := newAction(cmd)
792787
name := args[0]
793-
account := action.getString("account")
794-
adminUsername := action.getString("admin-username")
795-
adminPassword := action.getString("admin-password")
796-
proxyUsername := action.getString("proxy-username")
797-
proxyPassword := action.getString("proxy-password")
788+
account := action.getStringEnv("account", "SNOWSQL_ACCOUNT")
789+
adminUsername := action.getStringEnv("admin-username", "SNOWSQL_USER")
790+
adminPassword := action.getStringEnv("admin-password", "SNOWSQL_PWD")
791+
proxyUsername := action.getStringEnv("proxy-username", "SNOWSQL_USER")
792+
proxyPassword := action.getStringEnv("proxy-password", "SNOWSQL_PWD")
798793
adminCreds := rai.SnowflakeCredentials{
799794
Username: adminUsername, Password: adminPassword}
800795
proxyCreds := rai.SnowflakeCredentials{
@@ -808,12 +803,11 @@ func createSnowflakeIntegration(cmd *cobra.Command, args []string) {
808803
func deleteSnowflakeIntegration(cmd *cobra.Command, args []string) {
809804
action := newAction(cmd)
810805
name := args[0]
811-
adminUsername := action.getString("admin-username")
812-
adminPassword := action.getString("admin-password")
813-
adminCreds := rai.SnowflakeCredentials{
814-
Username: adminUsername, Password: adminPassword}
806+
username := action.getStringEnv("username", "SNOWSQL_USER")
807+
password := action.getStringEnv("password", "SNOWSQL_PWD")
808+
creds := rai.SnowflakeCredentials{Username: username, Password: password}
815809
action.Start("Delete Snowflake integration '%s'", name)
816-
err := action.Client().DeleteSnowflakeIntegration(name, &adminCreds)
810+
err := action.Client().DeleteSnowflakeIntegration(name, &creds)
817811
action.Exit(nil, err)
818812
}
819813

@@ -832,6 +826,61 @@ func listSnowflakeIntegrations(cmd *cobra.Command, _ []string) {
832826
action.Exit(rsp, err)
833827
}
834828

829+
//
830+
// Snowflake database links
831+
//
832+
833+
func createSnowflakeDatabaseLink(cmd *cobra.Command, args []string) {
834+
action := newAction(cmd)
835+
integration := args[0]
836+
database := action.getStringEnv("database", "SNOWSQL_DATABASE")
837+
schema := action.getStringEnv("schema", "SNOWSQL_SCHEMA")
838+
role := action.getStringEnv("role", "SNOWSQL_ROLE")
839+
username := action.getStringEnv("username", "SNOWSQL_USER")
840+
password := action.getStringEnv("username", "SNOWSQL_PWD")
841+
creds := rai.SnowflakeCredentials{Username: username, Password: password}
842+
name := fmt.Sprintf("%s.%s", database, schema)
843+
action.Start("Create Snowflake database link '%s' (%s)", name, integration)
844+
rsp, err := action.Client().CreateSnowflakeDatabaseLink(
845+
integration, database, schema, role, &creds)
846+
action.Exit(rsp, err)
847+
}
848+
849+
func deleteSnowflakeDatabaseLink(cmd *cobra.Command, args []string) {
850+
action := newAction(cmd)
851+
integration := args[0]
852+
database := action.getStringEnv("database", "SNOWSQL_DATABASE")
853+
schema := action.getStringEnv("schema", "SNOWSQL_SCHEMA")
854+
role := action.getStringEnv("role", "SNOWSQL_ROLE")
855+
username := action.getStringEnv("username", "SNOWSQL_USER")
856+
password := action.getStringEnv("username", "SNOWSQL_PWD")
857+
creds := rai.SnowflakeCredentials{Username: username, Password: password}
858+
name := fmt.Sprintf("%s.%s", database, schema)
859+
action.Start("Delete Snowflake database link '%s' (%s)", name, integration)
860+
err := action.Client().DeleteSnowflakeDatabaseLink(
861+
integration, database, schema, role, &creds)
862+
action.Exit(nil, err)
863+
}
864+
865+
func getSnowflakeDatabaseLink(cmd *cobra.Command, args []string) {
866+
action := newAction(cmd)
867+
integration := args[0]
868+
database := action.getStringEnv("database", "SNOWSQL_DATABASE")
869+
schema := action.getStringEnv("schema", "SNOWSQL_SCHEMA")
870+
name := fmt.Sprintf("%s.%s", database, schema)
871+
action.Start("Get Snowflake database link '%s' (%s)", name, integration)
872+
rsp, err := action.Client().GetSnowflakeDatabaseLink(integration, database, schema)
873+
action.Exit(rsp, err)
874+
}
875+
876+
func listSnowflakeDatabaseLinks(cmd *cobra.Command, args []string) {
877+
action := newAction(cmd)
878+
integration := args[0]
879+
action.Start("List Snowflake database links (%s)", integration)
880+
rsp, err := action.Client().ListSnowflakeDatabaseLinks(integration)
881+
action.Exit(rsp, err)
882+
}
883+
835884
//
836885
// Misc
837886
//

‎rai/main.go‎

Lines changed: 49 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -290,26 +290,26 @@ func addCommands(root *cobra.Command) {
290290
cmd.Flags().StringArray("role", nil, "user roles")
291291
root.AddCommand(cmd)
292292

293-
// Integrations
293+
// Snowflake integrations
294294
cmd = &cobra.Command{
295295
Use: "create-snowflake-integration name",
296296
Short: "Create a Snowflake integration",
297297
Args: cobra.ExactArgs(1),
298298
Run: createSnowflakeIntegration}
299-
cmd.Flags().String("account", "", "Snowflake account")
300-
cmd.Flags().String("admin-username", "", "admin username")
301-
cmd.Flags().String("admin-password", "", "admin password")
302-
cmd.Flags().String("proxy-username", "", "proxy username")
303-
cmd.Flags().String("proxy-password", "", "proxy password")
299+
cmd.Flags().String("account", "", "Snowflake account (default: SNOWSQL_ACCOUNT env var)")
300+
cmd.Flags().String("admin-username", "", "Snowflake admin username (default: SNOWSQL_USER env var")
301+
cmd.Flags().String("admin-password", "", "Snowflake admin password (default: SNOWSQL_PWD env var")
302+
cmd.Flags().String("proxy-username", "", "Snowflake proxy username (default: SNOWSQL_USER env var")
303+
cmd.Flags().String("proxy-password", "", "Snowflake proxy password (default: SNOWSQL_PWD env var)")
304304
root.AddCommand(cmd)
305305

306306
cmd = &cobra.Command{
307307
Use: "delete-snowflake-integration name",
308308
Short: "Delete a Snowflake integration",
309309
Args: cobra.ExactArgs(1),
310310
Run: deleteSnowflakeIntegration}
311-
cmd.Flags().String("admin-username", "", "admin username")
312-
cmd.Flags().String("admin-password", "", "admin password")
311+
cmd.Flags().String("username", "", "Snowflake username (default: SNOWSQL_USER env var)")
312+
cmd.Flags().String("password", "", "Snowflake password (default: SNOWSQL_PWD env var")
313313
root.AddCommand(cmd)
314314

315315
cmd = &cobra.Command{
@@ -326,20 +326,57 @@ func addCommands(root *cobra.Command) {
326326
Run: listSnowflakeIntegrations}
327327
root.AddCommand(cmd)
328328

329+
// Snowflake database links
330+
cmd = &cobra.Command{
331+
Use: "create-snowflake-database-link integration",
332+
Short: "Create a Snowflake database link",
333+
Args: cobra.ExactArgs(1),
334+
Run: createSnowflakeDatabaseLink}
335+
cmd.Flags().String("database", "", "Snowflake database (default: SNOWSQL_DATABASE)")
336+
cmd.Flags().String("schema", "", "Snowflake schema (default: SNOWSQL_SCHEMA)")
337+
cmd.Flags().String("role", "", "Snowflake role (default: SNOWSQL_ROLE)")
338+
cmd.Flags().String("username", "", "Snowflake username (default: SNOWSQL_USER)")
339+
cmd.Flags().String("password", "", "Snowflake password (default: SNOWSQL_PWD)")
340+
root.AddCommand(cmd)
341+
342+
cmd = &cobra.Command{
343+
Use: "delete-snowflake-database-link integration",
344+
Short: "Delete a Snowflake database link",
345+
Args: cobra.ExactArgs(1),
346+
Run: deleteSnowflakeDatabaseLink}
347+
cmd.Flags().String("database", "", "Snowflake database (default: SNOWSQL_DATABASE env var)")
348+
cmd.Flags().String("schema", "", "Snowflake schema (default: SNOWSQL_SCHEMA env var)")
349+
cmd.Flags().String("role", "", "Snowflake role (default: SNOWSQL_ROLE env var)")
350+
cmd.Flags().String("username", "", "Snowflake username (default: SNOWSQL_USER env var)")
351+
cmd.Flags().String("password", "", "Snowflake password (default: SNOWSQL_PWD env var)")
352+
root.AddCommand(cmd)
353+
354+
cmd = &cobra.Command{
355+
Use: "get-snowflake-database-link integration",
356+
Short: "Get information about the given Snowflake database link",
357+
Args: cobra.ExactArgs(1),
358+
Run: getSnowflakeDatabaseLink}
359+
cmd.Flags().String("database", "", "Snowflake database (default: SNOWSQL_DATABASE env var)")
360+
cmd.Flags().String("schema", "", "Snowflake schema (default: SNOWSQL_SCHEMA env var)")
361+
root.AddCommand(cmd)
362+
363+
cmd = &cobra.Command{
364+
Use: "list-snowflake-database-links integration",
365+
Short: "List Snowflake database links for the given integration",
366+
Args: cobra.ExactArgs(1),
367+
Run: listSnowflakeDatabaseLinks}
368+
root.AddCommand(cmd)
369+
329370
// Misc
330371
cmd = &cobra.Command{
331372
Use: "get-access-token",
332373
Short: "Get OAuth access token",
333374
Run: getAccessToken}
334375
root.AddCommand(cmd)
335-
336376
}
337377

338378
func main() {
339379
var root = &cobra.Command{Use: "rai"}
340-
// todo: additional root options
341-
// --request-timeout
342-
// --token : Bearer token for authenticating API request
343380
root.PersistentFlags().String("host", "", "host name")
344381
root.PersistentFlags().String("port", "", "port number")
345382
root.PersistentFlags().String("config", "~/.rai/config", "config file")

0 commit comments

Comments
 (0)