// Copyright 2019 Dolthub, Inc. // // 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 tblcmds import ( "context" "os" "path/filepath" "github.com/dolthub/go-mysql-server/sql" "github.com/fatih/color" "github.com/dolthub/dolt/go/cmd/dolt/cli" "github.com/dolthub/dolt/go/cmd/dolt/commands" "github.com/dolthub/dolt/go/cmd/dolt/commands/engine" "github.com/dolthub/dolt/go/cmd/dolt/errhand" "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" "github.com/dolthub/dolt/go/libraries/doltcore/env" "github.com/dolthub/dolt/go/libraries/doltcore/mvdata" "github.com/dolthub/dolt/go/libraries/doltcore/schema" "github.com/dolthub/dolt/go/libraries/doltcore/table" "github.com/dolthub/dolt/go/libraries/doltcore/table/editor" "github.com/dolthub/dolt/go/libraries/utils/argparser" "github.com/dolthub/dolt/go/libraries/utils/filesys" "github.com/dolthub/dolt/go/libraries/utils/iohelp" eventsapi "github.com/dolthub/eventsapi_schema/dolt/services/eventsapi/v1alpha1" ) var exportDocs = cli.CommandDocumentationContent{ ShortDesc: `Export the contents of a table to a file.`, LongDesc: `{{.EmphasisLeft}}dolt table export{{.EmphasisRight}} will export the contents of {{.LessThan}}table{{.GreaterThan}} to {{.LessThan}}|file{{.GreaterThan}} The output format is inferred from the file extension, or can be set explicitly with {{.EmphasisLeft}}--file-type{{.EmphasisRight}}. Supported file types: {{.EmphasisLeft}}csv{{.EmphasisRight}}, {{.EmphasisLeft}}psv{{.EmphasisRight}}, {{.EmphasisLeft}}json{{.EmphasisRight}}, {{.EmphasisLeft}}jsonl{{.EmphasisRight}}, {{.EmphasisLeft}}sql{{.EmphasisRight}}, {{.EmphasisLeft}}parquet{{.EmphasisRight}}. {{.EmphasisLeft}}.json{{.EmphasisRight}} exports a single JSON object containing a {{.EmphasisLeft}}rows{{.EmphasisRight}} array; {{.EmphasisLeft}}.jsonl{{.EmphasisRight}} exports one JSON object per line. See the help for {{.EmphasisLeft}}dolt table import{{.EmphasisRight}} as the options are the same. `, Synopsis: []string{ "[-f] [-pk {{.LessThan}}field{{.GreaterThan}}] [-schema {{.LessThan}}file{{.GreaterThan}}] [-map {{.LessThan}}file{{.GreaterThan}}] [-continue] [-file-type {{.LessThan}}type{{.GreaterThan}}] {{.LessThan}}table{{.GreaterThan}} {{.LessThan}}file{{.GreaterThan}}", }, } type exportOptions struct { tableName string force bool dest mvdata.DataLocation srcOptions interface{} } func (m exportOptions) checkOverwrite(ctx context.Context, root doltdb.RootValue, fs filesys.ReadableFS) (bool, error) { if _, isStream := m.dest.(mvdata.StreamDataLocation); isStream { return false, nil } if !m.force { return m.dest.Exists(ctx, root, fs) } return false, nil } func (m exportOptions) IsBatched() bool { return false } func (m exportOptions) IsAutocommitOff() bool { return false } func (m exportOptions) WritesToTable() bool { return false } func (m exportOptions) SrcName() string { return m.tableName } func (m exportOptions) DestName() string { if f, fileDest := m.dest.(mvdata.FileDataLocation); fileDest { return f.Path } return m.dest.String() } // getExportDestination returns an export destination corresponding to the input parameters func getExportDestination(apr *argparser.ArgParseResults) mvdata.DataLocation { path := "" if apr.NArg() > 1 { path = apr.Arg(1) } fType, _ := apr.GetValue(fileTypeParam) destLoc := mvdata.NewDataLocation(path, fType) switch val := destLoc.(type) { case mvdata.FileDataLocation: if val.Format == mvdata.InvalidDataFormat { cli.PrintErrln( color.RedString("Could not infer type file '%s'\n", path), "File extensions should match supported file types, or should be explicitly defined via the file-type parameter") return nil } case mvdata.StreamDataLocation: if val.Format == mvdata.InvalidDataFormat { val = mvdata.StreamDataLocation{Format: mvdata.CsvFile, Reader: os.Stdin, Writer: iohelp.NopWrCloser(cli.CliOut)} destLoc = val } else if val.Format != mvdata.CsvFile && val.Format != mvdata.PsvFile { cli.PrintErrln(color.RedString("Cannot export this format to stdout")) return nil } } return destLoc } func parseExportArgs(ap *argparser.ArgParser, commandStr string, args []string) (*exportOptions, errhand.VerboseError) { help, usage := cli.HelpAndUsagePrinters(cli.CommandDocsForCommandString(commandStr, exportDocs, ap)) apr := cli.ParseArgsOrDie(ap, args, help) if apr.NArg() == 0 { usage() return nil, errhand.BuildDError("missing required argument").Build() } else if apr.NArg() > 2 { usage() return nil, errhand.BuildDError("too many arguments").Build() } tableName := apr.Arg(0) if !doltdb.IsValidTableName(tableName) { usage() cli.PrintErrln(color.RedString("'%s' is not a valid table name", tableName)) return nil, errhand.BuildDError("invalid table name").Build() } fileLoc := getExportDestination(apr) if fileLoc == nil { return nil, errhand.BuildDError("could not validate table export args").Build() } return &exportOptions{ tableName: tableName, force: apr.Contains(forceParam), dest: fileLoc, }, nil } type ExportCmd struct{} // Name is returns the name of the Dolt cli command. This is what is used on the command line to invoke the command func (cmd ExportCmd) Name() string { return "export" } // Description returns a description of the command func (cmd ExportCmd) Description() string { return "Export a table to a file." } func (cmd ExportCmd) Docs() *cli.CommandDocumentation { ap := cmd.ArgParser() return cli.NewCommandDocumentation(exportDocs, ap) } func (cmd ExportCmd) ArgParser() *argparser.ArgParser { ap := argparser.NewArgParserWithMaxArgs(cmd.Name(), 2) ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"table", "The table being exported."}) ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"file", "The file being output to."}) ap.SupportsFlag(forceParam, "f", "If data already exists in the destination, the force flag will allow the target to be overwritten.") ap.SupportsString(fileTypeParam, "", "file_type", "Explicitly define the type of the file if it can't be inferred from the file extension.") return ap } // EventType returns the type of the event to log func (cmd ExportCmd) EventType() eventsapi.ClientEventType { return eventsapi.ClientEventType_TABLE_EXPORT } // Exec executes the command func (cmd ExportCmd) Exec(ctx context.Context, commandStr string, args []string, dEnv *env.DoltEnv, cliCtx cli.CliContext) int { ap := cmd.ArgParser() _, usage := cli.HelpAndUsagePrinters(cli.CommandDocsForCommandString(commandStr, exportDocs, ap)) exOpts, verr := parseExportArgs(ap, commandStr, args) if verr != nil { return commands.HandleVErrAndExitCode(verr, usage) } root, verr := commands.GetWorkingWithVErr(dEnv) if verr != nil { return commands.HandleVErrAndExitCode(verr, usage) } engine, dbName, berr := engine.NewSqlEngineForEnv(ctx, dEnv) if berr != nil { return commands.HandleVErrAndExitCode(errhand.VerboseErrorFromError(berr), usage) } defer engine.Close() sqlCtx, berr := engine.NewLocalContext(ctx) if berr != nil { return commands.HandleVErrAndExitCode(errhand.VerboseErrorFromError(berr), usage) } defer sql.SessionEnd(sqlCtx.Session) sql.SessionCommandBegin(sqlCtx.Session) defer sql.SessionCommandEnd(sqlCtx.Session) sqlCtx.SetCurrentDatabase(dbName) rd, err := mvdata.NewSqlEngineReader(sqlCtx, engine.GetUnderlyingEngine(), root, exOpts.tableName) if err != nil { return commands.HandleVErrAndExitCode(errhand.BuildDError("Error creating reader for %s.", exOpts.SrcName()).AddCause(err).Build(), usage) } wr, verr := getTableWriter(sqlCtx, root, dEnv, rd.GetSchema(), exOpts) if verr != nil { return commands.HandleVErrAndExitCode(verr, usage) } pipeline := mvdata.NewDataMoverPipeline(sqlCtx, rd, wr) err = pipeline.Execute() if err != nil { return commands.HandleVErrAndExitCode(errhand.BuildDError("Error opening writer for %s.", exOpts.DestName()).AddCause(err).Build(), usage) } cli.PrintErrln(color.CyanString("Successfully exported data.")) return 0 } func getTableWriter(ctx context.Context, root doltdb.RootValue, dEnv *env.DoltEnv, rdSchema schema.Schema, exOpts *exportOptions) (table.SqlRowWriter, errhand.VerboseError) { ow, err := exOpts.checkOverwrite(ctx, root, dEnv.FS) if err != nil { return nil, errhand.VerboseErrorFromError(err) } if ow { return nil, errhand.BuildDError("%s already exists. Use -f to overwrite.", exOpts.DestName()).Build() } err = dEnv.FS.MkDirs(filepath.Dir(exOpts.DestName())) if err != nil { return nil, errhand.VerboseErrorFromError(err) } filePath, err := dEnv.FS.Abs(exOpts.DestName()) if err != nil { return nil, errhand.VerboseErrorFromError(err) } writer, err := dEnv.FS.OpenForWrite(filePath, os.ModePerm) if err != nil { return nil, errhand.BuildDError("Error opening writer for %s.", exOpts.DestName()).AddCause(err).Build() } wr, err := exOpts.dest.NewCreatingWriter(ctx, exOpts, root, rdSchema, editor.Options{}, writer) if err != nil { return nil, errhand.BuildDError("Error opening writer for %s.", exOpts.DestName()).AddCause(err).Build() } return wr, nil }