github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/cmd/hkserver/commands/import.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package commands
     5  
     6  import (
     7  	"errors"
     8  	"fmt"
     9  	"os"
    10  
    11  	"github.com/spf13/cobra"
    12  
    13  	"github.com/masterhung0112/hk_server/v5/app/request"
    14  	"github.com/masterhung0112/hk_server/v5/audit"
    15  )
    16  
    17  var ImportCmd = &cobra.Command{
    18  	Use:   "import",
    19  	Short: "Import data.",
    20  }
    21  
    22  var SlackImportCmd = &cobra.Command{
    23  	Use:     "slack [team] [file]",
    24  	Short:   "Import a team from Slack.",
    25  	Long:    "Import a team from a Slack export zip file.",
    26  	Example: "  import slack myteam slack_export.zip",
    27  	RunE:    slackImportCmdF,
    28  }
    29  
    30  var BulkImportCmd = &cobra.Command{
    31  	Use:     "bulk [file]",
    32  	Short:   "Import bulk data.",
    33  	Long:    "Import data from a Mattermost Bulk Import File.",
    34  	Example: "  import bulk bulk_data.json",
    35  	RunE:    bulkImportCmdF,
    36  }
    37  
    38  func init() {
    39  	BulkImportCmd.Flags().Bool("apply", false, "Save the import data to the database. Use with caution - this cannot be reverted.")
    40  	BulkImportCmd.Flags().Bool("validate", false, "Validate the import data without making any changes to the system.")
    41  	BulkImportCmd.Flags().Int("workers", 2, "How many workers to run whilst doing the import.")
    42  	BulkImportCmd.Flags().String("import-path", "", "A path to the data directory to import files from.")
    43  
    44  	ImportCmd.AddCommand(
    45  		BulkImportCmd,
    46  		SlackImportCmd,
    47  	)
    48  	RootCmd.AddCommand(ImportCmd)
    49  }
    50  
    51  func slackImportCmdF(command *cobra.Command, args []string) error {
    52  	a, err := InitDBCommandContextCobra(command)
    53  	if err != nil {
    54  		return err
    55  	}
    56  	defer a.Srv().Shutdown()
    57  
    58  	if len(args) != 2 {
    59  		return errors.New("Incorrect number of arguments.")
    60  	}
    61  
    62  	team := getTeamFromTeamArg(a, args[0])
    63  	if team == nil {
    64  		return errors.New("Unable to find team '" + args[0] + "'")
    65  	}
    66  
    67  	fileReader, err := os.Open(args[1])
    68  	if err != nil {
    69  		return err
    70  	}
    71  	defer fileReader.Close()
    72  
    73  	fileInfo, err := fileReader.Stat()
    74  	if err != nil {
    75  		return err
    76  	}
    77  
    78  	CommandPrettyPrintln("Running Slack Import. This may take a long time for large teams or teams with many messages.")
    79  
    80  	importErr, log := a.SlackImport(&request.Context{}, fileReader, fileInfo.Size(), team.Id)
    81  
    82  	if importErr != nil {
    83  		return err
    84  	}
    85  
    86  	CommandPrettyPrintln("")
    87  	CommandPrintln(log.String())
    88  	CommandPrettyPrintln("")
    89  
    90  	CommandPrettyPrintln("Finished Slack Import.")
    91  	CommandPrettyPrintln("")
    92  
    93  	auditRec := a.MakeAuditRecord("slackImport", audit.Success)
    94  	auditRec.AddMeta("team", team)
    95  	auditRec.AddMeta("file", args[1])
    96  	a.LogAuditRec(auditRec, nil)
    97  
    98  	return nil
    99  }
   100  
   101  func bulkImportCmdF(command *cobra.Command, args []string) error {
   102  	a, err := InitDBCommandContextCobra(command)
   103  	if err != nil {
   104  		return err
   105  	}
   106  	defer a.Srv().Shutdown()
   107  
   108  	apply, err := command.Flags().GetBool("apply")
   109  	if err != nil {
   110  		return errors.New("Apply flag error")
   111  	}
   112  
   113  	validate, err := command.Flags().GetBool("validate")
   114  	if err != nil {
   115  		return errors.New("Validate flag error")
   116  	}
   117  
   118  	workers, err := command.Flags().GetInt("workers")
   119  	if err != nil {
   120  		return errors.New("Workers flag error")
   121  	}
   122  
   123  	importPath, err := command.Flags().GetString("import-path")
   124  	if err != nil {
   125  		return errors.New("import-path flag error")
   126  	}
   127  
   128  	if len(args) != 1 {
   129  		return errors.New("Incorrect number of arguments.")
   130  	}
   131  
   132  	fileReader, err := os.Open(args[0])
   133  	if err != nil {
   134  		return err
   135  	}
   136  	defer fileReader.Close()
   137  
   138  	if apply && validate {
   139  		CommandPrettyPrintln("Use only one of --apply or --validate.")
   140  		return nil
   141  	}
   142  
   143  	if apply && !validate {
   144  		CommandPrettyPrintln("Running Bulk Import. This may take a long time.")
   145  	} else {
   146  		CommandPrettyPrintln("Running Bulk Import Data Validation.")
   147  		CommandPrettyPrintln("** This checks the validity of the entities in the data file, but does not persist any changes **")
   148  		CommandPrettyPrintln("Use the --apply flag to perform the actual data import.")
   149  	}
   150  
   151  	CommandPrettyPrintln("")
   152  
   153  	if err, lineNumber := a.BulkImportWithPath(&request.Context{}, fileReader, nil, !apply, workers, importPath); err != nil {
   154  		CommandPrintErrorln(err.Error())
   155  		if lineNumber != 0 {
   156  			CommandPrintErrorln(fmt.Sprintf("Error occurred on data file line %v", lineNumber))
   157  		}
   158  		return err
   159  	}
   160  
   161  	if apply {
   162  		CommandPrettyPrintln("Finished Bulk Import.")
   163  		auditRec := a.MakeAuditRecord("bulkImport", audit.Success)
   164  		auditRec.AddMeta("file", args[0])
   165  		a.LogAuditRec(auditRec, nil)
   166  	} else {
   167  		CommandPrettyPrintln("Validation complete. You can now perform the import by rerunning this command with the --apply flag.")
   168  	}
   169  
   170  	return nil
   171  }