package grpc import ( "context" "crypto/tls" "encoding/json" goerr "errors" "fmt" "io" "net" "os" "strings" "time" "unicode/utf8" gcsstorage "cloud.google.com/go/storage" "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/credentials/stscreds" awss3 "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/aws/aws-sdk-go-v2/service/sts" "github.com/pingcap/errors" plog "github.com/pingcap/log" tdbrstorage "github.com/pingcap/tidb/br/pkg/storage" lnrestore "github.com/pingcap/tidb/lightning/pkg/importer" lnrestore_opts "github.com/pingcap/tidb/lightning/pkg/importer/opts" lnprecheck "github.com/pingcap/tidb/lightning/pkg/precheck" "github.com/pingcap/tidb/pkg/importsdk" lncommon "github.com/pingcap/tidb/pkg/lightning/common" lnconfig "github.com/pingcap/tidb/pkg/lightning/config" lnlog "github.com/pingcap/tidb/pkg/lightning/log" "github.com/pingcap/tidb/pkg/lightning/mydump" "github.com/pingcap/tidb/pkg/meta/metadef" "github.com/pingcap/tiflow/dm/checker" dmconfig "github.com/pingcap/tiflow/dm/config" "github.com/pingcap/tiflow/dm/config/dbconfig" "github.com/pingcap/tiflow/dm/config/security" "github.com/pingcap/tiflow/dm/pkg/conn" tcontext "github.com/pingcap/tiflow/dm/pkg/context" "github.com/pingcap/tiflow/dm/pkg/log" "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent" "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/common" "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow" commondm "github.com/tidbcloud/dataflow-service/internal/common/dm" "github.com/tidbcloud/dataflow-service/internal/crypto" "github.com/tidbcloud/dataflow-service/internal/db/entity" "github.com/tidbcloud/dataflow-service/internal/utils" "github.com/tidbcloud/dataflow-service/pkg/agent/constants" fileparser "github.com/tidbcloud/dataflow-service/pkg/agent/file-parser" "github.com/tidbcloud/dataflow-service/pkg/agent/lightning" "github.com/tidbcloud/dataflow-service/pkg/agent/matcher" "github.com/tidbcloud/dataflow-service/pkg/agent/tidbcloud" "github.com/tidbcloud/dataflow-service/pkg/azblob" "github.com/tidbcloud/dataflow-service/pkg/gcs" "github.com/tidbcloud/dataflow-service/pkg/s3" storagepkg "github.com/tidbcloud/dataflow-service/pkg/storage" "github.com/tidbcloud/pkgv2/import/cloud_storage" "github.com/tidbcloud/pkgv2/import/precheck" "github.com/tikv/pd/client/http" pdServiceDiscovery "github.com/tikv/pd/client/servicediscovery" "go.buf.build/protocolbuffers/go/linuxgit/tidbcloud/base" "go.uber.org/zap" "golang.org/x/sync/errgroup" "google.golang.org/api/option" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/emptypb" "google.golang.org/protobuf/types/known/timestamppb" ) func init() { conf := &plog.Config{Level: "info", File: plog.FileLogConfig{}} plogLogger, plogProps, err := plog.InitLogger(conf) if err == nil { plog.ReplaceGlobals(plogLogger, plogProps) } lnlog.SetAppLogger(plog.L()) lnlog.SetLevel(zap.InfoLevel) } var tooManySrcFilesErrMessage = fmt.Sprintf("too many source files, only scanned %d of them", lightning.MaxScanFilesCount) func (s *AgentServer) CheckS3Access(ctx context.Context, req *agent.CheckS3AccessReq) ( *agent.CheckS3AccessResp, error, ) { resp := &agent.CheckS3AccessResp{} if err := s3.CheckS3Access( ctx, req.S3Url, req.RoleArn, req.ExternalId, req.AwsAccessKeyId, req.AwsSecretAccessKey, ); err != nil { resp.ErrorMsg = err.Error() } return resp, nil } func (s *AgentServer) CheckGcsAccess(ctx context.Context, req *agent.CheckGcsAccessReq) ( *agent.CheckGcsAccessResp, error, ) { resp := &agent.CheckGcsAccessResp{} if err := gcs.CheckGcsAccess(ctx, req.GcsUrl, req.GetGcsOptions().GetImpersonateServiceAccount()); err != nil { resp.ErrorMsg = err.Error() } return resp, nil } func (s *AgentServer) CheckAzblobAccess(ctx context.Context, req *agent.CheckAzblobAccessReq) ( *agent.CheckAzblobAccessResp, error, ) { resp := &agent.CheckAzblobAccessResp{} if err := azblob.CheckAzblobAccess(ctx, req.AzblobUrl, req.SasToken); err != nil { resp.ErrorMsg = err.Error() } return resp, nil } func (s *AgentServer) getClusterPassword(ctx context.Context, cluster *agent.Cluster) (string, error) { if !cluster.GetPasswordEncrypted() { return cluster.GetPassword(), nil } pem := cluster.GetEncryptedPassword().GetPem() verified, err := s.encryptor.VerifyKeyMatch(pem) if err != nil { return "", errors.Trace(err) } if !verified { return "", crypto.PEMUpdatedError } data := cluster.GetEncryptedPassword().GetData() passwd, err := s.encryptor.Decrypt(data) if err != nil { return "", errors.Trace(err) } return string(passwd), nil } func (s *AgentServer) ExecSQLOnCluster(ctx context.Context, req *agent.ExecSQLReq) ( *agent.ExecSQLResp, error, ) { resp := &agent.ExecSQLResp{ Rows: make([]*agent.Row, 0), } opts := &tidbcloud.DBOptions{ ReadTimeout: req.Cluster.ReadTimeout, WriteTimeout: req.Cluster.WriteTimeout, } passwd, err := s.getClusterPassword(ctx, req.Cluster) if err != nil { if goerr.Is(err, crypto.PEMUpdatedError) { resp.ErrorMsg = crypto.PEMUpdatedErrorMsg return resp, nil } return nil, err } db, err := s.clusterProvider.GetDB(ctx, fmt.Sprint(req.Cluster.ClusterId), req.Cluster.User, passwd, opts) if err != nil { return nil, err } defer db.Close() // nolint: errcheck for _, sql := range req.SetupSqls { log.L().Info("execute setup SQL", zap.String("sql", sql)) _, err = db.Exec(sql) if err != nil { log.L().Error("exec setup SQL error", zap.String("sql", sql), zap.Error(err)) return nil, err } } // this represents need not return rows, so use exec if req.ColumnNum == 0 { _, err := db.ExecContext(ctx, req.GetSql()) if err != nil { log.L().Error("exec SQL error", zap.String("sql", req.GetSql()), zap.Error(err)) return nil, err } return resp, nil } rows, err := db.QueryContext(ctx, req.Sql) if err != nil { log.L().Error("query SQL error", zap.String("sql", req.GetSql()), zap.Error(err)) return nil, err } defer rows.Close() // nolint: errcheck resp.ColumnNames, err = rows.Columns() if err != nil { log.L().Error("get columns error", zap.String("sql", req.GetSql()), zap.Error(err)) return nil, err } if req.TypedColumns { for rows.Next() { row := &agent.Row{} values := make([]interface{}, len(resp.ColumnNames)) valuePtrs := make([]interface{}, len(resp.ColumnNames)) for i := range values { valuePtrs[i] = &values[i] } if err := rows.Scan(valuePtrs...); err != nil { log.L().Error("scan SQL error", zap.String("sql", req.GetSql()), zap.Error(err)) return nil, err } for _, value := range values { columnValue := &agent.ColumnValue{} switch v := value.(type) { case nil: columnValue.Value = nil case uint, uint8, uint16, uint32, uint64: columnValue.Value = &agent.ColumnValue_Int64Value{Int64Value: int64(v.(uint64))} case int, int8, int16, int32, int64: columnValue.Value = &agent.ColumnValue_Int64Value{Int64Value: int64(v.(int64))} case float32, float64: columnValue.Value = &agent.ColumnValue_DoubleValue{DoubleValue: float64(v.(float64))} case bool: columnValue.Value = &agent.ColumnValue_BoolValue{BoolValue: v} case string: columnValue.Value = &agent.ColumnValue_StringValue{StringValue: v} case []uint8: columnValue.Value = &agent.ColumnValue_StringValue{StringValue: string(v)} case time.Time: columnValue.Value = &agent.ColumnValue_TimestampValue{TimestampValue: timestamppb.New(v)} default: return nil, errors.Errorf("unsupported column type: %T", v) } row.TypedColumnValues = append(row.TypedColumnValues, columnValue) } resp.Rows = append(resp.Rows, row) } return resp, nil } for rows.Next() { row := &agent.Row{ ColumnValues: make([]string, req.ColumnNum), } columns := make([]any, req.ColumnNum) for i := range columns { columns[i] = any(&row.ColumnValues[i]) } if err := rows.Scan(columns...); err != nil { log.L().Error("scan SQL error", zap.String("sql", req.GetSql()), zap.Error(err)) return nil, err } resp.Rows = append(resp.Rows, row) } return resp, nil } func (s *AgentServer) LoadDataToCluster(ctx context.Context, req *agent.LoadDataReq) ( *agent.LoadDataResp, error, ) { resp := &agent.LoadDataResp{} cluster := tidbcloud.Cluster{ ClusterID: req.Cluster.ClusterId, Namespace: req.Cluster.Namespace, User: req.Cluster.User, Password: req.Cluster.Password, ReadTimeout: req.Cluster.ReadTimeout, WriteTimeout: req.Cluster.WriteTimeout, ProviderType: req.Cluster.ProviderType, } cli, err := tidbcloud.GetKubeClient(ctx) if err != nil { return nil, err } db, err := cluster.Connect(ctx, cli) if err != nil { return nil, err } defer db.Close() // nolint: errcheck conn, err := db.Conn(ctx) if err != nil { return nil, err } // nolint: errcheck defer conn.Close() // Execute SET statement _, err = conn.ExecContext(ctx, "set @@tidb_dml_batch_size=1000") if err != nil { log.L().Error("exec SET statement failed", zap.Error(err)) resp.ErrMessage = err.Error() return resp, nil } // Execute LOAD DATA statement loadDataSql := generateLoadDataSql( req.SourceUri, req.TargetDb, req.TargetTable, req.Terminator, req.Separator, req.Delimiter, req.BackslashEscape, req.IgnoreOneLine, ) _, err = conn.ExecContext(ctx, loadDataSql) if err != nil { log.L().Error("exec LOAD DATA statement failed", zap.Error(err), zap.String("sql", loadDataSql)) resp.ErrMessage = err.Error() } // It is not necessary to expose it to users temporarily. // We can comment it out for now. // // Execute SHOW WARNINGS // rows, err := db.QueryContext(ctx, "show warnings") // if err != nil { // log.L().Warn("show warnings failed", zap.Error(err)) // return resp, nil // } // defer rows.Close() // var ( // level string // code int // message string // ) // warnings := []string{} // for rows.Next() { // err := rows.Scan(&level, &code, &message) // if err != nil { // log.L().Error("scan show warnings failed", zap.Error(err)) // return resp, nil // } // warnings = append(warnings, fmt.Sprintf("%s (Code: %d): %s", level, code, message)) // } // err = rows.Err() // if err != nil { // log.L().Error("scan show warnings failed", zap.Error(err)) // return resp, nil // } // // Standardize the warnings as a single string // resp.WarningMessage = strings.Join(warnings, "; ") return resp, nil } func getLightningCheckItemID(item agent.CheckItem) lnprecheck.CheckItemID { switch item { case agent.CheckItem_CHECK_ITEM_TARGET_TABLE_EMPTY: return lnprecheck.CheckTargetTableEmpty case agent.CheckItem_CHECK_ITEM_SOURCE_SCHEMA: return lnprecheck.CheckSourceSchemaValid case agent.CheckItem_CHECK_ITEM_CSV_HEADER: return lnprecheck.CheckCSVHeader case agent.CheckItem_CHECK_ITEM_TARGET_CLUSTER_SIZE: return lnprecheck.CheckTargetClusterSize default: return lnprecheck.CheckItemID("") } } func (s *AgentServer) LightningPrecheck(ctx context.Context, req *agent.LightningPrecheckReq) ( *agent.LightningPrecheckResp, error, ) { // In order to avoid the context timeout and return directly, here we cancel 5 second in advance ctx, cancel := context.WithTimeout(ctx, 45*time.Second) defer cancel() plog.L().Debug("LightningPrecheckReq", zap.Any("req", req)) // init the lightning precheck builder theConfig := lnconfig.NewConfig() if err := theConfig.LoadFromTOML([]byte(req.LightningConfig)); err != nil { return nil, errors.Annotate(err, "load config data error") } // If the agent is running with flag `enableLocalClusterCert`, generate the security files // with the local cluster certificates instead of following the requests. // For now, this option is only used in the serverless environment. var pdTlsConfig *tls.Config if s.enableLocalClusterCert { lightning.AdjustSecFilesInConfig(theConfig, &tidbcloud.SecurityConfig{ CACert: constants.ClusterCAPath, TLSCert: constants.ClusterCertPath, TLSKey: constants.ClusterKeyPath, }) // Rebuild the TLS config from the adjusted file paths instead of reusing // whatever may have been derived from the original request payload. theConfig.Security.TLSConfig = nil err := theConfig.Security.BuildTLSConfig() if err != nil { return nil, errors.Annotate(err, "build tls config error") } // Certificates for clusters in a serverless environment are issued by `tidb-serverless-issuer`. // SQL certificates for lightnings are the same as cluster certificates. // In a serverless cluster, if a lightning tries to connect to a TiDB cluster, // it first connects to the inner TiDB gateway, which then forwards the request to an allocated TiDB pod. // However, the TiDB gateway's certificate is not issued by `tidb-serverless-issuer`, // so the lightning cannot connect to the TiDB gateway if it continues to verify the server certificate. // Therefore, the TLS mode should be set to at least `skip-verify` to skip server certificate verification. // We use `preferred` here to allow the inner gateway to fall back to insecure mode and convenient for service to upgrade in any order. // TODO(iosmanthus): Use an empty string instead of `preferred` when the inner gateway's TLS issues are resolved. theConfig.TiDB.TLS = "preferred" } else if len(theConfig.Security.CAPath) > 0 { // For some dev-tier cluster, the TLS is disabled on the TiDB cluster. // So the generated config has no security section. // In this situation, we don't need to generate the security files clusterClientCerts, err := tidbcloud.GenerateClusterClientSecurityFiles(ctx, req.GetClusterNamePrefix(), req.GetClusterNamespace()) defer tidbcloud.CleanupSecurityFiles(clusterClientCerts) if err != nil { return nil, errors.Annotate(err, "generate cluster-client mTLS security files error") } lightning.AdjustSecFilesInConfig(theConfig, clusterClientCerts) // Rebuild the TLS config from the generated temp files. Otherwise a stale // TLSConfig can survive from earlier config loading and the PD client may // fall back to the wrong trust chain. theConfig.Security.TLSConfig = nil if err := theConfig.Security.BuildTLSConfig(); err != nil { return nil, errors.Annotate(err, "build cluster-client tls config error") } // versions before 8.1 might get template with TLS="false" and lose all certs theConfig.TiDB.TLS = "" if theConfig.Security.TLSConfig == nil { return nil, errors.New("cluster-client tls config is nil after rebuilding") } tidbServerCerts, err := tidbcloud.GenerateTiDBServerSecurityFiles(ctx, req.GetClusterNamePrefix(), req.GetClusterNamespace()) defer tidbcloud.CleanupSecurityFiles(tidbServerCerts) if err != nil { return nil, errors.Annotate(err, "generate tidb-server security files error") } security := lnconfig.Security{ // this is tidb-server CA cert, we use it to verify tidb server cert CAPath: tidbServerCerts.CACert, // 1. we don't need to present client cert/key when connecting tidb-server, even if TLS is required/enabled. // 2. cert/key in `tidbServerCerts` is the server-side cert/key, can't be used as client-side cert/key //CertPath: tidbServerCerts.TLSCert, //KeyPath: tidbServerCerts.TLSKey, } if err := security.BuildTLSConfig(); err != nil { return nil, errors.Annotate(err, "build tidb-server tls config error") } theConfig.TiDB.Security = &security } if pdTlsConfig == nil { pdTlsConfig = theConfig.Security.TLSConfig } if pdTlsConfig != nil && pdTlsConfig.RootCAs != nil { plog.L().Info("rebuilt pd tls config for lightning precheck", zap.String("cluster_namespace", req.GetClusterNamespace()), zap.String("cluster_name_prefix", req.GetClusterNamePrefix()), zap.String("pd_addr", theConfig.TiDB.PdAddr), zap.String("ca_path", theConfig.Security.CAPath), zap.String("cert_path", theConfig.Security.CertPath), zap.String("key_path", theConfig.Security.KeyPath), zap.Int("root_ca_subjects", len(pdTlsConfig.RootCAs.Subjects()))) } else { plog.L().Warn("pd tls config is empty before lightning precheck", zap.String("cluster_namespace", req.GetClusterNamespace()), zap.String("cluster_name_prefix", req.GetClusterNamePrefix()), zap.String("pd_addr", theConfig.TiDB.PdAddr), zap.String("ca_path", theConfig.Security.CAPath), zap.String("cert_path", theConfig.Security.CertPath), zap.String("key_path", theConfig.Security.KeyPath)) } if req.GenGcpCredential { credentialsFile, err := tidbcloud.GenerateGcpCredentialsFile(ctx, req.GetClusterNamePrefix(), req.GetClusterNamespace()) if err != nil { return nil, errors.Annotate(err, "generate gcp credentials file error") } defer os.Remove(credentialsFile) // nolint: errcheck lightning.AdjustSourceDir(theConfig, credentialsFile) } lightning.AdjustTiDBHostNameInConfig(theConfig, req.GetClusterNamePrefix(), req.GetClusterNamespace()) if err := theConfig.Adjust(ctx); err != nil { return nil, errors.Annotate(err, "adjust lightning config error") } resp := &agent.LightningPrecheckResp{ BaseResp: &base.BaseResp{}, } fileIter, err := lightning.NewPartialFileIterator(ctx, theConfig) if err != nil { return nil, errors.Annotate(err, "new partial file iterator failed") } plog.L().Debug("adjusted config", zap.Any("theConfig", theConfig)) pdServiceDiscovery := pdServiceDiscovery.NewDefaultServiceDiscovery(ctx, cancel, []string{theConfig.TiDB.PdAddr}, pdTlsConfig) if err := pdServiceDiscovery.Init(); err != nil { return nil, errors.Annotate(err, "init pd service discovery failed") } pdHTTPCli := http.NewClientWithServiceDiscovery( "dataflow-agent", pdServiceDiscovery, http.WithTLSConfig(pdTlsConfig), ) defer pdHTTPCli.Close() precheckBuilder, err := lnrestore.NewPrecheckItemBuilderFromConfig( ctx, theConfig, pdHTTPCli, lnrestore_opts.WithPreInfoGetterOptions( lnrestore_opts.WithIgnoreDBNotExist(true), ), lnrestore_opts.WithMDLoaderSetupOptions( mydump.WithMaxScanFiles(lightning.MaxScanFilesCount), mydump.WithFileIterator(fileIter), ), ) if err != nil { if errors.ErrorEqual(err, lncommon.ErrTooManySourceFiles) { resp.BaseResp.ErrMsg = tooManySrcFilesErrMessage } else { return nil, errors.Annotate(err, "generate precheck builder from config error") } } for _, item := range req.CheckItems { lnCheckItemID := getLightningCheckItemID(item) if len(lnCheckItemID) <= 0 { // unknown or unsupported check item plog.L().Warn("unsupported check item. Skip it", zap.Any("check item", item)) continue } checker, err := precheckBuilder.BuildPrecheckItem(lnCheckItemID) if err != nil { return resp, errors.Annotatef(err, "build precheck item on '%s' error", item) } checkResult, err := checker.Check(ctx) if err != nil { if ctx.Err() == context.DeadlineExceeded || strings.Contains(err.Error(), "context deadline exceeded") { plog.L().Warn("meet the timeout limitation") return resp, nil } return resp, errors.Annotatef(err, "run precheck item '%s' error", item) } if checkResult != nil { resp.CheckResults = append(resp.CheckResults, &agent.CheckResult{ CheckItem: item, HasPassed: checkResult.Passed, Msg: checkResult.Message, }) if !checkResult.Passed && checkResult.Severity == lnprecheck.Critical { // this critical test has not passed, do not need to continue the rest break } } else { resp.CheckResults = append(resp.CheckResults, &agent.CheckResult{ CheckItem: item, HasPassed: true, Msg: "this check item has been skipped", }) } } return resp, nil } func (s *AgentServer) LightningSourcePreview(ctx context.Context, req *agent.LightningSourcePreviewReq) ( *agent.LightningSourcePreviewResp, error, ) { return nil, errors.New("not implemented") } func (s *AgentServer) GetImportSchemasMeta(ctx context.Context, req *agent.GetImportSchemasMetaReq) ( *agent.GetImportSchemasMetaResp, error, ) { return nil, errors.New("not implemented") } func (s *AgentServer) EstimateSourceDataSize(ctx context.Context, req *agent.EstimateSourceDataSizeReq) ( *agent.EstimateSourceDataSizeResp, error, ) { return nil, errors.New("not implemented") } func (a *AgentServer) ClearEmptyTables(ctx context.Context, req *agent.ClearEmptyTablesReq) (*agent.ClearEmptyTablesResp, error) { return nil, errors.New("not implemented") } func (s *AgentServer) getBaseDB(logger *zap.Logger, req *agent.CheckDBConnectivityReq) (*conn.BaseDB, error) { provider := conn.DefaultDBProvider dbCfg := &dbconfig.DBConfig{ Host: req.Host, Port: int(req.Port), User: req.User, Password: req.Password, } scopedDBCfg := conn.ScopedDBConfig{ DBConfig: dbCfg, } if len(req.SslCa) > 0 || len(req.SslCert) > 0 || len(req.SslKey) > 0 { dbCfg.Security = &security.Security{ SSLCABytes: []byte(req.SslCa), SSLCertBytes: []byte(req.SslCert), SSLKeyBytes: []byte(req.SslKey), } } baseDB, err := provider.Apply(scopedDBCfg) if err != nil { logger.Error("failed to connect db", zap.Error(err)) // DM error contains workaround about dm component, we don't want it exposed, we just return the raw error msg. err2 := errors.Cause(err) return nil, errors.New(err2.Error()) } return baseDB, nil } // CheckDBConnectivity check database connectivity. func (s *AgentServer) CheckDBConnectivity(ctx context.Context, req *agent.CheckDBConnectivityReq) (*agent.CheckDBConnectivityResp, error) { // todo: the middleware logger should print request id too. logger := plog.L().With(zap.String("request-id", req.RequestId)) baseDB, err := s.getBaseDB(logger, req) if err != nil { return nil, err } defer utils.LogClose(baseDB, logger) dbVersion, err := GetDBVersion(ctx, baseDB.DB) if err != nil { msg := "failed to get db version" logger.Error(msg, zap.Error(err)) return nil, errors.Annotate(err, msg) } return &agent.CheckDBConnectivityResp{DbVersion: dbVersion}, nil } func (s *AgentServer) GetTiDBServerSecret(ctx context.Context, req *agent.GetTiDBServerSecretReq) (*agent.GetTiDBServerSecretResp, error) { resp := &agent.GetTiDBServerSecretResp{} securityConfig, err := tidbcloud.GetTiDBServerSecurity(ctx, req.GetClusterNamePrefix(), req.GetClusterNamespace()) if err != nil { return nil, err } resp.SslCa = string(securityConfig.CACertBytes) resp.SslCert = string(securityConfig.TLSCertBytes) resp.SslKey = string(securityConfig.TLSKeyBytes) return resp, nil } func (s *AgentServer) GetClusterClientSecret(ctx context.Context, req *agent.GetClusterClientSecretReq) (*agent.GetClusterClientSecretResp, error) { resp := &agent.GetClusterClientSecretResp{} securityConfig, err := tidbcloud.GetClusterClientSecurity(ctx, req.GetClusterNamePrefix(), req.GetClusterNamespace()) if err != nil { return nil, err } resp.SslCa = string(securityConfig.CACertBytes) resp.SslCert = string(securityConfig.TLSCertBytes) resp.SslKey = string(securityConfig.TLSKeyBytes) return resp, nil } func (s *AgentServer) PrecheckDMJob(ctx context.Context, req *agent.PrecheckDMJobReq) (*agent.PrecheckDMJobResp, error) { plog.L().Info("validate dm job", zap.Reflect("req", req)) cfgs := make([]*dmconfig.SubTaskConfig, len(req.TaskConfig)) resp := &agent.PrecheckDMJobResp{} for i, cfg := range req.TaskConfig { subCfg := &dmconfig.SubTaskConfig{} if err := subCfg.Decode(cfg, false); err != nil { plog.L().Error("fail to build subtask config", zap.Error(err)) return nil, err } subCfg.From.Password = req.Password subCfg.To.Password = req.ClusterPassword cfgs[i] = subCfg } results, err := checker.RunCheckOnConfigs(ctx, cfgs, req.FullInstanceMigrate, 10, 10) if err != nil { plog.L().Error("fail to validate dm job", zap.Error(err)) return nil, err } plog.L().Info("precheck", zap.Reflect("precheck result", results)) retStr, err := json.Marshal(results) if err != nil { plog.L().Error("fail to marshal check results", zap.Error(err)) return nil, err } resp.CheckResults = retStr return resp, nil } func (s *AgentServer) GetUserFileContent(ctx context.Context, req *agent.GetUserFileReq) (*agent.GetUserFileResp, error) { logger := plog.L().With(zap.String("request-id", req.RequestId)) resp := &agent.GetUserFileResp{} var ( content string err error ) switch dataflow.ImportDataSource(req.DataSourceType) { case dataflow.ImportDataSource_S3_SOURCE: content, err = s3.GetFileContent(ctx, req.SourceUrl, req.RoleArn, req.ExternalId, req.StartOffset, req.EndOffset) case dataflow.ImportDataSource_GCS_SOURCE: content, err = gcs.GetFileContent(ctx, req.SourceUrl, req.GetGcsOptions().GetImpersonateServiceAccount(), req.StartOffset, req.EndOffset) case dataflow.ImportDataSource_AZBLOB_SOURCE: content, err = azblob.GetFileContent(ctx, req.SourceUrl, req.GetSasToken(), req.StartOffset, req.EndOffset) } if err != nil { logger.Error("failed to get user file content", zap.Error(err)) return nil, err } resp.Content = content return resp, nil } // GetSourceTables get all source tables. func (s *AgentServer) GetSourceTables(ctx context.Context, req *agent.CheckDBConnectivityReq) (*agent.GetSourceTablesResp, error) { // todo: the middleware logger should print request id too. logger := plog.L().With(zap.String("request-id", req.RequestId)) resp := &agent.GetSourceTablesResp{} baseDB, err := s.getBaseDB(logger, req) if err != nil { return nil, err } defer utils.LogClose(baseDB, logger) resp.Schemas, err = GetAllTables(ctx, baseDB.DB, commondm.IsMysqlSystemDB) if err != nil { logger.Error("failed to get all table", zap.Error(err)) return nil, err } return resp, nil } func (s *AgentServer) GetSourceBinlogStatus(ctx context.Context, req *agent.CheckDBConnectivityReq) (*agent.GetSourceBinlogStatusResp, error) { logger := plog.L().With(zap.String("request-id", req.RequestId)) baseDB, err := s.getBaseDB(logger, req) if err != nil { return nil, err } defer utils.LogClose(baseDB, logger) flavor, err := conn.GetFlavor(ctx, baseDB) if err != nil { message := "failed to get db flavor" logger.Error(message, zap.Error(err)) return nil, errors.New(message) } logger.Info("got db flavor", zap.String("flavor", flavor)) binlogName, binlogPos, doDBs, ignoreDBs, gtidStr, err := conn.GetMasterStatus(tcontext.NewContext(ctx, log.Logger{Logger: logger}), baseDB, flavor) if err != nil { msg := "failed to get master status" logger.Error(msg, zap.Error(err)) return nil, errors.New(msg) } return &agent.GetSourceBinlogStatusResp{ BinlogName: binlogName, BinlogPos: uint32(binlogPos), BinlogDoDbs: doDBs, BinlogIgnoreDbs: ignoreDBs, GtidStr: gtidStr, }, nil } // GetClusterTables get all cluster tables. func (s *AgentServer) GetClusterTables(ctx context.Context, req *agent.Cluster) (*agent.GetClusterTablesResp, error) { logger := plog.L().With(zap.Uint64("get cluster tables", req.ClusterId)) resp := &agent.GetClusterTablesResp{} cluster := tidbcloud.Cluster{ ClusterID: req.ClusterId, Namespace: req.Namespace, User: req.User, Password: req.Password, ReadTimeout: req.ReadTimeout, WriteTimeout: req.WriteTimeout, ProviderType: req.ProviderType, } cli, err := tidbcloud.GetKubeClient(ctx) if err != nil { return nil, err } db, err := cluster.Connect(ctx, cli) if err != nil { return nil, err } defer db.Close() // nolint: errcheck resp.Schemas, err = GetAllTables(ctx, db, metadef.IsMemOrSysDB) if err != nil { logger.Error("failed to get all table", zap.Error(err)) return nil, err } return resp, nil } func (s *AgentServer) NsLookup(ctx context.Context, req *agent.NsLookupReq) (*agent.NsLookupResp, error) { logger := plog.L().With(zap.String("request-id", req.RequestId)) addrs, err := net.DefaultResolver.LookupIPAddr(ctx, req.Host) if err != nil { return nil, err } if len(addrs) == 0 { // not sure whether will this happen msg := "ips returned by net.LookupIP is empty" logger.Error(msg) return nil, errors.New(msg) } resp := &agent.NsLookupResp{} for _, addr := range addrs { resp.Ips = append(resp.Ips, addr.IP.String()) } return resp, nil } func (s *AgentServer) GetImportFiles(ctx context.Context, req *agent.GetImportFilesReq) (*agent.GetImportFilesResp, error) { sourceUri := req.SourceUri sdk, err := importsdk.NewImportSDK(ctx, sourceUri, nil) if err != nil { return nil, err } metas, err := sdk.GetTableMetas(ctx) if err != nil { return nil, err } tableMetas := []*agent.GetImportFilesResp_TableMeta{} for _, meta := range metas { dataFiles := []*agent.GetImportFilesResp_DataFileMeta{} for _, f := range meta.DataFiles { df := &agent.GetImportFilesResp_DataFileMeta{ Path: f.Path, Size: uint64(f.Size), SourceType: agent.GetImportFilesResp_SourceType(f.Format), Compression: agent.GetImportFilesResp_Compression(f.Compression), } dataFiles = append(dataFiles, df) } tableMeta := &agent.GetImportFilesResp_TableMeta{ Database: meta.Database, Table: meta.Table, DataFiles: dataFiles, TotalSize: uint64(meta.TotalSize), WildcardPath: meta.WildcardPath, SchemaFile: meta.SchemaFile, } tableMetas = append(tableMetas, tableMeta) } resp := &agent.GetImportFilesResp{ TableMeta: tableMetas, } return resp, nil } func (s *AgentServer) CreateSchemas(ctx context.Context, req *agent.CreateSchemasReq) (*emptypb.Empty, error) { logger := plog.L().With(zap.Uint64("cluster-id", req.ClusterId)) db, err := s.clusterProvider.GetDB(ctx, fmt.Sprint(req.ClusterId), req.User, req.Password, nil) if err != nil { return nil, err } defer db.Close() // nolint: errcheck sourceUri := req.SourceUri sdk, err := importsdk.NewImportSDK(ctx, sourceUri, db) if err != nil { return nil, err } if err := sdk.CreateSchemasAndTables(ctx); err != nil { logger.Error("failed to create schemas and tables", zap.Error(err)) return nil, status.Errorf(codes.Internal, "failed to create schemas and tables: %v", err) } return &emptypb.Empty{}, nil } func (s *AgentServer) SearchImportFiles(ctx context.Context, req *agent.SearchImportFilesReq) (*agent.SearchImportFilesResp, error) { logger := plog.L().With(zap.Uint32("data format", req.DataFormat)) // In order to avoid the context timeout and return directly, here we cancel 5 second in advance ctx, cancel := context.WithTimeout(ctx, 45*time.Second) defer cancel() resp := &agent.SearchImportFilesResp{ TableMatches: make([]*agent.SearchImportFilesResp_TableMatch, 0, len(req.TargetSource)), } ch := make(chan *agent.SearchImportFilesReq_TargetSource) g, ctx := errgroup.WithContext(ctx) g.Go(func() error { defer close(ch) for _, targetSource := range req.TargetSource { logger.Info("checking target source", zap.Any("target source", targetSource)) ch <- targetSource } return nil }) var ( client storagepkg.StorageClient err error sourceUri = req.GetSourceUri() bucket string uriPrefix string ) switch dataflow.ImportDataSource(req.DataSourceType) { case dataflow.ImportDataSource_S3_SOURCE: client, err = storagepkg.NewS3Client(ctx, sourceUri, req.AwsRoleArn, req.AwsExternalId, req.AwsAccessKeyId, req.AwsSecretAccessKey, ) case dataflow.ImportDataSource_GCS_SOURCE: client, err = storagepkg.NewGcsClient(ctx, sourceUri, req.GetGcsOptions().GetImpersonateServiceAccount()) case dataflow.ImportDataSource_AZBLOB_SOURCE: client, err = storagepkg.NewAzblobClient(ctx, sourceUri, req.GetSasToken()) default: return nil, errors.New(fmt.Sprintf("unsupported data source type %d", req.DataSourceType)) } if err != nil { return nil, err } bucket, uriPrefix, err = storagepkg.ParseUri(sourceUri) if err != nil { return nil, err } forLoop: for { select { case <-ctx.Done(): if ctx.Err() == context.DeadlineExceeded { logger.Warn("meet the timeout limitation") return resp, nil } case targetSource, ok := <-ch: if !ok { logger.Info("all target sources are checked") break forLoop } var uriMatcher matcher.Matcher if targetSource.CustomFile != "" { uriMatcher = matcher.NewCustomMatcher(client, bucket, uriPrefix, sourceUri, req.DataFormat) } else { uriMatcher = matcher.NewDefaultMatcher(client, bucket, uriPrefix, req.DataFormat) } if matchedUri, err := uriMatcher.Match(ctx, targetSource); err != nil { if ctx.Err() == context.DeadlineExceeded || strings.Contains(err.Error(), "context deadline exceeded") { logger.Warn("meet the timeout limitation") return resp, nil } return nil, err } else { logger.Info("after matcher match", zap.String("matchedUri", matchedUri), zap.Any("table", targetSource.TargetTable)) resp.TableMatches = append(resp.TableMatches, &agent.SearchImportFilesResp_TableMatch{ TargetTable: targetSource.TargetTable, SourceUri: matchedUri, }) if len(matchedUri) > 0 { logger.Info("uri matched", zap.String("matchedUri", matchedUri), zap.Any("table", targetSource.TargetTable)) } } } } // nolint: errcheck g.Wait() return resp, nil } func (s *AgentServer) PreviewCsvFile(ctx context.Context, req *agent.PreviewCsvFileReq) (*agent.PreviewCsvFileResp, error) { logger := plog.L().With(zap.String("source uri", req.SourceUri), zap.Any("csv config", req.CsvConfig)) parser, err := fileparser.NewCsvParser(ctx, req.SourceUri, req.CsvConfig, &tdbrstorage.BackendOptions{ S3: tdbrstorage.S3BackendOptions{ RoleARN: req.AwsRoleArn, ExternalID: req.AwsExternalId, AccessKey: req.AwsAccessKeyId, SecretAccessKey: req.AwsSecretAccessKey, }, }) if err != nil { return nil, errors.Annotate(err, "create source data parser error") } //nolint: errcheck defer parser.Close() limitCount := 5 if req.FirstN > 0 { limitCount = int(req.FirstN) } sourceData := &agent.Table{ RowValues: make([]*agent.Row, limitCount), } for i := 0; i < limitCount; i++ { writeRow, err := fileparser.ReadRow(parser) if err != nil { logger.Error("read row error", zap.Error(err)) if errors.Cause(err) == io.EOF { sourceData.RowValues = sourceData.RowValues[:i] break } return nil, errors.Annotate(err, "read row error") } sourceData.RowValues[i] = writeRow } if len(parser.Columns()) != 0 { sourceData.RowHeaders = parser.Columns() } return &agent.PreviewCsvFileResp{Table: sourceData}, nil } func (s *AgentServer) PrecheckImportFile(ctx context.Context, req *agent.PrecheckImportFileReq) (*emptypb.Empty, error) { cluster := tidbcloud.Cluster{ ClusterID: req.Cluster.ClusterId, Namespace: req.Cluster.Namespace, User: req.Cluster.User, Password: req.Cluster.Password, ReadTimeout: req.Cluster.ReadTimeout, WriteTimeout: req.Cluster.WriteTimeout, ProviderType: req.Cluster.ProviderType, } logger := plog.L() logger.Info("PrecheckImportFile", zap.Any("req", req)) cli, err := tidbcloud.GetKubeClient(ctx) if err != nil { return nil, err } db, err := cluster.Connect(ctx, cli) if err != nil { return nil, err } defer db.Close() // nolint: errcheck var accountName string if entity.DataSource(req.DataSource) == entity.ImportDataSourceAzblob { azurl, err := azblob.ParseAzblobURL(req.SourceUrl) if err != nil { return nil, err } accountName = azurl.AccountName() } builder := fileparser.NewParserBuilder(ctx, entity.DataFormat(req.DataFormat), req.CsvConfig, &tdbrstorage.BackendOptions{ S3: tdbrstorage.S3BackendOptions{ RoleARN: req.AwsRoleArn, ExternalID: req.AwsExternalId, AccessKey: req.AwsAccessKeyId, SecretAccessKey: req.AwsSecretAccessKey, }, Azblob: tdbrstorage.AzblobBackendOptions{ SASToken: req.GetSasToken(), AccountName: accountName, }, }) supportedCheckColumnCompatitable := map[entity.DataFormat]struct{}{ entity.ImportDataFormatCsv: {}, entity.ImportDataFormatSql: {}, } warningTables := []*warningTable{} for _, source := range req.Sources { if err := checkTableEmpty(ctx, db, source.Table.Schema, source.Table.Table); err != nil { return nil, err } if _, ok := supportedCheckColumnCompatitable[entity.DataFormat(req.DataFormat)]; !ok { continue } sourceURL := source.File if entity.DataSource(req.DataSource) == entity.ImportDataSourceAzblob && sourceURL != "" { sourceURL, err = azblob.ConvertAzblobURL(sourceURL, false) if err != nil { return nil, err } } warnTable, err := checkColumnCompatitable(ctx, db, source.Table.Schema, source.Table.Table, sourceURL, builder) if err != nil { return nil, err } if warnTable != nil { warningTables = append(warningTables, warnTable) } } resp := &emptypb.Empty{} // generate warning message by warning tables if len(warningTables) != 0 { tblWarnMsg := "" for i, tbl := range warningTables { tblWarnMsg += fmt.Sprintf("table `%s`.`%s`, columns %s", tbl.schema, tbl.table, strings.Join(tbl.cols, ",")) if i == 3 { tblWarnMsg += "..." break } else { tblWarnMsg += "; " } } WarningMessage := fmt.Sprintf("Some columns in the target table are missing in your source and can not be generated. You can either remove the NOT_NULL constraint or add a default value for them. For example, %s", tblWarnMsg) return resp, status.Error(codes.PermissionDenied, WarningMessage) } return resp, nil } func (s *AgentServer) DropClusterTables(ctx context.Context, req *agent.DropClusterTablesReq) ( *emptypb.Empty, error, ) { cluster := tidbcloud.Cluster{ ClusterID: req.Cluster.ClusterId, Namespace: req.Cluster.Namespace, User: req.Cluster.User, Password: req.Cluster.Password, ReadTimeout: req.Cluster.ReadTimeout, WriteTimeout: req.Cluster.WriteTimeout, ProviderType: req.Cluster.ProviderType, } cli, err := tidbcloud.GetKubeClient(ctx) if err != nil { return nil, err } db, err := cluster.Connect(ctx, cli) if err != nil { return nil, err } defer db.Close() // nolint: errcheck conn, err := db.Conn(ctx) if err != nil { return nil, err } // nolint: errcheck defer conn.Close() dropTableSql := generateDropTableSql(req.Tables) _, err = conn.ExecContext(ctx, dropTableSql) if err != nil { log.L().Error("exec DROP TABLE statement failed", zap.Error(err), zap.String("sql", dropTableSql)) return nil, err } return &emptypb.Empty{}, nil } func (s *AgentServer) createPrecheckStorage(ctx context.Context, req *agent.PrecheckImportJobReq) (cloud_storage.CloudStorage, error) { // Check if source is provided if req.Source == nil { return nil, errors.New("source is required") } // Determine the source type and create appropriate storage switch req.Source.Type { case agent.ImportSourceType_S3: return s.createS3Storage(ctx, req.Source.GetS3()) case agent.ImportSourceType_GCS: // Pass cluster namespace to get GCP credentials from Secret return s.createGCSStorage(ctx, req.Source.GetGcs(), req.TargetCluster.GetNamespace()) case agent.ImportSourceType_AZURE_BLOB: return s.createAzureBlobStorage(ctx, req.Source.GetAzureBlob()) default: return nil, errors.Errorf("unsupported source type: %v", req.Source.Type) } } func (s *AgentServer) getS3Client(cfg aws.Config, s3Source *agent.S3Source) (*awss3.Client, error) { var s3Client *awss3.Client switch s3Source.AuthType { case agent.ImportS3AuthType_ROLE_ARN: if roleArn := s3Source.GetRoleArn(); roleArn.RoleArn != "" { plog.L().Info("start to init role-arn credentials", zap.String("roleArn", roleArn.RoleArn)) stsClient := sts.NewFromConfig(cfg) creds := stscreds.NewAssumeRoleProvider(stsClient, roleArn.GetRoleArn(), func(options *stscreds.AssumeRoleOptions) { externalID := roleArn.GetExternalId() options.ExternalID = &externalID }) cfg.Credentials = creds s3Client = awss3.NewFromConfig(cfg) } else { return nil, errors.New("role ARN is required for ROLE_ARN auth type") } case agent.ImportS3AuthType_ACCESS_KEY: if accessKey := s3Source.GetAccessKey(); accessKey != nil { plog.L().Info("start to init role-arn credentials") cfg.Credentials = credentials.NewStaticCredentialsProvider( accessKey.Id, accessKey.Secret, "", ) s3Client = awss3.NewFromConfig(cfg) } else { return nil, errors.New("access key is required for ACCESS_KEY auth type") } default: return nil, errors.Errorf("unsupported S3 auth type: %v", s3Source.AuthType) } return s3Client, nil } func (s *AgentServer) createS3Storage(ctx context.Context, s3Source *agent.S3Source) (cloud_storage.CloudStorage, error) { if s3Source == nil { return nil, errors.New("S3 source configuration is required") } // Parse S3 URL _, err := s3.NewS3URL(s3Source.Uri) if err != nil { return nil, errors.Annotate(err, "failed to parse S3 URL") } // Create S3 storage using cloud_storage.NewS3 with appropriate auth options // This will automatically detect and set the correct bucket region var opts []cloud_storage.S3Option switch s3Source.AuthType { case agent.ImportS3AuthType_ROLE_ARN: if roleArn := s3Source.GetRoleArn(); roleArn != nil && roleArn.RoleArn != "" { plog.L().Info("start to init role-arn credentials", zap.String("roleArn", roleArn.RoleArn)) opts = append(opts, cloud_storage.WithAssumeRole(roleArn.GetRoleArn(), roleArn.GetExternalId())) } else { return nil, errors.New("role ARN is required for ROLE_ARN auth type") } case agent.ImportS3AuthType_ACCESS_KEY: if accessKey := s3Source.GetAccessKey(); accessKey != nil { plog.L().Info("start to init access-key credentials") opts = append(opts, cloud_storage.WithAccessKey(accessKey.Id, accessKey.Secret)) } else { return nil, errors.New("access key is required for ACCESS_KEY auth type") } default: return nil, errors.Errorf("unsupported S3 auth type: %v", s3Source.AuthType) } // Create S3 storage - NewS3 will automatically detect bucket region plog.L().Info("finished init s3 storage", zap.String("s3Source.Uri", s3Source.Uri)) return cloud_storage.NewS3(ctx, s3Source.Uri, opts...) } func (s *AgentServer) createGCSStorage(ctx context.Context, gcsSource *agent.GCSSource, clusterNamespace string) (cloud_storage.CloudStorage, error) { if gcsSource == nil { return nil, errors.New("GCS source configuration is required") } // Parse GCS URL _, err := gcs.NewGcsURL(gcsSource.Uri) if err != nil { return nil, errors.Annotate(err, "failed to parse GCS URL") } // Create GCS storage based on auth type switch gcsSource.AuthType { case agent.ImportGcsAuthType_SERVICE_ACCOUNT_KEY: // User provides their own service account JSON key (base64 encoded) if serviceAccountKey := gcsSource.GetServiceAccountKey(); serviceAccountKey != "" { plog.L().Info("using user-provided service account key for GCS", zap.String("uri", gcsSource.Uri)) return cloud_storage.NewGCS(ctx, gcsSource.Uri, cloud_storage.WithGCSAccountKey(serviceAccountKey)) } return nil, errors.New("service account key is required for SERVICE_ACCOUNT_KEY auth type") case agent.ImportGcsAuthType_AUTH_TYPE_UNSPECIFIED: // For TiDB Dedicated: user grants permission to TiDB Cloud's Service Account // Get GCP credentials from cluster namespace's Secret (gcp-br-credentials) if clusterNamespace == "" { return nil, errors.New("cluster namespace is required for GCS credential retrieval") } plog.L().Info("using GCP credentials from Secret", zap.String("uri", gcsSource.Uri), zap.String("namespace", clusterNamespace)) credentialsJSON, err := tidbcloud.GetGcpCredentialsFromSecret(ctx, clusterNamespace) if err != nil { return nil, errors.Annotate(err, fmt.Sprintf("failed to get GCP credentials from Secret: ensure Kubernetes Secret %q exists in namespace %q and contains key %q with the GCP service account JSON used by TiDB Cloud's Service Account", "gcp-br-credentials", clusterNamespace, "credentials")) } client, err := gcsstorage.NewClient(ctx, option.WithCredentialsJSON(credentialsJSON)) if err != nil { return nil, errors.Annotate(err, "failed to create GCS client with credentials") } return cloud_storage.NewGCS(ctx, gcsSource.Uri, cloud_storage.WithGcsClient(client)) default: return nil, errors.Errorf("unsupported GCS auth type: %v", gcsSource.AuthType) } } func (s *AgentServer) createAzureBlobStorage(ctx context.Context, azSource *agent.AzureBlobSource) (cloud_storage.CloudStorage, error) { if azSource == nil { return nil, errors.New("Azure Blob source configuration is required") } // Parse Azure Blob URL azURL, err := azblob.ParseAzblobURL(azSource.Uri) if err != nil { return nil, errors.Annotate(err, "failed to parse Azure Blob URL") } // Create Azure Blob client based on auth type var containerClient *container.Client switch azSource.AuthType { case agent.ImportAzureBlobAuthType_SAS_TOKEN: if sasToken := azSource.GetSasToken(); sasToken != "" { azblobStorage, err := azblob.NewAzureBlobStorage(ctx, azSource.Uri, sasToken, nil) if err != nil { return nil, errors.Annotate(err, "failed to create Azure Blob storage") } client, err := azblobStorage.Container(ctx, azURL.Container()) if err != nil { return nil, errors.Annotate(err, "failed to get container client") } containerClient = client } else { return nil, errors.New("SAS token is required for SAS_TOKEN auth type") } default: return nil, errors.Errorf("unsupported Azure Blob auth type: %v", azSource.AuthType) } // Create Azure Blob storage return cloud_storage.NewAzBlob(ctx, azSource.Uri, cloud_storage.WithAzBlobContainerClient(containerClient)) } func (s *AgentServer) createFileRouter(ctx context.Context, req *agent.PrecheckImportJobReq) (precheck.FileRouter, error) { rules := []*precheck.FileRouteRule{ {Pattern: DefaultDBSchemaFilePattern, Schema: "$1", Type: precheck.SchemaSchema}, {Pattern: DefaultTableSchemaFilePattern, Schema: "$1", Table: "$2", Type: precheck.TableSchema}, } // Get target table information from source var targetTables []*agent.ImportTargetTableInfo source := req.GetSource() if source != nil { switch source.Type { case agent.ImportSourceType_S3: if s3Source := source.GetS3(); s3Source != nil { targetTables = s3Source.GetTargetTableInfos() } case agent.ImportSourceType_GCS: if gcsSource := source.GetGcs(); gcsSource != nil { targetTables = gcsSource.GetTargetTableInfos() } case agent.ImportSourceType_AZURE_BLOB: if azSource := source.GetAzureBlob(); azSource != nil { targetTables = azSource.GetTargetTableInfos() } } } if len(targetTables) == 0 { // No custom target tables, use default file rule rule, err := s.getDefaultFileRule(req.GetImportOptions()) if err != nil { return nil, err } if rule != nil { rules = append(rules, rule) } } else { // Has custom target tables, create custom file rules customRules, err := s.getCustomFileRules(targetTables, req.GetImportOptions()) if err != nil { return nil, err } rules = append(rules, customRules...) } plog.L().Debug("create file router rule", zap.Any("rules", rules)) return precheck.NewFileRouter(ctx, rules) } func (s *AgentServer) getDefaultFileRule(importOptions *agent.ImportOptions) (*precheck.FileRouteRule, error) { if importOptions == nil { return nil, errors.New("import options is required") } fileType := importOptions.GetFileType() switch fileType { case agent.ImportFileType_CSV: return &precheck.FileRouteRule{ Pattern: fmt.Sprintf(DefaultFilePattern, "csv"), Schema: "$1", Table: "$2", Type: precheck.TypeCSV, }, nil case agent.ImportFileType_PARQUET: return &precheck.FileRouteRule{ Pattern: DefaultRawParquetFilePattern, Schema: "$1", Table: "$2", Type: precheck.TypeParquet, }, nil case agent.ImportFileType_SQL: return &precheck.FileRouteRule{ Pattern: fmt.Sprintf(DefaultFilePattern, "sql"), Schema: "$1", Table: "$2", Type: precheck.TypeSQL, }, nil case agent.ImportFileType_AURORA_SNAPSHOT: return &precheck.FileRouteRule{ Pattern: DefaultAuroraSnapshotFilePattern, Schema: "$1", Table: "$2", Type: precheck.TypeParquet, }, nil default: return nil, errors.Errorf("unsupported file type: %v", fileType) } } func (s *AgentServer) getCustomFileRules(targetTables []*agent.ImportTargetTableInfo, importOptions *agent.ImportOptions) ([]*precheck.FileRouteRule, error) { if importOptions == nil { return nil, errors.New("import options is required") } rules := make([]*precheck.FileRouteRule, 0, len(targetTables)) fileType := importOptions.GetFileType() for _, tableInfo := range targetTables { customFile := tableInfo.GetCustomFile() if customFile == "" { return nil, errors.New("custom file pattern is empty") } var fileTypeStr string switch fileType { case agent.ImportFileType_CSV: fileTypeStr = precheck.TypeCSV case agent.ImportFileType_PARQUET, agent.ImportFileType_AURORA_SNAPSHOT: fileTypeStr = precheck.TypeParquet case agent.ImportFileType_SQL: fileTypeStr = precheck.TypeSQL default: return nil, errors.Errorf("unsupported file type: %v", fileType) } targetTable := tableInfo.GetTargetTable() if targetTable == nil { return nil, errors.New("target table is required") } // Generate custom naming rule regexp pattern // This would be similar to utils.GenCustomNamingRuleRegexp in serverless version customPattern := s.generateCustomFilePattern(customFile, targetTable.Schema, targetTable.Table, fileType) rules = append(rules, &precheck.FileRouteRule{ Pattern: customPattern, Schema: targetTable.Schema, Table: targetTable.Table, Type: fileTypeStr, }) } return rules, nil } func (s *AgentServer) generateCustomFilePattern(customFile, schema, table string, fileType agent.ImportFileType_Enum) string { // Generate custom file pattern based on the custom file name and file type // This is a simplified version that creates a pattern based on the custom file if customFile == "" { // Fallback to default pattern if custom file is empty switch fileType { case agent.ImportFileType_CSV: return fmt.Sprintf(DefaultCSVFilePatternWithTargetTable, schema, table) case agent.ImportFileType_SQL: return fmt.Sprintf(DefaultSQLFilePatternWithTargetTable, schema, table) case agent.ImportFileType_PARQUET: return fmt.Sprintf(DefaultRawParquetFilePatternWithTargetTable, schema, table) case agent.ImportFileType_AURORA_SNAPSHOT: return fmt.Sprintf(DefaultAuroraSnapshotFilePatternWithTargetTable, schema, schema, table) default: return customFile } } return genCustomNamingRuleRegexp(customFile) } func genCustomNamingRuleRegexp(wildcards string) string { var sb strings.Builder sb.WriteRune('^') for cursor := 0; cursor < len(wildcards); { r, size := utf8.DecodeRuneInString(wildcards[cursor:]) switch r { case '{', '}', '(', ')', '^', '$', '.', '|', '+', '\\': sb.WriteByte('\\') sb.WriteRune(r) case '?': sb.WriteByte('.') case '*': sb.WriteString(".*") default: sb.WriteRune(r) } cursor += size } sb.WriteRune('$') return sb.String() } func (s *AgentServer) PrecheckImportJob(ctx context.Context, req *agent.PrecheckImportJobReq) (*agent.PrecheckImportJobResp, error) { storage, err := s.createPrecheckStorage(ctx, req) if err != nil { return nil, err } plog.L().Info("start to createFileRouter") fileRouter, err := s.createFileRouter(ctx, req) if err != nil { return nil, err } opts := &tidbcloud.DBOptions{ ReadTimeout: req.TargetCluster.ReadTimeout, WriteTimeout: req.TargetCluster.WriteTimeout, } passwd, err := s.getClusterPassword(ctx, req.TargetCluster) if err != nil { if goerr.Is(err, crypto.PEMUpdatedError) { return &agent.PrecheckImportJobResp{ ErrorMessage: crypto.PEMUpdatedErrorMsg, }, nil } return nil, err } db, err := s.clusterProvider.GetDB(ctx, fmt.Sprint(req.TargetCluster.ClusterId), req.TargetCluster.User, passwd, opts) if err != nil { return nil, err } defer db.Close() // nolint: errcheck resp, err := precheck.NewChecker(storage, fileRouter, precheck.NewDBSchemaChecker(db), false).Run(ctx) if err != nil { return nil, err } plog.L().Info("finished run checker", zap.Any("resp", resp)) return convertToPrecheckResp(resp), nil } func convertToPrecheckResp(result *precheck.PrecheckResult) *agent.PrecheckImportJobResp { resp := &agent.PrecheckImportJobResp{} if len(result.TableResults) == 0 { // Keep the legacy empty-completed response unchanged, but preserve an // incomplete scan status so callers can tell an empty result apart from // a scan that did not examine every source file. switch result.Status { case precheck.PrecheckStatusTruncated: resp.IsTruncated = true resp.Status = agent.ImportPrecheckStatusType_TRUNCATED case precheck.PrecheckStatusTimeout: resp.Status = agent.ImportPrecheckStatusType_TIMEOUT } return resp } var totalFiles uint64 tablesWithErrors := make([]string, 0) // Process all results - each result represents a unique table for _, res := range result.TableResults { tableKey := res.Table.Schema + "." + res.Table.Name // Create table result and add to response tableResult := &agent.TableResult{ TargetTable: &common.Table{ Schema: res.Table.Schema, Table: res.Table.Name, }, MatchedDataFiles: res.DataFiles, MatchedSchemaFile: res.SchemaFile, UseSchemaFile: res.UseSchemaFile, } // Handle tables with errors if res.ErrorMsg != "" { tablesWithErrors = append(tablesWithErrors, tableKey) tableResult.ErrorMessage = res.ErrorMsg } resp.TargetTables = append(resp.TargetTables, tableResult) totalFiles += uint64(len(res.DataFiles)) } // Set counts and warning message resp.TotalTablesCount = uint64(len(result.TableResults)) resp.TotalDataFilesCount = totalFiles resp.IsTruncated = result.Status == precheck.PrecheckStatusTruncated if len(tablesWithErrors) > 0 { resp.ErrorMessage = fmt.Sprintf("Found %d table(s) with error: %s", len(tablesWithErrors), strings.Join(tablesWithErrors, ", ")) } resp.Status = convertPrecheckStatus(result.Status) return resp } func convertPrecheckStatus(status precheck.PrecheckStatus) agent.ImportPrecheckStatusType_Enum { switch status { case precheck.PrecheckStatusCompleted: return agent.ImportPrecheckStatusType_COMPLETED case precheck.PrecheckStatusTruncated: return agent.ImportPrecheckStatusType_TRUNCATED case precheck.PrecheckStatusTimeout: return agent.ImportPrecheckStatusType_TIMEOUT default: return agent.ImportPrecheckStatusType_UNSPECIFIED } } func (s *AgentServer) GetPublicKey(context.Context, *agent.GetPublicKeyRequest) (*agent.GetPublicKeyResponse, error) { pem, err := s.encryptor.ExportPublicKeyAsPEM() if err != nil { return nil, err } return &agent.GetPublicKeyResponse{Pem: pem}, nil }