Commit 1dfd1f28 authored by Pawel Rozlach's avatar Pawel Rozlach 💬
Browse files

feat: add /v1/gitlab api endpoint for asynchronous restarting of BBMs

parent 67d0eeb2

registry/api/gitlab/v1/routes.go

+8 −0
Changes for registry/api/gitlab/v1/routes.go: 8 added lines, 0 removed lines.
Original line number Diff line number Diff line
@@ -82,6 +82,13 @@ var ( 
		Path: BBM.Path + "resume/",

		ID:   BBM.Path + "resume",

	}



	// BBMRestart is the API route for restarting a specific background migration.

	BBMRestart = Route{

		Name: "background-migrations-restart",

		Path: Base.Path + "restart/{bbmId:" + reference.NumericRegexp.String() + "}/",

		ID:   Base.Path + "restart/{bbmId:" + reference.NumericRegexp.String() + "}",

	}

)



// Router returns a new *mux.Router for the Gitlab v1 API.
@@ -112,6 +119,7 @@ func RouterWithPrefix(prefix string) *mux.Router { 
	router.Path(BBMById.Path).Name(BBMById.Name)

	router.Path(BBMPause.Path).Name(BBMPause.Name)

	router.Path(BBMResume.Path).Name(BBMResume.Name)

	router.Path(BBMRestart.Path).Name(BBMRestart.Name)



	return rootRouter

}

registry/handlers/app.go

