Skip to content

Commit

Permalink
support getting a bearer token from auth0
Browse files Browse the repository at this point in the history
  • Loading branch information
@pmenglund
pmenglund committed Oct 8, 2023
1 parent 5401b21 commit adffaa1
Show file tree
Hide file tree
Showing 8 changed files with 352 additions and 157 deletions.
54 changes: 42 additions & 12 deletions cmd/auth.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,33 +2,40 @@ package cmd

import (
"fmt"
"github.com/rockset/cli/config"
"log/slog"
"time"

"github.com/pkg/browser"
devauth "github.com/rockset/device-authorization"
"github.com/spf13/cobra"
"log"
)

const Auth0ClientID = "0dJNiGWClbLjg7AdtXtAyPCeE0jKOFet"

func newAuthCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "auth",
Args: cobra.NoArgs,
Use: "auth NAME CLUSTER ORGANIZATION",
Args: cobra.ExactArgs(3),
Short: "authenticate",
RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
name := args[0]
cluster := args[1]
org := args[2]

p := devauth.NewProvider("auth0")
cfg := p.Config("rockset", Auth0ClientID)
cfg.Audience = "https://rockset.sh/"
cfg.OAuth2Config.Endpoint.AuthURL = "https://auth.rockset.com/oauth/device/code"
cfg.OAuth2Config.Endpoint.TokenURL = "https://auth.rockset.com/oauth/token"
cfg.OAuth2Config.Scopes = append(cfg.OAuth2Config.Scopes, "email")
authCfg := p.Config("rockset", Auth0ClientID)
authCfg.Audience = "https://rockset.sh/"
authCfg.OAuth2Config.Endpoint.AuthURL = "https://auth.rockset.com/oauth/device/code"
authCfg.OAuth2Config.Endpoint.TokenURL = "https://auth.rockset.com/oauth/token"
authCfg.OAuth2Config.Scopes = append(authCfg.OAuth2Config.Scopes, "email")

a := devauth.NewAuthorizer(cfg)
a := devauth.NewAuthorizer(authCfg)

code, err := a.RequestCode(ctx)
if err != nil {
log.Fatalf("failed to request a device code: %v", err)
slog.Error("failed to request a device code", "err", err)
}

fmt.Printf(`Attempting to automatically open the SSO authorization page in your default browser.
Expand All @@ -40,14 +47,37 @@ Then enter the code:
%s
`, code.VerificationURIComplete, code.UserCode)

if err := browser.OpenURL(code.VerificationURIComplete); err != nil {
slog.Warn("could not open", "url", code.VerificationURIComplete, "err", err)
}

token, err := a.WaitForAuthorization(ctx, code)
if err != nil {
log.Fatalf("failed to wait for authorization: %v", err)
slog.Error("failed to wait for authorization", "err", err)
}

fmt.Printf("Successfully logged in!\n")

fmt.Printf("token:\n%s\n", token.AccessToken)
fmt.Printf("token:\n%s\n", token.IDToken)

// TODO save token & expiration in config
cfg, err := config.Load()
if err != nil {
return err
}

if err = cfg.AddToken(name, config.Token{
Token: token.IDToken,
Org: org,
Server: cluster,
Expiration: time.Now().Add(time.Duration(token.ExpiresIn) * time.Second),
}); err != nil {
return err
}

if err = config.Store(cfg); err != nil {
return err
}

return nil
},
Expand Down
185 changes: 47 additions & 138 deletions cmd/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,13 @@ package cmd
import (
"errors"
"fmt"
"log/slog"
"os"
"path"
"sort"

"github.com/mitchellh/go-homedir"
"github.com/spf13/cobra"
"golang.org/x/exp/slog"
"gopkg.in/yaml.v3"

"github.com/rockset/cli/config"
"github.com/rockset/rockset-go-client"
"github.com/spf13/cobra"
)

