Skip to content

Commit

Permalink
feat: Server switcher command (#5)
Browse files Browse the repository at this point in the history
feat: Server switcher command

feat: Tenant, Database and defaults flag

- Updated docs
  • Loading branch information
tazarov authored Mar 5, 2024
2 parents 1fb51ed + 39b3a3b commit 3a05fc9
Show file tree
Hide file tree
Showing 4 changed files with 159 additions and 3 deletions.
Binary file removed chroma
Binary file not shown.
8 changes: 5 additions & 3 deletions cmd/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,16 @@ import (

// serverCmd represents the server command
var serverCmd = &cobra.Command{
Use: "server",
Short: "Manage Chroma servers",
Long: ``,
Use: "server",
Aliases: []string{"s"},
Short: "Manage Chroma servers",
Long: ``,
}

func init() {
rootCmd.AddCommand(serverCmd)
serverCmd.AddCommand(server.AddCommand)
serverCmd.AddCommand(server.ListCommand)
serverCmd.AddCommand(server.RmCommand)
rootCmd.AddCommand(server.SwitchCommand)
}
126 changes: 126 additions & 0 deletions cmd/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,55 @@ func getDatabase(changed bool) (string, error) {
return Database, nil
}

// setActiveServer sets the active server to the one with the given alias
func setActiveServer(alias string) error {
var servers = viper.GetStringMap("servers")
if servers == nil {
servers = make(map[string]interface{})
}
if _, ok := servers[alias]; ok {
viper.Set("active_server", alias)
err := viper.WriteConfig()
if err != nil {
return fmt.Errorf("unable to write to config file: %v", err)
}
} else {
return fmt.Errorf("server with alias %v does not exist", alias)
}
return nil
}

// setActiveDatabase sets the active database to the one with the given name
func setActiveDatabase(database string) error {
viper.Set("active_db", database)
err := viper.WriteConfig()
if err != nil {
return fmt.Errorf("unable to write to config file: %v", err)
}
return nil
}

// setActiveTenant sets the active tenant to the one with the given name
func setActiveTenant(tenant string) error {
viper.Set("active_tenant", tenant)
err := viper.WriteConfig()
if err != nil {
return fmt.Errorf("unable to write to config file: %v", err)
}
return nil
}

func getServer(alias string) (map[string]interface{}, error) {
var servers = viper.GetStringMap("servers")
if servers == nil {
servers = make(map[string]interface{})
}
if server, ok := servers[alias]; ok {
return server.(map[string]interface{}), nil
}
return nil, fmt.Errorf("server with alias %v does not exist", alias)
}

var Host string
var Port string
var Overwrite bool
Expand Down Expand Up @@ -219,6 +268,10 @@ var RmCommand = &cobra.Command{
os.Exit(0)
}
delete(servers, alias)
if viper.GetString("active") == alias {
viper.Set("active", "")
fmt.Println(alias, "was the active server. You will need to set a new active server.")
}
viper.Set("servers", servers)
err := viper.WriteConfig()
if err != nil {
Expand Down Expand Up @@ -249,6 +302,74 @@ var ListCommand = &cobra.Command{
},
}

var DBAndTenantDefaults bool

var SwitchCommand = &cobra.Command{
Use: "switch",
Aliases: []string{"sw"},
Short: "Set active server",
Args: cobra.MinimumNArgs(1),
Run: func(cmd *cobra.Command, args []string) {
alias := args[0]
err := setActiveServer(alias)
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
if cmd.Flags().Changed("tenant") {
err := setActiveTenant(Tenant)
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
fmt.Printf("Tenant '%v' set as active!\n", Tenant)
} else if cmd.Flags().Changed("defaults") {
getSrv, err := getServer(alias)
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
if getSrv["tenant"] == nil {
getSrv["tenant"] = DefaultTenant
}
if _, ok := getSrv["tenant"]; ok {
err := setActiveTenant(getSrv["tenant"].(string))
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
fmt.Printf("Tenant '%v' set as active!\n", getSrv["tenant"])
}
}
if cmd.Flags().Changed("database") {
err := setActiveDatabase(Database)
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
fmt.Printf("Database '%v' set as active!\n", Database)
} else if cmd.Flags().Changed("defaults") {
getSrv, err := getServer(alias)
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
if getSrv["database"] == nil {
getSrv["database"] = DefaultDatabase
}
if _, ok := getSrv["database"]; ok {
err := setActiveDatabase(getSrv["database"].(string))
if err != nil {
fmt.Printf("%v\n", err)
os.Exit(1)
}
}
fmt.Printf("Database '%v' set as active!\n", getSrv["database"])
}
fmt.Printf("Server '%v' set as active!\n", alias)
},
}

func init() {
AddCommand.Flags().StringVarP(&Host, "host", "H", "", "Chroma server host")
AddCommand.Flags().StringVarP(&Port, "port", "p", "", "Chroma server port")
Expand All @@ -260,4 +381,9 @@ func init() {
AddCommand.ValidArgs = []string{"alias"}
RmCommand.ValidArgs = []string{"alias"}
RmCommand.Flags().BoolVarP(&ForceDelete, "force", "f", false, "Force remove server without confirmation")
SwitchCommand.Flags().StringVarP(&Tenant, "tenant", "t", "", "Default tenant for the server")
SwitchCommand.Flags().StringVarP(&Database, "database", "d", "", "Default database for the server")
SwitchCommand.Flags().BoolVarP(&DBAndTenantDefaults, "defaults", "r", false, "Reset active tenant and database to defaults")
SwitchCommand.MarkFlagsMutuallyExclusive("tenant", "defaults")
SwitchCommand.MarkFlagsMutuallyExclusive("database", "defaults")
}
28 changes: 28 additions & 0 deletions docs/docs/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,3 +25,31 @@ Flags:
```bash
chroma server add <server-alias> -h <host> -p <port> -o -s -t <tenant> -d <database>
```

### Switch Server

Arguments:

- `server-alias` - Alias for the server

Flags:

- `-t` or `--tenant` - Changes the default tenant
- `-d` or `--database` - Changes the default database
- `-r` or `--defaults` - Uses the default tenant and database (this is mutually exclusive with `-t` and `-d` flags)

```bash
chroma switch <server-alias>
```

With shorthands:

```bash
chroma sw <server-alias> -t <tenant> -d <database>
```

With defaults:

```bash
chroma sw <server-alias> -r
```

0 comments on commit 3a05fc9

Please sign in to comment.