Implement gtctl cluster connect mysql subcommand #3

pkg/cmd/gtctl/cluster/cluster.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/spf13/cobra"
 
+	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/cluster/connect"
 	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/cluster/create"
 	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/cluster/delete"
 	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/cluster/get"
@@ -47,6 +48,7 @@ func NewClusterCommand(l logger.Logger) *cobra.Command {
 	cmd.AddCommand(scale.NewScaleClusterCommand(l))
 	cmd.AddCommand(get.NewGetClusterCommand(l))
 	cmd.AddCommand(list.NewListClustersCommand(l))
+	cmd.AddCommand(connect.NewConnectCommand(l))
 
 	return cmd
 }

pkg/cmd/gtctl/cluster/connect/connect.go

@@ -0,0 +1,93 @@
+// Copyright 2023 Greptime Team
+//
+// 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 connect
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/types"
+
+	greptimedbclusterv1alpha1 "github.com/GreptimeTeam/greptimedb-operator/apis/v1alpha1"
+	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/cluster/connect/mysql"
+	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/cluster/connect/pg"
+	"github.com/GreptimeTeam/gtctl/pkg/deployer/k8s"
+	"github.com/GreptimeTeam/gtctl/pkg/logger"
+)
+
+type getClusterCliOptions struct {
+	Namespace string
+}
+
+func NewConnectCommand(l logger.Logger) *cobra.Command {
+	var options getClusterCliOptions
+	cmd := &cobra.Command{
+		Use:   "connect",
+		Short: "Connect to a GreptimeDB cluster",
+		Long:  `Connect to a GreptimeDB cluster`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return fmt.Errorf("cluster name should be set")
+			}
+
+			k8sDeployer, err := k8s.NewDeployer(l)
+			if err != nil {
+				return err
+			}
+
+			var (
+				ctx         = context.TODO()
+				clusterName = args[0]
+				namespace   = options.Namespace
+			)
+
+			name := types.NamespacedName{
+				Namespace: options.Namespace,
+				Name:      clusterName,
+			}.String()
+			cluster, err := k8sDeployer.GetGreptimeDBCluster(ctx, name, nil)
+			if err != nil && errors.IsNotFound(err) {
+				l.Errorf("cluster %s in %s not found\n", clusterName, namespace)
+				return nil
+			}
+			rawCluster, ok := cluster.Raw.(*greptimedbclusterv1alpha1.GreptimeDBCluster)
+			if !ok {
+				return fmt.Errorf("invalid cluster type")
+			}
+			dbType := cmd.Flag("p")
+			switch strings.ToLower(dbType.Value.String()) {
+			case "mysql":
+				err := mysql.ConnectCommand(rawCluster, l)
+				if err != nil {
+					_ = fmt.Errorf("error connecting to mysql: %v", err)
+				}
+			case "pg":
+				err := pg.ConnectCommand(rawCluster, l)
+				if err != nil {
+					_ = fmt.Errorf("error connecting to postgres: %v", err)
+				}
+			default:
+				return fmt.Errorf("database type not supported")
+			}
+			return nil
+		},
+	}
+	cmd.Flags().String("p", "mysql", "Specify a database")
+	cmd.Flags().StringVarP(&options.Namespace, "namespace", "n", "default", "Namespace of GreptimeDB cluster.")
+	return cmd
+}

pkg/cmd/gtctl/cluster/connect/mysql/mysql.go

@@ -0,0 +1,108 @@
+// Copyright 2023 Greptime Team
+//
+// 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 mysql
+
+import (
+	"context"
+	"database/sql"
+	"os"
+	"os/exec"
+	"strconv"
+	"sync"
+
+	"github.com/go-sql-driver/mysql"
+
+	greptimedbclusterv1alpha1 "github.com/GreptimeTeam/greptimedb-operator/apis/v1alpha1"
+	"github.com/GreptimeTeam/gtctl/pkg/logger"
+)
+
+const (
+	ArgHost = "-h"
+	ArgPort = "-P"
+	MySQL   = "mysql"
+)
+
+// ConnectCommand connects to a GreptimeDB cluster
+func ConnectCommand(rawCluster *greptimedbclusterv1alpha1.GreptimeDBCluster, l logger.Logger) error {
+	return connect("127.0.0.1", strconv.Itoa(int(rawCluster.Spec.MySQLServicePort)), rawCluster.Name, l)
+}
+
+// connect connects to a GreptimeDB cluster
+func connect(host, port, clusterName string, l logger.Logger) error {
+	waitGroup := sync.WaitGroup{}
+	cmd := exec.CommandContext(context.Background(), "kubectl", "port-forward", "-n", "default", "svc/"+clusterName+"-frontend", port+":"+port)
+	err := cmd.Start()
+	if err != nil {
+		l.Errorf("Error starting port-forwarding: %v", err)
+		return err
+	}
+	go func() {
+		waitGroup.Add(1)
+		defer waitGroup.Done()
+		if err = cmd.Wait(); err != nil {
+		}
+	}()
+	for {
+		cfg := mysql.Config{
+			Net:                  "tcp",
+			Addr:                 "127.0.0.1:4002",
+			User:                 "",
+			Passwd:               "",
+			DBName:               "",
+			AllowNativePasswords: true,
+		}
+
+		db, err := sql.Open("mysql", cfg.FormatDSN())
+		if err != nil {
+			continue
+		}
+
+		_, err = db.Conn(context.Background())
+		if err != nil {
+			continue
+		}
+
+		err = db.Close()
+		if err != nil {
+			l.V(1).Infof("Error closing connection: %v", err)
+			return err
+		}
+		break
+	}
+
+	cmd = exec.Command(MySQL, ArgHost, host, ArgPort, port)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.Stdin = os.Stdin
+	err = cmd.Start()
+	if err != nil {
+		l.Errorf("Error starting mysql client: %v", err)
+		return err
+	}
+	if err = cmd.Wait(); err != nil {
+		l.Errorf("Error waiting for mysql client to finish: %v", err)
+		return err
+	}
+	// gracefully stop port-forwarding
+	err = cmd.Process.Kill()
+	if err != nil {
+		if err.Error() != "os: process already finished" {
+			l.V(1).Info("Shutting down port-forwarding successfully")
+		}
+		return err
+	}
+	waitGroup.Wait()
+	return nil
+}

pkg/cmd/gtctl/cluster/connect/pg/pg.go

@@ -0,0 +1,24 @@
+// Copyright 2023 Greptime Team
+//
+// 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 pg
+
+import (
+	greptimedbclusterv1alpha1 "github.com/GreptimeTeam/greptimedb-operator/apis/v1alpha1"
+	"github.com/GreptimeTeam/gtctl/pkg/logger"
+)
+
+func ConnectCommand(rawCluster *greptimedbclusterv1alpha1.GreptimeDBCluster, l logger.Logger) error {
+	return nil
+}

pkg/cmd/gtctl/root.go

@@ -16,7 +16,6 @@ package gtctl
 
 import (
 	"fmt"
-	"github.com/GreptimeTeam/gtctl/pkg/cmd/gtctl/connect"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -63,7 +62,6 @@ func NewRootCommand() *cobra.Command {
 	// Add all top level subcommands.
 	cmd.AddCommand(version.NewVersionCommand(l))
 	cmd.AddCommand(cluster.NewClusterCommand(l))
-	cmd.AddCommand(connect.NewConnectCommand(l))
 
 	return cmd
 }