func newListConfigCmd() *cobra.Command {
Expand All @@ -34,17 +31,17 @@ configs:
apiserver: api.usw2a1.rockset.com
prod:
apikey: ...
apiserver: api.use1a1.rockset.com`, APIKeysFile),
apiserver: api.use1a1.rockset.com`, config.FileName),
RunE: func(cmd *cobra.Command, args []string) error {
cfg, err := loadAPIKeys()
cfg, err := config.Load()
if err != nil {
if !errors.Is(err, os.ErrNotExist) {
return fmt.Errorf("config file %s not readable available: %v", APIKeysFile, err)
return fmt.Errorf("config file %s not readable available: %v", config.FileName, err)
}
return err
}

_, _ = fmt.Fprintf(cmd.OutOrStdout(), "available configs:\n")
_, _ = fmt.Fprintf(cmd.OutOrStdout(), "available contexts:\n")
var names []string
for name := range cfg.Keys {
names = append(names, name)
Expand Down Expand Up @@ -83,7 +80,7 @@ func newCreateConfigCmd() *cobra.Command {
return fmt.Errorf("both --apikey and --server are required")
}

cfg, err := loadAPIKeys()
cfg, err := config.Load()
if err != nil {
return err
}
Expand Down Expand Up @@ -111,12 +108,12 @@ func newCreateConfigCmd() *cobra.Command {
logger.Warn("skipping auth check due to --force option")
}

cfg.Keys[args[0]] = APIKey{
cfg.Keys[args[0]] = config.APIKey{
Key: key,
Server: server,
}

if err = storeAPIKeys(cfg); err != nil {
if err = config.Store(cfg); err != nil {
return err
}

Expand All @@ -132,6 +129,7 @@ func newCreateConfigCmd() *cobra.Command {

return &cmd
}

func newUseConfigCmd() *cobra.Command {
cmd := cobra.Command{
Use: "config NAME",
Expand All @@ -141,17 +139,16 @@ func newUseConfigCmd() *cobra.Command {
Annotations: group("config"),
Args: cobra.ExactArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
cfg, err := loadAPIKeys()
cfg, err := config.Load()
if err != nil {
return err
}

if _, found := cfg.Keys[args[0]]; !found {
return fmt.Errorf("configuration %s not found", args[0])
if err = cfg.Use(args[0]); err != nil {
return err
}

cfg.Current = args[0]
if err = storeAPIKeys(cfg); err != nil {
if err = config.Store(cfg); err != nil {
return err
}

Expand All @@ -164,134 +161,46 @@ func newUseConfigCmd() *cobra.Command {
return &cmd
}

const (
APIKeysFileName = "apikeys.yaml"
HistoryFile = "cli.hist"
)

var APIKeysFile string

func init() {
file, err := apikeysPath()
if err != nil {
panic(fmt.Sprintf("unable to locate apikey file %s: %v", APIKeysFileName, err))
}
APIKeysFile = file
}

type APIKeys struct {
Current string `yaml:"current"`
Keys map[string]APIKey `yaml:"configs"`
}

type APIKey struct {
Key string `yaml:"apikey"`
Server string `yaml:"apiserver"`
}

func apikeysPath() (string, error) {
return rocksetFile(APIKeysFileName)
}

func historyFile() (string, error) {
return rocksetFile(HistoryFile)
}

func rocksetFile(name string) (string, error) {
home, err := homedir.Dir()
if err != nil {
return "", err
}
return path.Join(home, ".config", "rockset", name), nil
}

func storeAPIKeys(cfg APIKeys) error {
file, err := apikeysPath()
if err != nil {
return err
}

// TODO create directory

f, err := os.OpenFile(file, os.O_TRUNC|os.O_WRONLY, 0600)
if err != nil {
return err
}

enc := yaml.NewEncoder(f)
if err = enc.Encode(cfg); err != nil {
return err
}

if err = enc.Close(); err != nil {
logger.Error("failed to close config", "err", err)
}

return nil
}

func loadAPIKeys() (APIKeys, error) {
var keys APIKeys

cfg, err := apikeysPath()
if err != nil {
return keys, err
}

f, err := os.Open(cfg)
if err != nil {
return keys, fmt.Errorf("failed to read apikey config file: %w", err)
}

dec := yaml.NewDecoder(f)
err = dec.Decode(&keys)
if err != nil {
return keys, err
}
return keys, nil
}

func rockClient(cmd *cobra.Command) (*rockset.RockClient, error) {
var apikey, apiserver string

// load from config, ok if none is found
cfg, err := loadAPIKeys()
cfg, err := config.Load()
if err != nil {
return nil, err
if !errors.Is(err, config.NotFoundErr) {
return nil, err
}
}

context, _ := cmd.Flags().GetString(ContextFLag)
if context == "" {
slog.Debug("using context from file")
context = cfg.Current
}
if ctx, found := cfg.Keys[context]; found {
slog.Debug("using", context, context)
apikey = ctx.Key
apiserver = ctx.Server
} else {
slog.Debug("not found", "context", context)
override, _ := cmd.Flags().GetString(ContextFLag)
if override != "" {
slog.Debug("using override", "name", override)
}

// load from environment
if key, found := os.LookupEnv(rockset.APIKeyEnvironmentVariableName); found {
slog.Debug("set apikey")
apikey = key
}
if server, found := os.LookupEnv(rockset.APIServerEnvironmentVariableName); found {
slog.Debug("set apiserver")
apiserver = server
var options = []rockset.RockOption{
rockset.WithCustomHeader("rockset-go-cli", Version),
}

// let the --cluster flag override the apiserver
if cluster, _ := cmd.Flags().GetString(ClusterFLag); cluster != "" {
slog.Debug("override apiserver")
apiserver = fmt.Sprintf("https://api.%s.rockset.com", cluster)
opts, err := cfg.AsOptions(override)
if err != nil {
return nil, err
}

return rockset.NewClient(
rockset.WithAPIKey(apikey),
rockset.WithAPIServer(apiserver),
// TODO add rockset.WithCustomHeader("rockset-go-cli", Version) once the new client it released
)
options = append(options, opts...)

/*
// load from environment
if key, found := os.LookupEnv(rockset.APIKeyEnvironmentVariableName); found {
slog.Debug("set apikey")
apikey = key
}
if server, found := os.LookupEnv(rockset.APIServerEnvironmentVariableName); found {
slog.Debug("set apiserver")
apiserver = server
}
// let the --cluster flag override the apiserver
if cluster, _ := cmd.Flags().GetString(ClusterFLag); cluster != "" {
slog.Debug("override apiserver")
apiserver = fmt.Sprintf("https://api.%s.rockset.com", cluster)
}
*/
return rockset.NewClient(options...)
}
5 changes: 3 additions & 2 deletions cmd/query.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package cmd
import (
"context"
"fmt"
"github.com/rockset/cli/config"
"github.com/rockset/cli/format"
"io"
"os"
Expand All @@ -24,7 +25,7 @@ func newListQueryCmd() *cobra.Command {
Short: "list queries",
Long: "list queries on a virtual instance",
Args: cobra.ExactArgs(1),
Annotations: group("query"), // TODO shoudl this be in the VI group too?
Annotations: group("query"), // TODO should this be in the VI group too?
RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

Expand Down Expand Up @@ -187,7 +188,7 @@ var (
)

func interactiveQuery(ctx context.Context, in io.ReadCloser, out io.Writer, rs *rockset.RockClient) error {
histFile, err := historyFile()
histFile, err := config.HistoryFile()
if err != nil {
return err
}
Expand Down
Loading

0 comments on commit adffaa1

Please sign in to comment.