go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/cv/internal/rpc/admin/dsmapper.go (about)

     1  // Copyright 2021 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package admin
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"strings"
    21  
    22  	"google.golang.org/grpc/codes"
    23  	"google.golang.org/protobuf/types/known/emptypb"
    24  
    25  	"go.chromium.org/luci/common/tsmon/field"
    26  	"go.chromium.org/luci/common/tsmon/metric"
    27  	"go.chromium.org/luci/grpc/appstatus"
    28  	"go.chromium.org/luci/server/dsmapper"
    29  
    30  	adminpb "go.chromium.org/luci/cv/internal/rpc/admin/api"
    31  )
    32  
    33  var (
    34  	metricUpgraded = metric.NewCounter(
    35  		"cv/internal/dsmapper/upgraded",
    36  		"Number of Datastore entities upgraded with dsmapper",
    37  		nil,
    38  		field.String("name"), // dsmapper.ID
    39  		field.Int("id"),      // dsmapper.JobID
    40  		field.String("kind"), // entity Kind
    41  	)
    42  )
    43  
    44  type dsMapper struct {
    45  	ctrl  *dsmapper.Controller
    46  	cfgs  map[dsmapper.ID]*dsmapper.JobConfig
    47  	final bool
    48  }
    49  
    50  func newDSMapper(ctrl *dsmapper.Controller) *dsMapper {
    51  	s := &dsMapper{ctrl: ctrl}
    52  	// Add your dsmapper factory here,
    53  	// e.g. s.register(&yourConfig, yourFactory)
    54  
    55  	// TODO(crbug/1260615): remove descriptions once CQDaemon is gone.
    56  	s.register(&removeCLDescriptionsCOnfig, removeCLDescriptionsFactory)
    57  	s.register(&multiCLAnalysisConfig, multiCLAnalysisMapperFactory)
    58  	s.register(&backfillRetentionKey, backfillRetentionKeyFactory)
    59  	s.register(&deleteEntitiesKey, deleteEntitiesFactory)
    60  
    61  	s.final = true
    62  	return s
    63  }
    64  
    65  func (d *dsMapper) register(cfg *dsmapper.JobConfig, f dsmapper.Factory) {
    66  	if d.final {
    67  		panic("register must be called in New()")
    68  	}
    69  	if err := cfg.Validate(); err != nil {
    70  		panic(err)
    71  	}
    72  	if f == nil {
    73  		panic("dsmapper.Factory must not be nil")
    74  	}
    75  	d.ctrl.RegisterFactory(cfg.Mapper, f) // panics if id not unique
    76  	if d.cfgs == nil {
    77  		d.cfgs = make(map[dsmapper.ID]*dsmapper.JobConfig, 1)
    78  	}
    79  	d.cfgs[cfg.Mapper] = cfg
    80  }
    81  
    82  func (a *AdminServer) DSMLaunchJob(ctx context.Context, req *adminpb.DSMLaunchJobRequest) (resp *adminpb.DSMJobID, err error) {
    83  	defer func() { err = appstatus.GRPCifyAndLog(ctx, err) }()
    84  	if err = checkAllowed(ctx, "DSMLaunchJob"); err != nil {
    85  		return
    86  	}
    87  
    88  	cfg, exists := a.dsmapper.cfgs[dsmapper.ID(req.GetName())]
    89  	if !exists {
    90  		var sb strings.Builder
    91  		_, _ = fmt.Fprintf(&sb, "Name %q is not pregistered; %d known names are: ", req.GetName(), len(a.dsmapper.cfgs))
    92  		for name := range a.dsmapper.cfgs {
    93  			sb.WriteString(string(name))
    94  			sb.WriteRune(',')
    95  		}
    96  		return nil, appstatus.Error(codes.NotFound, strings.TrimRight(sb.String(), ","))
    97  	}
    98  
    99  	jobID, err := a.dsmapper.ctrl.LaunchJob(ctx, cfg)
   100  	if err != nil {
   101  		return nil, err
   102  	}
   103  	return &adminpb.DSMJobID{Id: int64(jobID)}, nil
   104  }
   105  
   106  func (a *AdminServer) DSMGetJob(ctx context.Context, req *adminpb.DSMJobID) (_ *adminpb.DSMJob, err error) {
   107  	defer func() { err = appstatus.GRPCifyAndLog(ctx, err) }()
   108  	if err = checkAllowed(ctx, "DSMGetJob"); err != nil {
   109  		return
   110  	}
   111  
   112  	id := dsmapper.JobID(req.GetId())
   113  	job, err := a.dsmapper.ctrl.GetJob(ctx, id)
   114  	switch {
   115  	case err == dsmapper.ErrNoSuchJob:
   116  		return nil, appstatus.Error(codes.NotFound, "not found")
   117  	case err != nil:
   118  		return nil, err
   119  	}
   120  
   121  	info, err := job.FetchInfo(ctx)
   122  	if err != nil {
   123  		return nil, err
   124  	}
   125  
   126  	return &adminpb.DSMJob{
   127  		Name: string(job.Config.Mapper),
   128  		Info: info,
   129  	}, nil
   130  }
   131  
   132  func (a *AdminServer) DSMAbortJob(ctx context.Context, req *adminpb.DSMJobID) (_ *emptypb.Empty, err error) {
   133  	defer func() { err = appstatus.GRPCifyAndLog(ctx, err) }()
   134  	if err = checkAllowed(ctx, "DSMAbortJob"); err != nil {
   135  		return
   136  	}
   137  
   138  	id := dsmapper.JobID(req.GetId())
   139  	switch _, err := a.dsmapper.ctrl.AbortJob(ctx, id); {
   140  	case err == dsmapper.ErrNoSuchJob:
   141  		return nil, appstatus.Error(codes.NotFound, "not found")
   142  	case err != nil:
   143  		return nil, err
   144  	default:
   145  		return &emptypb.Empty{}, nil
   146  	}
   147  }