@@ -25,6 +25,7 @@ type ImportOptions struct {
2525	Index               string 
2626	ForwardToReplicas   bool 
2727	ClearExistingRules  bool 
28+ 	Wait                bool 
2829	Scanner             * bufio.Scanner 
2930
3031	DoConfirm  bool 
@@ -102,6 +103,7 @@ func NewImportCmd(f *cmdutil.Factory, runF func(*ImportOptions) error) *cobra.Co
102103		BoolVarP (& opts .ForwardToReplicas , "forward-to-replicas" , "f" , true , "Forward the rules to the index replicas" )
103104	cmd .Flags ().
104105		BoolVarP (& opts .ClearExistingRules , "clear-existing-rules" , "c" , false , "Clear existing rules before importing new ones" )
106+ 	cmd .Flags ().BoolVarP (& opts .Wait , "wait" , "w" , false , "wait for the operation to complete" )
105107
106108	return  cmd 
107109}
@@ -146,8 +148,8 @@ func runImportCmd(opts *ImportOptions) error {
146148
147149		var  rule  search.Rule 
148150		if  err  :=  json .Unmarshal ([]byte (line ), & rule ); err  !=  nil  {
149- 			err   :=   fmt . Errorf ( "failed to parse JSON rule on line %d: %s" ,  count ,  err )
150- 			return  err 
151+ 			opts . IO . StopProgressIndicator ( )
152+ 			return  fmt . Errorf ( "failed to parse JSON rule on line %d: %s" ,  count ,  err ) 
151153		}
152154
153155		rules  =  append (rules , rule )
@@ -156,33 +158,66 @@ func runImportCmd(opts *ImportOptions) error {
156158		// If requested, only clear existing rules the first time 
157159		clearExistingRules  :=  opts .ClearExistingRules 
158160		if  count  ==  batchSize  {
159- 			_ , err  :=  client .SaveRules (
161+ 			res , err  :=  client .SaveRules (
160162				client .NewApiSaveRulesRequest (opts .Index , rules ).
161163					WithClearExistingRules (clearExistingRules ).
162164					WithForwardToReplicas (opts .ForwardToReplicas ),
163165			)
164166			if  err  !=  nil  {
167+ 				opts .IO .StopProgressIndicator ()
165168				return  err 
166169			}
167- 			rules  =  make ([]search.Rule , 0 , batchSize )
170+ 			if  opts .Wait  {
171+ 				_ , err  :=  client .WaitForTask (opts .Index , res .TaskID )
172+ 				if  err  !=  nil  {
173+ 					opts .IO .StopProgressIndicator ()
174+ 					return  err 
175+ 				}
176+ 			}
168177			totalCount  +=  count 
169178			opts .IO .UpdateProgressIndicatorLabel (fmt .Sprintf ("Imported %d rules" , totalCount ))
179+ 
180+ 			rules  =  make ([]search.Rule , 0 , batchSize )
170181			count  =  0 
171182			clearExistingRules  =  false 
172183		}
173184	}
174185
175186	if  count  >  0  {
176187		totalCount  +=  count 
177- 		if  _ , err  :=  client .SaveRules (client .NewApiSaveRulesRequest (opts .Index , rules ).WithForwardToReplicas (opts .ForwardToReplicas )); err  !=  nil  {
188+ 		res , err  :=  client .SaveRules (
189+ 			client .NewApiSaveRulesRequest (opts .Index , rules ).
190+ 				WithForwardToReplicas (opts .ForwardToReplicas ),
191+ 		)
192+ 		if  err  !=  nil  {
193+ 			opts .IO .StopProgressIndicator ()
178194			return  err 
179195		}
196+ 		if  opts .Wait  {
197+ 			_ , err  :=  client .WaitForTask (opts .Index , res .TaskID )
198+ 			if  err  !=  nil  {
199+ 				opts .IO .StopProgressIndicator ()
200+ 				return  err 
201+ 			}
202+ 		}
180203	}
181204	// Clear rules if 0 rules are imported and the clear existing is set 
182205	if  totalCount  ==  0  &&  opts .ClearExistingRules  {
183- 		if  _ , err  :=  client .ClearRules (client .NewApiClearRulesRequest (opts .Index ).WithForwardToReplicas (opts .ForwardToReplicas )); err  !=  nil  {
206+ 		res , err  :=  client .ClearRules (
207+ 			client .NewApiClearRulesRequest (opts .Index ).
208+ 				WithForwardToReplicas (opts .ForwardToReplicas ),
209+ 		)
210+ 		if  err  !=  nil  {
211+ 			opts .IO .StopProgressIndicator ()
184212			return  err 
185213		}
214+ 		if  opts .Wait  {
215+ 			_ , err  :=  client .WaitForTask (opts .Index , res .TaskID )
216+ 			if  err  !=  nil  {
217+ 				opts .IO .StopProgressIndicator ()
218+ 				return  err 
219+ 			}
220+ 		}
186221	}
187222
188223	opts .IO .StopProgressIndicator ()
0 commit comments