Skip to content

Commit

Permalink
polish export prompts and use more SDK methods (#233)
Browse files Browse the repository at this point in the history
  • Loading branch information
@zhangyangyu
zhangyangyu authored Sep 5, 2024
1 parent 2247735 commit e1a0def
Show file tree
Hide file tree
Showing 9 changed files with 126 additions and 201 deletions.
30 changes: 11 additions & 19 deletions internal/cli/serverless/dataimport/start/azblob.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ package start

import (
"fmt"
"slices"

"tidbcloud-cli/internal"
"tidbcloud-cli/internal/flag"
Expand All @@ -36,18 +35,10 @@ type AzBlobOpts struct {
clusterId string
}

func (o AzBlobOpts) SupportedFileTypes() []string {
return []string{
string(imp.IMPORTFILETYPEENUM_CSV),
string(imp.IMPORTFILETYPEENUM_PARQUET),
string(imp.IMPORTFILETYPEENUM_SQL),
string(imp.IMPORTFILETYPEENUM_AURORA_SNAPSHOT),
}
}

func (o AzBlobOpts) Run(cmd *cobra.Command) error {
ctx := cmd.Context()
var fileType, uri, sasToken string
var uri, sasToken string
var fileType imp.ImportFileTypeEnum
var authType imp.ImportAzureBlobAuthTypeEnum
var format *imp.CSVFormat
d, err := o.h.Client()
Expand Down Expand Up @@ -91,7 +82,7 @@ func (o AzBlobOpts) Run(cmd *cobra.Command) error {
}

var fileTypes []interface{}
for _, f := range o.SupportedFileTypes() {
for _, f := range imp.AllowedImportFileTypeEnumEnumValues {
fileTypes = append(fileTypes, f)
}
model, err = ui.InitialSelectModel(fileTypes, "Choose the source file type:")
Expand All @@ -106,22 +97,23 @@ func (o AzBlobOpts) Run(cmd *cobra.Command) error {
if m, _ := fileTypeModel.(ui.SelectModel); m.Interrupted {
return util.InterruptError
}
fileType = fileTypeModel.(ui.SelectModel).Choices[fileTypeModel.(ui.SelectModel).Selected].(string)
fileType = fileTypeModel.(ui.SelectModel).Choices[fileTypeModel.(ui.SelectModel).Selected].(imp.ImportFileTypeEnum)

if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
format, err = getCSVFormat()
if err != nil {
return err
}
}
} else {
// non-interactive mode
fileType, err = cmd.Flags().GetString(flag.FileType)
fileTypeStr, err := cmd.Flags().GetString(flag.FileType)
if err != nil {
return errors.Trace(err)
}
if !slices.Contains(o.SupportedFileTypes(), fileType) {
return fmt.Errorf("file type \"%s\" is not supported, please use one of %q", fileType, o.SupportedFileTypes())
fileType = imp.ImportFileTypeEnum(fileTypeStr)
if !fileType.IsValid() {
return fmt.Errorf("file type \"%s\" is not supported, please use one of %q", fileTypeStr, imp.AllowedImportFileTypeEnumEnumValues)
}

uri, err = cmd.Flags().GetString(flag.AzureBlobURI)
Expand All @@ -142,7 +134,7 @@ func (o AzBlobOpts) Run(cmd *cobra.Command) error {
authType = imp.IMPORTAZUREBLOBAUTHTYPEENUM_SAS_TOKEN

// optional flags
if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
format, err = getCSVFlagValue(cmd)
if err != nil {
return errors.Trace(err)
Expand All @@ -155,7 +147,7 @@ func (o AzBlobOpts) Run(cmd *cobra.Command) error {
source.AzureBlob.AuthType = authType
source.AzureBlob.SasToken = &sasToken
options := imp.NewImportOptions(imp.ImportFileTypeEnum(fileType))
if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
options.CsvFormat = format
}
body := imp.NewImportServiceCreateImportBody(*options, *source)
Expand Down
2 changes: 1 addition & 1 deletion internal/cli/serverless/dataimport/start/azblob_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ func (suite *AzblobImportSuite) TestAzblobImportArgs() {
{
name: "start import with unsupported data format",
args: []string{"--source-type", "AZURE_BLOB", "--cluster-id", clusterID, "--file-type", "yaml"},
err: fmt.Errorf("file type \"yaml\" is not supported, please use one of [\"CSV\" \"PARQUET\" \"SQL\" \"AURORA_SNAPSHOT\"]"),
err: fmt.Errorf("file type \"yaml\" is not supported, please use one of [\"CSV\" \"SQL\" \"AURORA_SNAPSHOT\" \"PARQUET\"]"),
},
{
name: "start import without required cluster id",
Expand Down
30 changes: 11 additions & 19 deletions internal/cli/serverless/dataimport/start/gcs.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ package start

import (
"fmt"
"slices"

"tidbcloud-cli/internal"
"tidbcloud-cli/internal/flag"
Expand All @@ -37,18 +36,10 @@ type GCSOpts struct {
clusterId string
}

func (o GCSOpts) SupportedFileTypes() []string {
return []string{
string(imp.IMPORTFILETYPEENUM_CSV),
string(imp.IMPORTFILETYPEENUM_PARQUET),
string(imp.IMPORTFILETYPEENUM_SQL),
string(imp.IMPORTFILETYPEENUM_AURORA_SNAPSHOT),
}
}

func (o GCSOpts) Run(cmd *cobra.Command) error {
ctx := cmd.Context()
var fileType, gcsUri, accountKey string
var gcsUri, accountKey string
var fileType imp.ImportFileTypeEnum
var authType imp.ImportGcsAuthTypeEnum
var format *imp.CSVFormat
d, err := o.h.Client()
Expand Down Expand Up @@ -92,7 +83,7 @@ func (o GCSOpts) Run(cmd *cobra.Command) error {
}

var fileTypes []interface{}
for _, f := range o.SupportedFileTypes() {
for _, f := range imp.AllowedImportFileTypeEnumEnumValues {
fileTypes = append(fileTypes, f)
}
model, err = ui.InitialSelectModel(fileTypes, "Choose the source file type:")
Expand All @@ -107,22 +98,23 @@ func (o GCSOpts) Run(cmd *cobra.Command) error {
if m, _ := fileTypeModel.(ui.SelectModel); m.Interrupted {
return util.InterruptError
}
fileType = fileTypeModel.(ui.SelectModel).Choices[fileTypeModel.(ui.SelectModel).Selected].(string)
fileType = fileTypeModel.(ui.SelectModel).Choices[fileTypeModel.(ui.SelectModel).Selected].(imp.ImportFileTypeEnum)

if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
format, err = getCSVFormat()
if err != nil {
return err
}
}
} else {
// non-interactive mode
fileType, err = cmd.Flags().GetString(flag.FileType)
fileTypeStr, err := cmd.Flags().GetString(flag.FileType)
if err != nil {
return errors.Trace(err)
}
if !slices.Contains(o.SupportedFileTypes(), fileType) {
return fmt.Errorf("file type \"%s\" is not supported, please use one of %q", fileType, o.SupportedFileTypes())
fileType = imp.ImportFileTypeEnum(fileTypeStr)
if !fileType.IsValid() {
return fmt.Errorf("file type \"%s\" is not supported, please use one of %q", fileTypeStr, imp.AllowedImportFileTypeEnumEnumValues)
}
gcsUri, err = cmd.Flags().GetString(flag.GCSURI)
if err != nil {
Expand All @@ -142,7 +134,7 @@ func (o GCSOpts) Run(cmd *cobra.Command) error {
authType = imp.IMPORTGCSAUTHTYPEENUM_SERVICE_ACCOUNT_KEY

// optional flags
if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
format, err = getCSVFlagValue(cmd)
if err != nil {
return errors.Trace(err)
Expand All @@ -154,7 +146,7 @@ func (o GCSOpts) Run(cmd *cobra.Command) error {
source.Gcs = imp.NewGCSSource(gcsUri, authType)
source.Gcs.ServiceAccountKey = aws.String(accountKey)
options := imp.NewImportOptions(imp.ImportFileTypeEnum(fileType))
if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
options.CsvFormat = format
}
body := imp.NewImportServiceCreateImportBody(*options, *source)
Expand Down
2 changes: 1 addition & 1 deletion internal/cli/serverless/dataimport/start/gcs_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ func (suite *GCSImportSuite) TestGCSImportArgs() {
{
name: "start import with unsupported data format",
args: []string{"--source-type", "GCS", "--cluster-id", clusterID, "--file-type", "yaml"},
err: fmt.Errorf("file type \"yaml\" is not supported, please use one of [\"CSV\" \"PARQUET\" \"SQL\" \"AURORA_SNAPSHOT\"]"),
err: fmt.Errorf("file type \"yaml\" is not supported, please use one of [\"CSV\" \"SQL\" \"AURORA_SNAPSHOT\" \"PARQUET\"]"),
},
{
name: "start import without required cluster id",
Expand Down
30 changes: 11 additions & 19 deletions internal/cli/serverless/dataimport/start/s3.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ package start

import (
"fmt"
"slices"

"tidbcloud-cli/internal"
"tidbcloud-cli/internal/flag"
Expand All @@ -35,18 +34,10 @@ type S3Opts struct {
clusterId string
}

func (o S3Opts) SupportedFileTypes() []string {
return []string{
string(imp.IMPORTFILETYPEENUM_CSV),
string(imp.IMPORTFILETYPEENUM_PARQUET),
string(imp.IMPORTFILETYPEENUM_SQL),
string(imp.IMPORTFILETYPEENUM_AURORA_SNAPSHOT),
}
}

func (o S3Opts) Run(cmd *cobra.Command) error {
ctx := cmd.Context()
var fileType, s3Uri, s3Arn, accessKeyID, secretAccessKey string
var s3Uri, s3Arn, accessKeyID, secretAccessKey string
var fileType imp.ImportFileTypeEnum
var authType imp.ImportS3AuthTypeEnum
var format *imp.CSVFormat
d, err := o.h.Client()
Expand Down Expand Up @@ -108,7 +99,7 @@ func (o S3Opts) Run(cmd *cobra.Command) error {
}

var fileTypes []interface{}
for _, f := range o.SupportedFileTypes() {
for _, f := range imp.AllowedImportFileTypeEnumEnumValues {
fileTypes = append(fileTypes, f)
}
model, err = ui.InitialSelectModel(fileTypes, "Choose the source file type:")
Expand All @@ -123,22 +114,23 @@ func (o S3Opts) Run(cmd *cobra.Command) error {
if m, _ := fileTypeModel.(ui.SelectModel); m.Interrupted {
return util.InterruptError
}
fileType = fileTypeModel.(ui.SelectModel).Choices[fileTypeModel.(ui.SelectModel).Selected].(string)
fileType = fileTypeModel.(ui.SelectModel).Choices[fileTypeModel.(ui.SelectModel).Selected].(imp.ImportFileTypeEnum)

if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
format, err = getCSVFormat()
if err != nil {
return err
}
}
} else {
// non-interactive mode
fileType, err = cmd.Flags().GetString(flag.FileType)
fileTypeStr, err := cmd.Flags().GetString(flag.FileType)
if err != nil {
return errors.Trace(err)
}
if !slices.Contains(o.SupportedFileTypes(), fileType) {
return fmt.Errorf("file type \"%s\" is not supported, please use one of %q", fileType, o.SupportedFileTypes())
fileType = imp.ImportFileTypeEnum(fileTypeStr)
if !fileType.IsValid() {
return fmt.Errorf("file type \"%s\" is not supported, please use one of %q", fileTypeStr, imp.AllowedImportFileTypeEnumEnumValues)
}
s3Uri, err = cmd.Flags().GetString(flag.S3URI)
if err != nil {
Expand All @@ -149,7 +141,7 @@ func (o S3Opts) Run(cmd *cobra.Command) error {
}

// optional flags
if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
format, err = getCSVFlagValue(cmd)
if err != nil {
return errors.Trace(err)
Expand Down Expand Up @@ -189,7 +181,7 @@ func (o S3Opts) Run(cmd *cobra.Command) error {
}
}
options := imp.NewImportOptions(imp.ImportFileTypeEnum(fileType))
if fileType == string(imp.IMPORTFILETYPEENUM_CSV) {
if fileType == imp.IMPORTFILETYPEENUM_CSV {
options.CsvFormat = format
}
body := imp.NewImportServiceCreateImportBody(*options, *source)
Expand Down
2 changes: 1 addition & 1 deletion internal/cli/serverless/dataimport/start/s3_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ func (suite *S3ImportSuite) TestS3ImportArgs() {
{
name: "start import with unsupported data format",
args: []string{"--source-type", "S3", "--cluster-id", clusterID, "--file-type", "yaml"},
err: fmt.Errorf("file type \"yaml\" is not supported, please use one of [\"CSV\" \"PARQUET\" \"SQL\" \"AURORA_SNAPSHOT\"]"),
err: fmt.Errorf("file type \"yaml\" is not supported, please use one of [\"CSV\" \"SQL\" \"AURORA_SNAPSHOT\" \"PARQUET\"]"),
},
{
name: "start import without required cluster id",
Expand Down
24 changes: 4 additions & 20 deletions internal/cli/serverless/dataimport/start/start.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,26 +66,10 @@ var inputDescription = map[string]string{
flag.CSVNotNull: "Input the CSV not-null: whether the CSV can contains any NULL value, skip to use default value (false)",
}

var sourceTypes = []imp.ImportSourceTypeEnum{
imp.IMPORTSOURCETYPEENUM_S3,
imp.IMPORTSOURCETYPEENUM_LOCAL,
imp.IMPORTSOURCETYPEENUM_GCS,
imp.IMPORTSOURCETYPEENUM_AZURE_BLOB,
}

type StartOpts struct {
interactive bool
}

func (o StartOpts) SupportedFileTypes() []string {
return []string{
string(imp.IMPORTFILETYPEENUM_CSV),
string(imp.IMPORTFILETYPEENUM_PARQUET),
string(imp.IMPORTFILETYPEENUM_SQL),
string(imp.IMPORTFILETYPEENUM_AURORA_SNAPSHOT),
}
}

func (o StartOpts) NonInteractiveFlags() []string {
return []string{
flag.ClusterID,
Expand Down Expand Up @@ -249,8 +233,8 @@ func StartCmd(h *internal.Helper) *cobra.Command {
}

startCmd.Flags().StringP(flag.ClusterID, flag.ClusterIDShort, "", "Cluster ID.")
startCmd.Flags().String(flag.SourceType, "LOCAL", fmt.Sprintf("The import source type, one of %q.", sourceTypes))
startCmd.Flags().String(flag.FileType, "", fmt.Sprintf("The import file type, one of %q.", opts.SupportedFileTypes()))
startCmd.Flags().String(flag.SourceType, "LOCAL", fmt.Sprintf("The import source type, one of %q.", imp.AllowedImportSourceTypeEnumEnumValues))
startCmd.Flags().String(flag.FileType, "", fmt.Sprintf("The import file type, one of %q.", imp.AllowedImportFileTypeEnumEnumValues))

startCmd.Flags().String(flag.LocalFilePath, "", "The local file path to import.")
startCmd.Flags().String(flag.LocalTargetDatabase, "", "Target database to which import data.")
Expand Down Expand Up @@ -282,8 +266,8 @@ func StartCmd(h *internal.Helper) *cobra.Command {
}

func getSelectedSourceType() (imp.ImportSourceTypeEnum, error) {
SourceTypes := make([]interface{}, 0, len(sourceTypes))
for _, sourceType := range sourceTypes {
SourceTypes := make([]interface{}, 0, len(imp.AllowedImportSourceTypeEnumEnumValues))
for _, sourceType := range imp.AllowedImportSourceTypeEnumEnumValues {
SourceTypes = append(SourceTypes, sourceType)
}
model, err := ui.InitialSelectModel(SourceTypes, "Choose import source type:")
Expand Down
Loading

0 comments on commit e1a0def

Please sign in to comment.