+3 −2
Changes for registry/handlers/app.go: 3 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -1325,6 +1325,7 @@ func (app *App) initMetaRouter() error { 
	app.registerGitlab(v1.BBMById, h.wrap(backgroundMigrationDispatcher))

	app.registerGitlab(v1.BBMPause, h.wrap(backgroundMigrationsPauseDispatcher))

	app.registerGitlab(v1.BBMResume, h.wrap(backgroundMigrationsResumeDispatcher))

	app.registerGitlab(v1.BBMRestart, h.wrap(backgroundMigrationsRestartDispatcher))



	var err error

	v1PathWithPrefix := fmt.Sprintf("^%s%s.*", strings.TrimSuffix(app.Config.HTTP.Prefix, "/"), v1.Base.Path)
@@ -1793,7 +1794,7 @@ func (*App) nameRequired(r *http.Request) bool { 
	switch routeName {

	case v2.RouteNameBase, v2.RouteNameCatalog, v1.Base.Name, v1.Statistics.Name:

		return false

	case v1.BBM.Name, v1.BBMById.Name, v1.BBMPause.Name, v1.BBMResume.Name:

	case v1.BBM.Name, v1.BBMById.Name, v1.BBMPause.Name, v1.BBMResume.Name, v1.BBMRestart.Name:

		return false

	}
@@ -1910,7 +1911,7 @@ func appendBBMAccessRecord(accessRecords []auth.Access, r *http.Request) []auth. 
	routeName := route.GetName()



	switch routeName {

	case v1.BBM.Name, v1.BBMById.Name, v1.BBMPause.Name, v1.BBMResume.Name:

	case v1.BBM.Name, v1.BBMById.Name, v1.BBMPause.Name, v1.BBMResume.Name, v1.BBMRestart.Name:

		accessRecords = append(accessRecords,

			auth.Access{

				Resource: auth.Resource{

registry/handlers/background_migrations.go

+66 −1
Changes for registry/handlers/background_migrations.go: 66 added lines, 1 removed line.
Original line number Diff line number Diff line
@@ -38,7 +38,8 @@ type BackgroundMigrationGetResponse struct { 
	Migration BackgroundMigrationResponse `json:"migration"`

}



// BackgroundMigrationActionResponse is the response for action endpoints (pause/resume).

// BackgroundMigrationActionResponse is the response for action endpoints

// (pause/resume/restart).

type BackgroundMigrationActionResponse struct {

	Success bool   `json:"success"`

	Message string `json:"message,omitempty"`
@@ -104,6 +105,17 @@ func backgroundMigrationsResumeDispatcher(ctx *Context, _ *http.Request) http.Ha 
	}

}



// backgroundMigrationsRestartDispatcher routes requests to the appropriate handler for the restart endpoint.

func backgroundMigrationsRestartDispatcher(ctx *Context, _ *http.Request) http.Handler {

	handler := &backgroundMigrationsHandler{

		Context: ctx,

	}



	return handlers.MethodHandler{

		http.MethodPost: http.HandlerFunc(handler.RestartBackgroundMigration),

	}

}



// GetBackgroundMigrations handles GET requests to get the status of all

// background migrations.

func (h *backgroundMigrationsHandler) GetBackgroundMigrations(w http.ResponseWriter, r *http.Request) {
@@ -214,6 +226,59 @@ func (h *backgroundMigrationsHandler) ResumeBackgroundMigrations(w http.Response 
	}

}



// RestartBackgroundMigration handles POST requests to restart a specific background migration.

// It resets the failure_error_code to NULL and status to active (1).

func (h *backgroundMigrationsHandler) RestartBackgroundMigration(w http.ResponseWriter, r *http.Request) {

	w.Header().Set("Content-Type", "application/json")



	bbmId, err := getBBMId(r.Context())

	if err != nil {

		h.Errors = append(

			h.Errors,

			v2.ErrorCodeBBMIdInvalid.WithDetail(err),

		)

		return

	}



	// First, fetch the BBM to ensure it exists

	bbm, err := datastore.NewBackgroundMigrationStore(h.db.Primary()).

		FindById(r.Context(), bbmId)

	if err != nil {

		h.Errors = append(h.Errors, errcode.FromUnknownError(err))

		return

	}

	if bbm == nil {

		h.Errors = append(

			h.Errors,

			v2.ErrorCodeBBMNotFound.WithDetail(fmt.Sprintf("BBM with ID %d does not exist", bbmId)),

		)

		return

	}



	// Set status to active and clear error code

	bbm.Status = models.BackgroundMigrationActive

	bbm.ErrorCode = models.NullErrCode



	// Update the BBM status

	err = datastore.NewBackgroundMigrationStore(h.db.Primary()).

		UpdateStatus(r.Context(), bbm)

	if err != nil {

		h.Errors = append(h.Errors, errcode.FromUnknownError(err))

		return

	}



	resp := BackgroundMigrationActionResponse{

		Success: true,

		Message: fmt.Sprintf("Background migration %d has been restarted", bbmId),

	}



	enc := json.NewEncoder(w)

	if err := enc.Encode(resp); err != nil {

		h.Errors = append(h.Errors, errcode.FromUnknownError(err))

		return

	}

}



// convertToBackgroundMigrationResponse converts a single models.BackgroundMigration to API response format.

func convertToBackgroundMigrationResponse(m *models.BackgroundMigration) BackgroundMigrationResponse {

	if m == nil {

registry/handlers/background_migrations_test.go

+264 −0
Changes for registry/handlers/background_migrations_test.go: 264 added lines, 0 removed lines.
Original line number Diff line number Diff line
@@ -652,6 +652,270 @@ func (s *BackgroundMigrationsHandlerTestSuite) TestResumeBackgroundMigrations_Da 
	require.Equal(s.T(), http.StatusInternalServerError, resp.StatusCode)

}



func (s *BackgroundMigrationsHandlerTestSuite) TestRestartBackgroundMigration_Success() {

	// Prepare test data - a failed migration

	migration := &models.BackgroundMigration{

		ID:               1,

		Name:             "migration_1",

		Status:           models.BackgroundMigrationFailed,

		JobName:          "job_1",

		TargetTable:      "public.repositories",

		TargetColumn:     "id",

		BatchSize:        100,

		StartID:          1,

		EndID:            1000,

		BatchingStrategy: models.SerialKeySetBatchingBBMStrategy,

		TotalTupleCount:  sql.NullInt64{Int64: 500, Valid: true},

		ErrorCode:        models.InvalidTableBBMErrCode,

	}



	// Setup sqlmock rows for FindById

	findRows := sqlmock.NewRows([]string{

		"id",

		"name",

		"min_value",

		"max_value",

		"batch_size",

		"status",

		"job_signature_name",

		"table_name",

		"column_name",

		"failure_error_code",

		"batching_strategy",

		"total_tuple_count",

	})



	errVal, _ := migration.ErrorCode.Value()

	bStrategy, _ := migration.BatchingStrategy.Value()

	findRows.AddRow(

		migration.ID,

		migration.Name,

		migration.StartID,

		migration.EndID,

		migration.BatchSize,

		int(migration.Status),

		migration.JobName,

		migration.TargetTable,

		migration.TargetColumn,

		errVal,

		bStrategy,

		migration.TotalTupleCount,

	)



	// Setup sqlmock rows for UpdateStatus

	updateRows := sqlmock.NewRows([]string{

		"status",

		"failure_error_code",

	}).AddRow(int(models.BackgroundMigrationActive), nil)



	s.mockAccessCtrl.EXPECT().

		Authorized(gomock.Any(), auth.Access{

			Resource: auth.Resource{

				Type: "registry",

				Name: "background-migrations",

			},

			Action: "*",

		}).

		DoAndReturn(func(ctx any, _ ...auth.Access) (any, error) {

			return ctx, nil

		})



	// Expect the FindById query

	s.mockQuerier.ExpectQuery(`SELECT.*FROM.*batched_background_migrations.*WHERE.*id = \$1`).

		WithArgs(1).

		WillReturnRows(findRows)



	// Expect the UpdateStatus query

	s.mockQuerier.ExpectQuery(`UPDATE.*batched_background_migrations.*SET.*status.*failure_error_code.*WHERE.*id.*`).

		WithArgs(int(models.BackgroundMigrationActive), nil, 1, int(models.BackgroundMigrationRunning), int(models.BackgroundMigrationFinished)).

		WillReturnRows(updateRows)



	s.mockLB.EXPECT().Primary().Return(&datastore.DB{DB: s.mockPrimaryDB}).Times(2)



	// Make request

	resp, err := http.Post(fmt.Sprintf("%s/gitlab/v1/restart/1/", s.server.URL), "application/json", nil)

	require.NoError(s.T(), err)

	defer resp.Body.Close()



	// Verify response

	require.Equal(s.T(), http.StatusOK, resp.StatusCode)

	require.Equal(s.T(), "application/json", resp.Header.Get("Content-Type"))



	var result BackgroundMigrationActionResponse

	err = json.NewDecoder(resp.Body).Decode(&result)

	require.NoError(s.T(), err)



	// Verify response content

	assert.True(s.T(), result.Success)

	assert.Equal(s.T(), "Background migration 1 has been restarted", result.Message)

}



func (s *BackgroundMigrationsHandlerTestSuite) TestRestartBackgroundMigration_Forbidden() {

	mockChallenge := amocks.NewMockChallenge(s.ctrl)

	mockChallenge.EXPECT().SetHeaders(gomock.Any(), gomock.Any()).Times(1)



	s.mockAccessCtrl.EXPECT().

		Authorized(gomock.Any(), auth.Access{

			Resource: auth.Resource{

				Type: "registry",

				Name: "background-migrations",

			},

			Action: "*",

		}).

		Return(nil, mockChallenge)



	// Make request

	resp, err := http.Post(fmt.Sprintf("%s/gitlab/v1/restart/1/", s.server.URL), "application/json", nil)

	require.NoError(s.T(), err)

	defer resp.Body.Close()



	// Verify 403 Forbidden response

	require.Equal(s.T(), http.StatusUnauthorized, resp.StatusCode)

}



func (s *BackgroundMigrationsHandlerTestSuite) TestRestartBackgroundMigration_InvalidID() {

	// Expect authorization check

	s.mockAccessCtrl.EXPECT().

		Authorized(gomock.Any(), gomock.Any()).

		DoAndReturn(func(ctx any, _ ...auth.Access) (any, error) {

			return ctx, nil

		})



	// Make request with invalid ID

	resp, err := http.Post(fmt.Sprintf("%s/gitlab/v1/restart/99999999999999999999999999999999999999999999999/", s.server.URL), "application/json", nil)

	require.NoError(s.T(), err)

	defer resp.Body.Close()



	// Verify error response

	require.Equal(s.T(), http.StatusBadRequest, resp.StatusCode)

}



func (s *BackgroundMigrationsHandlerTestSuite) TestRestartBackgroundMigration_NotFound() {

	// Expect authorization check

	s.mockAccessCtrl.EXPECT().

		Authorized(gomock.Any(), gomock.Any()).

		DoAndReturn(func(ctx any, _ ...auth.Access) (any, error) {

			return ctx, nil

		})



	// Expect the query to return no rows

	s.mockQuerier.ExpectQuery(`SELECT.*FROM.*batched_background_migrations.*WHERE.*id.*`).

		WithArgs(999).

		WillReturnError(sql.ErrNoRows)



	s.mockLB.EXPECT().Primary().Return(&datastore.DB{DB: s.mockPrimaryDB}).Times(1)



	// Make request

	resp, err := http.Post(fmt.Sprintf("%s/gitlab/v1/restart/999/", s.server.URL), "application/json", nil)

	require.NoError(s.T(), err)

	defer resp.Body.Close()



	// Verify error response

	require.Equal(s.T(), http.StatusNotFound, resp.StatusCode)

}



func (s *BackgroundMigrationsHandlerTestSuite) TestRestartBackgroundMigration_DatabaseError() {

	// Expect authorization check

	s.mockAccessCtrl.EXPECT().

		Authorized(gomock.Any(), gomock.Any()).

		DoAndReturn(func(ctx any, _ ...auth.Access) (any, error) {

			return ctx, nil

		})



	// Expect the query to return an error

	s.mockQuerier.ExpectQuery(`SELECT.*FROM.*batched_background_migrations.*WHERE.*id.*`).

		WithArgs(1).

		WillReturnError(fmt.Errorf("database connection error"))



	s.mockLB.EXPECT().Primary().Return(&datastore.DB{DB: s.mockPrimaryDB}).Times(1)



	// Make request

	resp, err := http.Post(fmt.Sprintf("%s/gitlab/v1/restart/1/", s.server.URL), "application/json", nil)

	require.NoError(s.T(), err)

	defer resp.Body.Close()



	// Verify error response

	require.Equal(s.T(), http.StatusInternalServerError, resp.StatusCode)

}



func (s *BackgroundMigrationsHandlerTestSuite) TestRestartBackgroundMigration_UpdateError() {

	// Prepare test data

	migration := &models.BackgroundMigration{

		ID:               1,

		Name:             "migration_1",

		Status:           models.BackgroundMigrationFailed,

		JobName:          "job_1",

		TargetTable:      "public.repositories",

		TargetColumn:     "id",

		BatchSize:        100,

		StartID:          1,

		EndID:            1000,

		BatchingStrategy: models.SerialKeySetBatchingBBMStrategy,

		TotalTupleCount:  sql.NullInt64{Int64: 500, Valid: true},

		ErrorCode:        models.InvalidTableBBMErrCode,

	}



	// Setup sqlmock rows for FindById

	findRows := sqlmock.NewRows([]string{

		"id",

		"name",

		"min_value",

		"max_value",

		"batch_size",

		"status",

		"job_signature_name",

		"table_name",

		"column_name",

		"failure_error_code",

		"batching_strategy",

		"total_tuple_count",

	})



	errVal, _ := migration.ErrorCode.Value()

	bStrategy, _ := migration.BatchingStrategy.Value()

	findRows.AddRow(

		migration.ID,

		migration.Name,

		migration.StartID,

		migration.EndID,

		migration.BatchSize,

		int(migration.Status),

		migration.JobName,

		migration.TargetTable,

		migration.TargetColumn,

		errVal,

		bStrategy,

		migration.TotalTupleCount,

	)



	// Expect authorization check

	s.mockAccessCtrl.EXPECT().

		Authorized(gomock.Any(), gomock.Any()).

		DoAndReturn(func(ctx any, _ ...auth.Access) (any, error) {

			return ctx, nil

		})



	// Expect the FindById query

	s.mockQuerier.ExpectQuery(`SELECT.*FROM.*batched_background_migrations.*WHERE.*id = \$1`).

		WithArgs(1).

		WillReturnRows(findRows)



	// Expect the UpdateStatus query to fail

	s.mockQuerier.ExpectQuery(`UPDATE.*batched_background_migrations.*SET.*status.*failure_error_code.*WHERE.*id.*`).

		WithArgs(int(models.BackgroundMigrationActive), nil, 1, int(models.BackgroundMigrationRunning), int(models.BackgroundMigrationFinished)).

		WillReturnError(fmt.Errorf("database update error"))



	s.mockLB.EXPECT().Primary().Return(&datastore.DB{DB: s.mockPrimaryDB}).Times(2)



	// Make request

	resp, err := http.Post(fmt.Sprintf("%s/gitlab/v1/restart/1/", s.server.URL), "application/json", nil)

	require.NoError(s.T(), err)

	defer resp.Body.Close()



	// Verify error response

	require.Equal(s.T(), http.StatusInternalServerError, resp.StatusCode)

}



func TestBackgroundMigrationsHandlerTestSuite(t *testing.T) {

	suite.Run(t, new(BackgroundMigrationsHandlerTestSuite))

}