Export support ali cluster and OSS target (#275)

Author: shiyuhang0

Diff for: docs/generate_doc/ticloud_serverless_export_create.md

@@ -46,13 +46,16 @@ ticloud serverless export create [flags]
       --gcs.service-account-key string   The base64 encoded service account key of GCS.
       --gcs.uri string                   The GCS URI in gs://<bucket>/<path> format. Required when target type is GCS.
   -h, --help                             help for create
+      --oss.access-key-id string         The access key ID of the OSS.
+      --oss.access-key-secret string     The access key secret of the OSS.
+      --oss.uri string                   The OSS URI in oss://<bucket>/<path> format. Required when target type is OSS.
       --parquet.compression string       The parquet compression algorithm. One of ["GZIP" "SNAPPY" "ZSTD" "NONE"]. (default "ZSTD")
       --s3.access-key-id string          The access key ID of the S3. You only need to set one of the s3.role-arn and [s3.access-key-id, s3.secret-access-key].
       --s3.role-arn string               The role arn of the S3. You only need to set one of the s3.role-arn and [s3.access-key-id, s3.secret-access-key].
       --s3.secret-access-key string      The secret access key of the S3. You only need to set one of the s3.role-arn and [s3.access-key-id, s3.secret-access-key].
       --s3.uri string                    The S3 URI in s3://<bucket>/<path> format. Required when target type is S3.
       --sql string                       Filter the exported data with SQL SELECT statement.
-      --target-type string               The export target. One of ["LOCAL" "S3" "GCS" "AZURE_BLOB"]. (default "LOCAL")
+      --target-type string               The export target. One of ["LOCAL" "S3" "GCS" "AZURE_BLOB" "OSS"]. (default "LOCAL")
       --where string                     Filter the exported table(s) with the where condition.
 ```
 

Diff for: internal/cli/serverless/export/create.go

@@ -70,6 +70,9 @@ func (c CreateOpts) NonInteractiveFlags() []string {
 		flag.AzureBlobSASToken,
 		flag.ParquetCompression,
 		flag.DisplayName,
+		flag.OSSURI,
+		flag.OSSAccessKeyID,
+		flag.OSSAccessKeySecret,
 	}
 }
 
@@ -155,6 +158,8 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 			var gcsURI, gcsServiceAccountKey string
 			// azure
 			var azBlobURI, azBlobSasToken string
+			// oss
+			var ossURI, ossAccessKeyID, ossAccessKeySecret string
 			var displayName string
 
 			if opts.interactive {
@@ -187,28 +192,50 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 				if err != nil {
 					return err
 				}
-				selectedAuthType, err := GetSelectedAuthType(targetType)
+				selectedAuthType, err := GetSelectedAuthType(targetType, *cluster.CloudProvider)
 				if err != nil {
 					return err
 				}
 				switch selectedAuthType {
+				// Both S3 and OSS supports ACCESS_KEY
 				case string(export.EXPORTS3AUTHTYPEENUM_ACCESS_KEY):
-					inputs := []string{flag.S3URI, flag.S3AccessKeyID, flag.S3SecretAccessKey}
-					textInput, err := ui.InitialInputModel(inputs, inputDescription)
-					if err != nil {
-						return err
-					}
-					s3URI = textInput.Inputs[0].Value()
-					if s3URI == "" {
-						return errors.New("empty S3 URI")
-					}
-					accessKeyID = textInput.Inputs[1].Value()
-					if accessKeyID == "" {
-						return errors.New("empty S3 access key Id")
+					if targetType == export.EXPORTTARGETTYPEENUM_S3 {
+						inputs := []string{flag.S3URI, flag.S3AccessKeyID, flag.S3SecretAccessKey}
+						textInput, err := ui.InitialInputModel(inputs, inputDescription)
+						if err != nil {
+							return err
+						}
+						s3URI = textInput.Inputs[0].Value()
+						if s3URI == "" {
+							return errors.New("empty S3 URI")
+						}
+						accessKeyID = textInput.Inputs[1].Value()
+						if accessKeyID == "" {
+							return errors.New("empty S3 access key Id")
+						}
+						secretAccessKey = textInput.Inputs[2].Value()
+						if secretAccessKey == "" {
+							return errors.New("empty S3 secret access key")
+						}
 					}
-					secretAccessKey = textInput.Inputs[2].Value()
-					if secretAccessKey == "" {
-						return errors.New("empty S3 secret access key")
+					if targetType == export.EXPORTTARGETTYPEENUM_OSS {
+						inputs := []string{flag.OSSURI, flag.OSSAccessKeyID, flag.OSSAccessKeySecret}
+						textInput, err := ui.InitialInputModel(inputs, inputDescription)
+						if err != nil {
+							return err
+						}
+						ossURI = textInput.Inputs[0].Value()
+						if ossURI == "" {
+							return errors.New("empty OSS URI")
+						}
+						ossAccessKeyID = textInput.Inputs[1].Value()
+						if ossAccessKeyID == "" {
+							return errors.New("empty OSS access key Id")
+						}
+						ossAccessKeySecret = textInput.Inputs[2].Value()
+						if ossAccessKeySecret == "" {
+							return errors.New("empty OSS access key secret")
+						}
 					}
 				case string(export.EXPORTS3AUTHTYPEENUM_ROLE_ARN):
 					inputs := []string{flag.S3URI, flag.S3RoleArn}
@@ -463,6 +490,25 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 					if azBlobSasToken == "" {
 						return errors.New("Azure Blob SAS token is required when target type is AZURE_BLOB")
 					}
+				case export.EXPORTTARGETTYPEENUM_OSS:
+					ossURI, err = cmd.Flags().GetString(flag.OSSURI)
+					if err != nil {
+						return errors.Trace(err)
+					}
+					if ossURI == "" {
+						return errors.New("OSS URI is required when target type is OSS")
+					}
+					ossAccessKeyID, err = cmd.Flags().GetString(flag.OSSAccessKeyID)
+					if err != nil {
+						return errors.Trace(err)
+					}
+					ossAccessKeySecret, err = cmd.Flags().GetString(flag.OSSAccessKeySecret)
+					if err != nil {
+						return errors.Trace(err)
+					}
+					if ossAccessKeyID == "" || ossAccessKeySecret == "" {
+						return errors.New("OSS access key id and access key secret are required when target type is OSS")
+					}
 				}
 
 				compressionStr, err := cmd.Flags().GetString(flag.Compression)
@@ -603,6 +649,15 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 					AuthType: export.EXPORTAZUREBLOBAUTHTYPEENUM_SAS_TOKEN,
 					SasToken: &azBlobSasToken,
 				}
+			case export.EXPORTTARGETTYPEENUM_OSS:
+				params.Target.Oss = &export.OSSTarget{
+					Uri:      ossURI,
+					AuthType: export.EXPORTOSSAUTHTYPEENUM_ACCESS_KEY,
+					AccessKey: &export.OSSTargetAccessKey{
+						Id:     ossAccessKeyID,
+						Secret: ossAccessKeySecret,
+					},
+				}
 			}
 			// add compression
 			if compression != "" {
@@ -669,6 +724,9 @@ func CreateCmd(h *internal.Helper) *cobra.Command {
 	createCmd.Flags().String(flag.GCSServiceAccountKey, "", "The base64 encoded service account key of GCS.")
 	createCmd.Flags().String(flag.AzureBlobURI, "", "The Azure Blob URI in azure://<account>.blob.core.windows.net/<container>/<path> format. Required when target type is AZURE_BLOB.")
 	createCmd.Flags().String(flag.AzureBlobSASToken, "", "The SAS token of Azure Blob.")
+	createCmd.Flags().String(flag.OSSURI, "", "The OSS URI in oss://<bucket>/<path> format. Required when target type is OSS.")
+	createCmd.Flags().String(flag.OSSAccessKeyID, "", "The access key ID of the OSS.")
+	createCmd.Flags().String(flag.OSSAccessKeySecret, "", "The access key secret of the OSS.")
 	createCmd.Flags().String(flag.ParquetCompression, "ZSTD", fmt.Sprintf("The parquet compression algorithm. One of %q.", export.AllowedExportParquetCompressionTypeEnumEnumValues))
 	createCmd.Flags().String(flag.DisplayName, "", "The display name of the export. (default \"SNAPSHOT_<snapshot_time>\")")
 

Diff for: internal/cli/serverless/export/create_test.go

@@ -373,6 +373,86 @@ func (suite *CreateExportSuite) TestCreateExportWithTableFilter() {
 	suite.AssertTest(ctx, tests)
 }
 
+func (suite *CreateExportSuite) TestCreateExportToOSS() {
+	ctx := context.Background()
+
+	clusterId := "fake-cluster-id"
+	exportId := "fake-export-id"
+	targetType := export.EXPORTTARGETTYPEENUM_OSS
+	uri := "s3://fake-bucket/fake-prefix"
+	accessKeyId := "fake-id"
+	secretAccess := "fake-secret"
+
+	body := getDefaultCreateExportBody()
+	body.Target = &export.ExportTarget{
+		Type: &targetType,
+		Oss: &export.OSSTarget{
+			Uri:      uri,
+			AuthType: export.EXPORTOSSAUTHTYPEENUM_ACCESS_KEY,
+			AccessKey: &export.OSSTargetAccessKey{
+				Id:     accessKeyId,
+				Secret: secretAccess,
+			},
+		},
+	}
+	suite.mockClient.On("CreateExport", ctx, clusterId, body).
+		Return(&export.Export{
+			ExportId: &exportId,
+		}, nil)
+
+	tests := []Test{
+		{
+			name: "export data to oss without uri",
+			args: []string{"-c", clusterId, "--target-type", "OSS", "--oss.access-key-id", accessKeyId, "--oss.access-key-secret", secretAccess, "--force"},
+			err:  errors.New("OSS URI is required when target type is OSS"),
+		},
+		{
+			name: "export data to oss without auth",
+			args: []string{"-c", clusterId, "--target-type", "OSS", "--oss.uri", uri, "--oss.access-key-id", accessKeyId, "--force"},
+			err:  errors.New("OSS access key id and access key secret are required when target type is OSS"),
+		},
+		{
+			name:         "export all data to oss using access key",
+			args:         []string{"-c", clusterId, "--target-type", "OSS", "--oss.uri", uri, "--oss.access-key-id", accessKeyId, "--oss.access-key-secret", secretAccess, "--force"},
+			stdoutString: fmt.Sprintf("export %s is running now\n", exportId),
+		},
+	}
+	suite.AssertTest(ctx, tests)
+}
+
+func (suite *CreateExportSuite) TestAliClusterExportToS3WithRoleArn() {
+	ctx := context.Background()
+
+	clusterId := "fake-cluster-id"
+	targetType := export.EXPORTTARGETTYPEENUM_S3
+	uri := "s3://fake-bucket/fake-prefix"
+	roleArn := "fake-role-arn"
+
+	body := getDefaultCreateExportBody()
+	body.Target = &export.ExportTarget{
+		Type: &targetType,
+		S3: &export.S3Target{
+			Uri:      &uri,
+			AuthType: export.EXPORTS3AUTHTYPEENUM_ROLE_ARN,
+			RoleArn:  &roleArn,
+		},
+	}
+
+	errMsg := "export to s3 with role arn is not supported in alicloud"
+
+	suite.mockClient.On("CreateExport", ctx, clusterId, body).
+		Return(nil, errors.New(errMsg))
+
+	tests := []Test{
+		{
+			name: "export ali cluster with S3 role arn",
+			args: []string{"-c", clusterId, "--target-type", "S3", "--s3.uri", uri, "--s3.role-arn", roleArn, "--force"},
+			err:  errors.New(errMsg),
+		},
+	}
+	suite.AssertTest(ctx, tests)
+}
+
 func getDefaultCreateExportBody() *export.ExportServiceCreateExportBody {
 	defaultFileType := export.EXPORTFILETYPEENUM_CSV
 	defaultTargetType := export.EXPORTTARGETTYPEENUM_LOCAL

Diff for: internal/cli/serverless/export/ui.go

@@ -15,15 +15,15 @@
 package export
 
 import (
+	"github.com/charmbracelet/bubbles/textinput"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/juju/errors"
 	"github.com/tidbcloud/tidbcloud-cli/internal/config"
 	"github.com/tidbcloud/tidbcloud-cli/internal/flag"
 	"github.com/tidbcloud/tidbcloud-cli/internal/ui"
 	"github.com/tidbcloud/tidbcloud-cli/internal/util"
+	"github.com/tidbcloud/tidbcloud-cli/pkg/tidbcloud/v1beta1/serverless/cluster"
 	"github.com/tidbcloud/tidbcloud-cli/pkg/tidbcloud/v1beta1/serverless/export"
-
-	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
-	"github.com/juju/errors"
 )
 
 var inputDescription = map[string]string{
@@ -43,6 +43,9 @@ var inputDescription = map[string]string{
 	flag.CSVNullValue:         "Input the CSV null value: representation of null values in CSV files, skip to use default value (\\N). If you want to set empty string, please use non-interactive mode",
 	flag.CSVSkipHeader:        "Input the CSV skip header: export CSV files of the tables without header. Type `true` to skip header, others will not skip header",
 	flag.DisplayName:          "Input the name of export. You can skip and use the default name SNAPSHOT_<snapshot_time> by pressing Enter",
+	flag.OSSURI:               "Input your OSS URI in oss://<bucket>/<path> format",
+	flag.OSSAccessKeyID:       "Input your OSS access key id",
+	flag.OSSAccessKeySecret:   "Input your OSS access key secret",
 }
 
 func GetSelectedParquetCompression() (export.ExportParquetCompressionTypeEnum, error) {
@@ -70,7 +73,9 @@ func GetSelectedParquetCompression() (export.ExportParquetCompressionTypeEnum, e
 
 func GetSelectedTargetType() (export.ExportTargetTypeEnum, error) {
 	targetTypes := make([]interface{}, 0, 4)
-	targetTypes = append(targetTypes, export.EXPORTTARGETTYPEENUM_LOCAL, export.EXPORTTARGETTYPEENUM_S3, export.EXPORTTARGETTYPEENUM_GCS, export.EXPORTTARGETTYPEENUM_AZURE_BLOB)
+	for _, v := range export.AllowedExportTargetTypeEnumEnumValues {
+		targetTypes = append(targetTypes, v)
+	}
 	model, err := ui.InitialSelectModel(targetTypes, "Choose the export target:")
 	if err != nil {
 		return "", errors.Trace(err)
@@ -91,22 +96,28 @@ func GetSelectedTargetType() (export.ExportTargetTypeEnum, error) {
 	return targetType.(export.ExportTargetTypeEnum), nil
 }
 
-func GetSelectedAuthType(target export.ExportTargetTypeEnum) (_ string, err error) {
+func GetSelectedAuthType(target export.ExportTargetTypeEnum, provider cluster.V1beta1RegionCloudProvider) (_ string, err error) {
 	var model *ui.SelectModel
 	switch target {
 	case export.EXPORTTARGETTYPEENUM_S3:
+		if provider != cluster.V1BETA1REGIONCLOUDPROVIDER_AWS {
+			return string(export.EXPORTS3AUTHTYPEENUM_ACCESS_KEY), nil
+		}
 		authTypes := make([]interface{}, 0, 2)
 		authTypes = append(authTypes, string(export.EXPORTS3AUTHTYPEENUM_ROLE_ARN), string(export.EXPORTS3AUTHTYPEENUM_ACCESS_KEY))
 		model, err = ui.InitialSelectModel(authTypes, "Choose and input the S3 auth:")
 		if err != nil {
 			return "", errors.Trace(err)
 		}
+
 	case export.EXPORTTARGETTYPEENUM_GCS:
 		return string(export.EXPORTGCSAUTHTYPEENUM_SERVICE_ACCOUNT_KEY), nil
 	case export.EXPORTTARGETTYPEENUM_AZURE_BLOB:
 		return string(export.EXPORTAZUREBLOBAUTHTYPEENUM_SAS_TOKEN), nil
 	case export.EXPORTTARGETTYPEENUM_LOCAL:
 		return "", nil
+	case export.EXPORTTARGETTYPEENUM_OSS:
+		return string(export.EXPORTOSSAUTHTYPEENUM_ACCESS_KEY), nil
 	}
 	if model == nil {
 		return "", errors.New("unknown auth type")