diff --git a/client.go b/client.go index 799b6bf4d..01a0d5bd9 100644 --- a/client.go +++ b/client.go @@ -339,6 +339,22 @@ func (c *Client) Import(ctx context.Context, index, frame string, slice uint64, return nil } +func (c *Client) EnsureIndex(ctx context.Context, name string, options IndexOptions) error { + err := c.CreateIndex(ctx, name, options) + if err == nil || err == ErrIndexExists { + return nil + } + return err +} + +func (c *Client) EnsureFrame(ctx context.Context, indexName string, frameName string, options FrameOptions) error { + err := c.CreateFrame(ctx, indexName, frameName, options) + if err == nil || err == ErrFrameExists { + return nil + } + return err +} + // MarshalImportPayload marshalls the import parameters into a protobuf byte slice. func MarshalImportPayload(index, frame string, slice uint64, bits []Bit) ([]byte, error) { // Separate row and column IDs to reduce allocations. diff --git a/cmd/import.go b/cmd/import.go index 944b8b9e6..a4276efa3 100644 --- a/cmd/import.go +++ b/cmd/import.go @@ -20,6 +20,7 @@ import ( "github.com/spf13/cobra" + "github.com/pilosa/pilosa" "github.com/pilosa/pilosa/ctl" ) @@ -49,12 +50,20 @@ omitted. If it is present then its format should be YYYY-MM-DDTHH:MM. return nil }, } + flags := importCmd.Flags() flags.StringVarP(&Importer.Host, "host", "", "localhost:10101", "host:port of Pilosa.") flags.StringVarP(&Importer.Index, "index", "i", "", "Pilosa index to import into.") flags.StringVarP(&Importer.Frame, "frame", "f", "", "Frame to import into.") flags.IntVarP(&Importer.BufferSize, "buffer-size", "s", 10000000, "Number of bits to buffer/sort before importing.") flags.BoolVarP(&Importer.Sort, "sort", "", false, "Enables sorting before import.") + flags.BoolVarP(&Importer.CreateSchema, "create", "e", false, "Create the schema if it does not exist before import.") + flags.Var(&Importer.IndexOptions.TimeQuantum, "index-time-quantum", "Time quantum for the index") + flags.Var(&Importer.FrameOptions.TimeQuantum, "frame-time-quantum", "Time quantum for the frame") + flags.BoolVar(&Importer.FrameOptions.InverseEnabled, "frame-inverse-enabled", false, "Enable inverse frame") + flags.BoolVar(&Importer.FrameOptions.RangeEnabled, "frame-range-enabled", false, "Enabled range encoded frame") + flags.StringVar(&Importer.FrameOptions.CacheType, "frame-cache-type", pilosa.CacheTypeRanked, "Cache type for the frame; valid values: none, lru, ranked") + flags.Uint32Var(&Importer.FrameOptions.CacheSize, "frame-cache-size", 50000, "Cache size for the frame") return importCmd } diff --git a/ctl/import.go b/ctl/import.go index a12cb0d06..68fbff691 100644 --- a/ctl/import.go +++ b/ctl/import.go @@ -38,6 +38,13 @@ type ImportCommand struct { Index string `json:"index"` Frame string `json:"frame"` + // Options for index & frame to be created if they don't exist + IndexOptions pilosa.IndexOptions + FrameOptions pilosa.FrameOptions + + // CreateSchema ensures the schema exists before import + CreateSchema bool + // Filenames to import from. Paths []string `json:"paths"` @@ -57,8 +64,7 @@ type ImportCommand struct { // NewImportCommand returns a new instance of ImportCommand. func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand { return &ImportCommand{ - CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr), - + CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr), BufferSize: 10000000, } } @@ -83,6 +89,13 @@ func (cmd *ImportCommand) Run(ctx context.Context) error { } cmd.Client = client + if cmd.CreateSchema { + err := cmd.ensureSchema(ctx) + if err != nil { + return err + } + } + // Import each path and import by slice. for _, path := range cmd.Paths { // Parse path into bits. @@ -95,6 +108,18 @@ func (cmd *ImportCommand) Run(ctx context.Context) error { return nil } +func (cmd *ImportCommand) ensureSchema(ctx context.Context) error { + err := cmd.Client.EnsureIndex(ctx, cmd.Index, cmd.IndexOptions) + if err != nil { + return fmt.Errorf("Error Creating Index: %s", err) + } + err = cmd.Client.EnsureFrame(ctx, cmd.Index, cmd.Frame, cmd.FrameOptions) + if err != nil { + return fmt.Errorf("Error Creating Frame: %s", err) + } + return nil +} + // importPath parses a path into bits and imports it to the server. func (cmd *ImportCommand) importPath(ctx context.Context, path string) error { a := make([]pilosa.Bit, 0, cmd.BufferSize) diff --git a/ctl/import_test.go b/ctl/import_test.go index dd9eba675..0023ad79a 100644 --- a/ctl/import_test.go +++ b/ctl/import_test.go @@ -71,11 +71,9 @@ func TestImportCommand_Run(t *testing.T) { s.Handler.Holder = hldr.Holder cm.Host = s.Host() - http.DefaultClient.Do(MustNewHTTPRequest("POST", s.URL+"/index/i", strings.NewReader(""))) - http.DefaultClient.Do(MustNewHTTPRequest("POST", s.URL+"/index/i/frame/f", strings.NewReader(""))) - cm.Index = "i" cm.Frame = "f" + cm.CreateSchema = true cm.Paths = []string{file.Name()} err = cm.Run(ctx) if err != nil { diff --git a/time.go b/time.go index ed6f62bf0..517292071 100644 --- a/time.go +++ b/time.go @@ -53,6 +53,23 @@ func (q TimeQuantum) Valid() bool { } } +// The following methods are required to implement pflag Value interface. + +// Set sets the time quantum value. +func (q *TimeQuantum) Set(value string) error { + *q = TimeQuantum(value) + return nil +} + +func (q TimeQuantum) String() string { + return string(q) +} + +// Type returns the type of a time quantum value. +func (q TimeQuantum) Type() string { + return "TimeQuantum" +} + // ParseTimeQuantum parses v into a time quantum. func ParseTimeQuantum(v string) (TimeQuantum, error) { q := TimeQuantum(strings.ToUpper(v))