From 03c2c0df39cd3cacc0386440a0e6c0ad36e8f70b Mon Sep 17 00:00:00 2001 From: Pavel Sviderski Date: Wed, 11 Sep 2024 18:40:06 +1000 Subject: [PATCH] refactor cluster connection in CLI --- cmd/uncloud/machine/add.go | 42 +++-- internal/cli/cli.go | 159 +++++++----------- internal/cli/client/client.go | 5 + .../cli/config/{machine.go => connection.go} | 0 4 files changed, 95 insertions(+), 111 deletions(-) rename internal/cli/config/{machine.go => connection.go} (100%) diff --git a/cmd/uncloud/machine/add.go b/cmd/uncloud/machine/add.go index 479cb9a7..fcff4a6c 100644 --- a/cmd/uncloud/machine/add.go +++ b/cmd/uncloud/machine/add.go @@ -1,15 +1,14 @@ package machine import ( - "context" + "fmt" "github.com/spf13/cobra" "uncloud/internal/cli" + "uncloud/internal/cli/config" ) type addOptions struct { name string - user string - port int sshKey string cluster string } @@ -17,25 +16,34 @@ type addOptions struct { func NewAddCommand() *cobra.Command { opts := addOptions{} cmd := &cobra.Command{ - // TODO: add support for [USER@]HOST[:PORT] syntax - Use: "add HOST", - Short: "Add a new machine to a cluster.", + Use: "add [USER@]HOST[:PORT]", + Short: "Add a remote machine to a cluster.", Args: cobra.ExactArgs(1), RunE: func(cmd *cobra.Command, args []string) error { uncli := cmd.Context().Value("cli").(*cli.CLI) - return add(cmd.Context(), uncli, args[0], opts) + + user, host, port, err := config.SSHDestination(args[0]).Parse() + if err != nil { + return fmt.Errorf("parse remote machine: %w", err) + } + remoteMachine := cli.RemoteMachine{ + User: user, + Host: host, + Port: port, + KeyPath: opts.sshKey, + } + + return uncli.AddMachine(cmd.Context(), remoteMachine, opts.cluster, opts.name) }, } cmd.Flags().StringVarP(&opts.name, "name", "n", "", "Assign a name to the machine") - cmd.Flags().StringVarP(&opts.user, "user", "u", "root", "Username for SSH remote login") - cmd.Flags().IntVarP(&opts.port, "port", "p", 22, "Port for SSH remote login") - cmd.Flags().StringVarP(&opts.sshKey, "ssh-key", "i", "", - "path to SSH private key for SSH remote login (default ~/.ssh/id_*)") - cmd.Flags().StringVarP(&opts.cluster, "cluster", "c", "", - "Name of the cluster to add the machine to (default is the current cluster)") + cmd.Flags().StringVarP( + &opts.sshKey, "ssh-key", "i", "", + "path to SSH private key for SSH remote login (default ~/.ssh/id_*)", + ) + cmd.Flags().StringVarP( + &opts.cluster, "cluster", "c", "", + "Name of the cluster to add the machine to (default is the current cluster)", + ) return cmd } - -func add(ctx context.Context, uncli *cli.CLI, host string, opts addOptions) error { - return uncli.AddMachine(ctx, opts.cluster, opts.name, opts.user, host, opts.port, opts.sshKey) -} diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 91ee97c3..942e0a9b 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -32,14 +32,14 @@ func New(configPath string) (*CLI, error) { }, nil } -func (cli *CLI) CreateCluster(name string, userPrivateKey secret.Secret) (*client.ClusterClient, error) { +func (cli *CLI) CreateCluster(name string, userPrivateKey secret.Secret) error { if _, ok := cli.config.Clusters[name]; ok { - return nil, fmt.Errorf("cluster %q already exists", name) + return fmt.Errorf("cluster %q already exists", name) } if userPrivateKey == nil { user, err := client.NewUser(nil) if err != nil { - return nil, fmt.Errorf("generate user: %w", err) + return fmt.Errorf("generate user: %w", err) } userPrivateKey = user.PrivateKey() } @@ -48,42 +48,14 @@ func (cli *CLI) CreateCluster(name string, userPrivateKey secret.Secret) (*clien Name: name, UserPrivateKey: userPrivateKey, } - if err := cli.config.Save(); err != nil { - return nil, err - } - - return cli.GetCluster(name) + return cli.config.Save() } -func (cli *CLI) CreateDefaultCluster() (*client.ClusterClient, error) { - c, err := cli.CreateCluster("default", nil) - if err != nil { - return nil, err +func (cli *CLI) CreateDefaultCluster() error { + if err := cli.CreateCluster(defaultClusterName, nil); err != nil { + return err } - if err = cli.SetCurrentCluster(c.Name()); err != nil { - return nil, err - } - return c, nil -} - -func (cli *CLI) GetCluster(name string) (*client.ClusterClient, error) { - clusterCfg, ok := cli.config.Clusters[name] - if !ok { - return nil, ErrNotFound - } - clusterCfg.Name = name - - user, err := client.NewUser(clusterCfg.UserPrivateKey) - if err != nil { - return nil, fmt.Errorf("create user: %w", err) - } - wgConnector := connector.NewWireGuardConnector(user, clusterCfg.Connections) - - return client.NewClusterClient(clusterCfg, wgConnector) -} - -func (cli *CLI) GetCurrentCluster() (*client.ClusterClient, error) { - return cli.GetCluster(cli.config.CurrentCluster) + return cli.SetCurrentCluster(defaultClusterName) } func (cli *CLI) SetCurrentCluster(name string) error { @@ -94,18 +66,6 @@ func (cli *CLI) SetCurrentCluster(name string) error { return cli.config.Save() } -func (cli *CLI) ListClusters() ([]*client.ClusterClient, error) { - var clusters []*client.ClusterClient - for name := range cli.config.Clusters { - c, err := cli.GetCluster(name) - if err != nil { - return nil, fmt.Errorf("get cluster %q: %w", name, err) - } - clusters = append(clusters, c) - } - return clusters, nil -} - func (cli *CLI) InitCluster( ctx context.Context, remoteMachine *RemoteMachine, clusterName, machineName string, netPrefix netip.Prefix, ) error { @@ -164,8 +124,7 @@ func (cli *CLI) initRemoteMachine( } fmt.Printf("Cluster %q initialised with machine %q\n", clusterName, resp.Machine.Name) - _, err = cli.CreateCluster(clusterName, user.PrivateKey()) - if err != nil { + if err = cli.CreateCluster(clusterName, user.PrivateKey()); err != nil { return fmt.Errorf("save cluster to config: %w", err) } // Set the current cluster to the just created one if it is the only cluster in the config. @@ -185,59 +144,71 @@ func (cli *CLI) initRemoteMachine( return nil } -func (cli *CLI) AddMachine( - ctx context.Context, clusterName, machineName, user, host string, port int, sshKeyPath string, -) error { - var ( - cluster *client.ClusterClient - err error - ) +func (cli *CLI) ConnectCluster(ctx context.Context, clusterName string) (*client.Client, error) { + if len(cli.config.Clusters) == 0 { + return nil, errors.New( + "no clusters found in the Uncloud config. " + + "Please initialise a cluster with `uncloud machine init` first", + ) + } if clusterName == "" { - // If the cluster is not specified, use the current cluster. If there are no clusters, create a default one. - cluster, err = cli.GetCurrentCluster() - if err != nil { - if errors.Is(err, ErrNotFound) { - // Do not create a default cluster if there are already clusters but the current cluster is not set. - clusters, cErr := cli.ListClusters() - if cErr != nil { - return fmt.Errorf("list clusters: %w", cErr) - } - if len(clusters) > 0 { - return errors.New( - "the current cluster is not set in the Uncloud config. " + - "Please specify a cluster with the --cluster flag or set current_cluster in the config", - ) - } - - cluster, err = cli.CreateDefaultCluster() - if err != nil { - return fmt.Errorf("create default cluster: %w", err) - } - fmt.Printf("Created %q cluster\n", cluster.Name()) - } else { - return fmt.Errorf("get current cluster: %w", err) - } + // If the cluster is not specified, use the current cluster if set. + if cli.config.CurrentCluster == "" { + return nil, errors.New( + "the current cluster is not set in the Uncloud config. " + + "Please specify a cluster with the --cluster flag or set current_cluster in the config", + ) } - } else { - cluster, err = cli.GetCluster(clusterName) - if err != nil { - return fmt.Errorf("get cluster %q: %w", clusterName, err) + if _, ok := cli.config.Clusters[cli.config.CurrentCluster]; !ok { + return nil, fmt.Errorf( + "current cluster %q not found in the config. "+ + "Please specify a cluster with the --cluster flag or update current_cluster in the config", + cli.config.CurrentCluster, + ) } + clusterName = cli.config.CurrentCluster } - defer func() { - _ = cluster.Close() - }() - name, connCfg, err := cluster.AddMachine(ctx, machineName, user, host, port, sshKeyPath) + cfg, ok := cli.config.Clusters[clusterName] + if !ok { + return nil, fmt.Errorf("cluster %q not found in the config", clusterName) + } + if len(cfg.Connections) == 0 { + return nil, fmt.Errorf("no connection configurations found for cluster %q in the config", clusterName) + } + + // TODO: iterate over all connections and try to connect to the cluster using the first successful connection. + conn := cfg.Connections[0] + user, host, port, err := conn.SSH.Parse() if err != nil { - return fmt.Errorf("add machine to cluster %q: %w", cluster.Name(), err) + return nil, fmt.Errorf("parse SSH connection %q: %w", conn.SSH, err) } - fmt.Printf("Machine %q added to cluster %q\n", name, cluster.Name()) + sshConfig := &connector.SSHConnectorConfig{ + User: user, + Host: host, + Port: port, + } + return client.New(ctx, connector.NewSSHConnector(sshConfig)) +} - cli.config.Clusters[cluster.Name()].Connections = append(cli.config.Clusters[cluster.Name()].Connections, connCfg) - if err = cli.config.Save(); err != nil { - return fmt.Errorf("save config: %w", err) +func (cli *CLI) AddMachine(ctx context.Context, remoteMachine RemoteMachine, clusterName, machineName string) error { + c, err := cli.ConnectCluster(ctx, clusterName) + if err != nil { + return fmt.Errorf("connect to cluster: %w", err) } + fmt.Println("Adding machine to cluster...", c) + // TODO + //name, connCfg, err := c.AddMachine(ctx, machineName, user, host, port, sshKeyPath) + //if err != nil { + // return fmt.Errorf("add machine to cluster %q: %w", cluster.Name(), err) + //} + //fmt.Printf("Machine %q added to cluster %q\n", name, cluster.Name()) + // + //cli.config.Clusters[cluster.Name()].Connections = append(cli.config.Clusters[cluster.Name()].Connections, connCfg) + //if err = cli.config.Save(); err != nil { + // return fmt.Errorf("save config: %w", err) + //} + return nil } diff --git a/internal/cli/client/client.go b/internal/cli/client/client.go index 1e225bf0..fdd225cf 100644 --- a/internal/cli/client/client.go +++ b/internal/cli/client/client.go @@ -8,11 +8,13 @@ import ( "uncloud/internal/machine/api/pb" ) +// Client is a client for the machine API. type Client struct { connector Connector conn *grpc.ClientConn pb.MachineClient + pb.ClusterClient } // Connector is an interface for establishing a connection to the machine API. @@ -21,6 +23,8 @@ type Connector interface { Close() error } +// New creates a new client for the machine API. The connector is used to establish the connection +// either locally or remotely. The client is responsible for closing the connector. func New(ctx context.Context, connector Connector) (*Client, error) { c := &Client{ connector: connector, @@ -32,6 +36,7 @@ func New(ctx context.Context, connector Connector) (*Client, error) { } c.MachineClient = pb.NewMachineClient(c.conn) + c.ClusterClient = pb.NewClusterClient(c.conn) return c, nil } diff --git a/internal/cli/config/machine.go b/internal/cli/config/connection.go similarity index 100% rename from internal/cli/config/machine.go rename to internal/cli/config/connection.go