enhance: add CommitImport/AbortImport SDK wrappers for import 2PC (#50177)

## Summary Adds Go SDK RESTful wrappers for the Import 2PC commit/abort endpoints introduced server-side in #48524, mirroring the existing `GetImportProgress` wrapper in `client/bulkwriter/bulk_import.go`: - `CommitImport` → `POST /v2/vectordb/jobs/import/commit` - `AbortImport` → `POST /v2/vectordb/jobs/import/abort` Both take a `jobId` (plus optional `clusterId`/`apiKey`) and delegate status handling to the existing `ResponseBase.CheckStatus` helper. The 2PC opt-in (`auto_commit=false`) continues to flow through the existing import-create `options` map, so no create-path change is needed. issue: #48525 ### Changes **Modified:** - `client/bulkwriter/bulk_import.go`: add `CommitImportOption`/`AbortImportOption` (+ constructors, `WithAPIKey`, `GetRequest`), `CommitImportResponse`/`AbortImportResponse`, and the `CommitImport`/`AbortImport` functions. - `client/bulkwriter/bulk_import_test.go`: add `TestCommitImport`/`TestAbortImport` (normal / status-error / server-closed cases). ### Test Plan - [x] `go test ./bulkwriter/` — `TestBulkImportAPIs` suite passes (incl. new `TestCommitImport`/`TestAbortImport`). - [x] `gofumpt` clean; imports unchanged. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: bigsheeper <yihao.dai@zilliz.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 02:05:41 +00:00 · 2026-06-08 11:28:19 +08:00
parent c1815947bf
commit 8b04c2f960
2 changed files with 192 additions and 0 deletions
@@ -307,6 +307,128 @@ func GetImportProgress(ctx context.Context, option *GetImportProgressOption) (*G
 	return result, result.CheckStatus()
 }

+type CommitImportOption struct {
+	URL       string `json:"-"`
+	JobID     string `json:"jobId"`
+	ClusterID string `json:"clusterId,omitempty"`
+	APIKey    string `json:"-"`
+}
+
+func (opt *CommitImportOption) GetRequest() ([]byte, error) {
+	return json.Marshal(opt)
+}
+
+func (opt *CommitImportOption) WithAPIKey(key string) *CommitImportOption {
+	opt.APIKey = key
+	return opt
+}
+
+func NewCommitImportOption(uri string, jobID string) *CommitImportOption {
+	return &CommitImportOption{
+		URL:   uri,
+		JobID: jobID,
+	}
+}
+
+func NewCloudCommitImportOption(uri string, jobID string, apiKey string, clusterID string) *CommitImportOption {
+	return &CommitImportOption{
+		URL:       uri,
+		JobID:     jobID,
+		APIKey:    apiKey,
+		ClusterID: clusterID,
+	}
+}
+
+type CommitImportResponse struct {
+	ResponseBase
+}
+
+// CommitImport is the API wrapper for the restful import commit API.
+func CommitImport(ctx context.Context, option *CommitImportOption) (*CommitImportResponse, error) {
+	url := option.URL + "/v2/vectordb/jobs/import/commit"
+
+	bs, err := option.GetRequest()
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bs))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	if option.APIKey != "" {
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", option.APIKey))
+	}
+
+	result := &CommitImportResponse{}
+	if err := doPostRequest(req, result); err != nil {
+		return nil, err
+	}
+	return result, result.CheckStatus()
+}
+
+type AbortImportOption struct {
+	URL       string `json:"-"`
+	JobID     string `json:"jobId"`
+	ClusterID string `json:"clusterId,omitempty"`
+	APIKey    string `json:"-"`
+}
+
+func (opt *AbortImportOption) GetRequest() ([]byte, error) {
+	return json.Marshal(opt)
+}
+
+func (opt *AbortImportOption) WithAPIKey(key string) *AbortImportOption {
+	opt.APIKey = key
+	return opt
+}
+
+func NewAbortImportOption(uri string, jobID string) *AbortImportOption {
+	return &AbortImportOption{
+		URL:   uri,
+		JobID: jobID,
+	}
+}
+
+func NewCloudAbortImportOption(uri string, jobID string, apiKey string, clusterID string) *AbortImportOption {
+	return &AbortImportOption{
+		URL:       uri,
+		JobID:     jobID,
+		APIKey:    apiKey,
+		ClusterID: clusterID,
+	}
+}
+
+type AbortImportResponse struct {
+	ResponseBase
+}
+
+// AbortImport is the API wrapper for the restful import abort API.
+func AbortImport(ctx context.Context, option *AbortImportOption) (*AbortImportResponse, error) {
+	url := option.URL + "/v2/vectordb/jobs/import/abort"
+
+	bs, err := option.GetRequest()
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bs))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	if option.APIKey != "" {
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", option.APIKey))
+	}
+
+	result := &AbortImportResponse{}
+	if err := doPostRequest(req, result); err != nil {
+		return nil, err
+	}
+	return result, result.CheckStatus()
+}
+
 func doPostRequest(req *http.Request, response any) error {
 	client := &http.Client{}
 	resp, err := client.Do(req)
@@ -163,6 +163,76 @@ func (s *BulkImportSuite) TestGetImportProgress() {
 	})
 }

