github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/pkg/ha/task_cli_args.go (about)

     1  // Copyright 2019 PingCAP, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package ha
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  
    20  	"github.com/pingcap/tiflow/dm/common"
    21  	"github.com/pingcap/tiflow/dm/config"
    22  	"github.com/pingcap/tiflow/dm/pkg/etcdutil"
    23  	"github.com/pingcap/tiflow/dm/pkg/terror"
    24  	clientv3 "go.etcd.io/etcd/client/v3"
    25  )
    26  
    27  func putTaskCliArgsOp(taskname string, sources []string, args config.TaskCliArgs) ([]clientv3.Op, error) {
    28  	data, err := args.ToJSON()
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  
    33  	ops := []clientv3.Op{}
    34  	for _, source := range sources {
    35  		key := common.TaskCliArgsKeyAdapter.Encode(taskname, source)
    36  		ops = append(ops, clientv3.OpPut(key, data))
    37  	}
    38  	return ops, nil
    39  }
    40  
    41  // PutTaskCliArgs puts TaskCliArgs into etcd.
    42  func PutTaskCliArgs(cli *clientv3.Client, taskName string, sources []string, args config.TaskCliArgs) error {
    43  	ops, err := putTaskCliArgsOp(taskName, sources, args)
    44  	if err != nil {
    45  		return err
    46  	}
    47  	_, _, err = etcdutil.DoTxnWithRepeatable(cli, etcdutil.ThenOpFunc(ops...))
    48  	return err
    49  }
    50  
    51  // GetTaskCliArgs gets the command line arguments for the specified task.
    52  func GetTaskCliArgs(cli *clientv3.Client, taskName, source string) (*config.TaskCliArgs, error) {
    53  	ctx, cancel := context.WithTimeout(cli.Ctx(), etcdutil.DefaultRequestTimeout)
    54  	defer cancel()
    55  
    56  	resp, err := cli.Get(ctx, common.TaskCliArgsKeyAdapter.Encode(taskName, source))
    57  	if err != nil {
    58  		return nil, terror.ErrHAFailTxnOperation.Delegate(err, fmt.Sprintf("fail to get task cli args, taskName: %s, source: %s", taskName, source))
    59  	}
    60  
    61  	if resp.Count == 0 {
    62  		return nil, nil
    63  	} else if resp.Count > 1 {
    64  		// this should not happen.
    65  		return nil, terror.ErrConfigMoreThanOne.Generate(resp.Count, "TaskCliArgs", "task: "+taskName)
    66  	}
    67  
    68  	args := &config.TaskCliArgs{}
    69  	err = args.Decode(resp.Kvs[0].Value)
    70  	if err != nil {
    71  		return nil, err
    72  	}
    73  
    74  	return args, nil
    75  }
    76  
    77  // DeleteAllTaskCliArgs deleted the command line arguments of this task.
    78  func DeleteAllTaskCliArgs(cli *clientv3.Client, taskName string) error {
    79  	key := common.TaskCliArgsKeyAdapter.Encode(taskName)
    80  	op := clientv3.OpDelete(key, clientv3.WithPrefix())
    81  	_, _, err := etcdutil.DoTxnWithRepeatable(cli, etcdutil.ThenOpFunc(op))
    82  	return err
    83  }
    84  
    85  // DeleteTaskCliArgs deleted the command line arguments of this task.
    86  func DeleteTaskCliArgs(cli *clientv3.Client, taskName string, sources []string) error {
    87  	if len(sources) == 0 {
    88  		return nil
    89  	}
    90  	ops := []clientv3.Op{}
    91  	for _, source := range sources {
    92  		key := common.TaskCliArgsKeyAdapter.Encode(taskName, source)
    93  		ops = append(ops, clientv3.OpDelete(key))
    94  	}
    95  	_, _, err := etcdutil.DoTxnWithRepeatable(cli, etcdutil.ThenOpFunc(ops...))
    96  	return err
    97  }