+func (s *BulkImportSuite) TestCommitImport() {
+	s.Run("normal_case", func() {
+		svr := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			authHeader := req.Header.Get("Authorization")
+			s.Equal("Bearer root:Milvus", authHeader)
+			s.True(strings.Contains(req.URL.Path, "/v2/vectordb/jobs/import/commit"))
+			rw.Write([]byte(`{"status":0, "data":{}}`))
+		}))
+		defer svr.Close()
+
+		resp, err := CommitImport(context.Background(),
+			NewCommitImportOption(svr.URL, "123").WithAPIKey("root:Milvus"))
+		s.NoError(err)
+		s.EqualValues(0, resp.Status)
+	})
+
+	s.Run("status_error", func() {
+		svr := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			s.True(strings.Contains(req.URL.Path, "/v2/vectordb/jobs/import/commit"))
+			rw.Write([]byte(`{"status":1100, "message": "commit failed"}`))
+		}))
+		defer svr.Close()
+
+		_, err := CommitImport(context.Background(), NewCommitImportOption(svr.URL, "123"))
+		s.Error(err)
+	})
+
+	s.Run("server_closed", func() {
+		svr := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}))
+		svr.Close()
+		_, err := CommitImport(context.Background(), NewCommitImportOption(svr.URL, "123"))
+		s.Error(err)
+	})
+}
+
+func (s *BulkImportSuite) TestAbortImport() {
+	s.Run("normal_case", func() {
+		svr := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			authHeader := req.Header.Get("Authorization")
+			s.Equal("Bearer root:Milvus", authHeader)
+			s.True(strings.Contains(req.URL.Path, "/v2/vectordb/jobs/import/abort"))
+			rw.Write([]byte(`{"status":0, "data":{}}`))
+		}))
+		defer svr.Close()
+
+		resp, err := AbortImport(context.Background(),
+			NewAbortImportOption(svr.URL, "123").WithAPIKey("root:Milvus"))
+		s.NoError(err)
+		s.EqualValues(0, resp.Status)
+	})
+
+	s.Run("status_error", func() {
+		svr := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			s.True(strings.Contains(req.URL.Path, "/v2/vectordb/jobs/import/abort"))
+			rw.Write([]byte(`{"status":1100, "message": "abort failed"}`))
+		}))
+		defer svr.Close()
+
+		_, err := AbortImport(context.Background(), NewAbortImportOption(svr.URL, "123"))
+		s.Error(err)
+	})
+
+	s.Run("server_closed", func() {
+		svr := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}))
+		svr.Close()
+		_, err := AbortImport(context.Background(), NewAbortImportOption(svr.URL, "123"))
+		s.Error(err)
+	})
+}
+
 func TestBulkImportAPIs(t *testing.T) {
 	suite.Run(t, new(BulkImportSuite))
 }