package pimport import ( "context" "fmt" "path/filepath" "strconv" "strings" "time" "github.com/aws/aws-sdk-go/aws" awss3 "github.com/aws/aws-sdk-go/service/s3" "github.com/docker/go-units" "github.com/dustin/go-humanize" "github.com/pingcap/errcode" "github.com/pingcap/errors" agentpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent" "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow" "github.com/tidbcloud/dataflow-service/internal/bizutils" "github.com/tidbcloud/dataflow-service/internal/common" "github.com/tidbcloud/dataflow-service/internal/db/entity" "github.com/tidbcloud/dataflow-service/internal/grpc/grpcli" "github.com/tidbcloud/dataflow-service/internal/model/cluster" pimport "github.com/tidbcloud/dataflow-service/internal/model/import" "github.com/tidbcloud/dataflow-service/internal/service" "github.com/tidbcloud/dataflow-service/internal/utils" "github.com/tidbcloud/dataflow-service/pkg/account/accountdb" infraapi "github.com/tidbcloud/dataflow-service/pkg/infra-api" "github.com/tidbcloud/dataflow-service/pkg/s3" infra_config "github.com/tidbcloud/infra-api/pkg/pingcap/util/config" "github.com/tidbcloud/pkgv2/log" "github.com/tidbcloud/tiinf/sdk/auditlog" "go.uber.org/zap" "google.golang.org/protobuf/types/known/emptypb" "gorm.io/gorm" ) const expectedPreparationSeconds = 45 type importService struct { dataflow.UnimplementedImportServiceServer service.OperationsService importModel pimport.Model clustersModel cluster.Model } // NewService constructs a importService func NewService(s service.OperationsService, model pimport.Model, clustersModel cluster.Model) dataflow.ImportServiceServer { return &importService{OperationsService: s, importModel: model, clustersModel: clustersModel} } func (s *importService) userImport(ctx context.Context) pimport.UserContext { innerAuth := ctx.Value(utils.CtxInnerCallerKey) if innerAuth != nil { if user, ok := innerAuth.(*accountdb.User); ok { return s.importModel.User(ctx, user) } } user := s.CurrentUser(ctx) return s.importModel.User(ctx, user) } func (s *importService) Log(ctx context.Context) *zap.Logger { return log.AppendCtxFields(ctx, s.userImport(ctx).Log()) } func (s *importService) Create(ctx context.Context, req *dataflow.ImportCreateReq) (_ *dataflow.ClusterImportId, rErr error) { defer func() { s.recordAuditLogForImport(ctx, req, rErr) }() logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) orgID := entity.ID(req.OrgId) projectID := entity.ID(req.ProjectId) userCtx := s.userImport(ctx) if rErr := bizutils.CheckClusterStatus(ctx, req.ClusterId); rErr != nil { logger.Warn("check cluster status error", zap.Error(rErr)) return nil, wrapNonErrorCodeAsInternalError(rErr, "") } // impersonated user cannot do post request, see GetUserClaim // todo: we should remove check about impersonated user for POST request rErr = userCtx.CheckOrgIDAndProjectID(orgID, projectID) if rErr != nil { return nil, rErr } sk := req.GetAwsSecretAccessKey() req.AwsSecretAccessKey = "REDACTED" logger.Info("create import", zap.Reflect("request", req)) req.AwsSecretAccessKey = sk rErr = validateReq(req, logger) if rErr != nil { return nil, errcode.NewInvalidInputErr(rErr) } imp, rErr := buildImport(ctx, req) if rErr != nil { logger.Error("build import error", zap.Error(rErr)) return nil, wrapNonErrorCodeAsInternalError(rErr, "") } if req.LocalImportCfg != nil { // todo: merge into AdjustImport if err := userCtx.AdjustLocalImport(ctx, imp, req); err != nil { logger.Error("Adjust local import error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } rErr := userCtx.CreateTargetTable(ctx, req.OrgId, req.ProjectId, req.ClusterId, req.LocalImportCfg) if rErr != nil { return nil, wrapNonErrorCodeAsInternalError(rErr, "Internal error, failed to pre-create target table, please retry later or contact support.") } } id, rErr := userCtx.CreateImport(ctx, orgID, projectID, imp) if rErr != nil { errMsg := "Internal error when creating import, please retry later or feel free to contact support." logger.Warn(errMsg, zap.Reflect("req", req), zap.Reflect("import", imp), zap.Error(rErr)) return nil, wrapNonErrorCodeAsInternalError(rErr, errMsg) } resp := &dataflow.ClusterImportId{ Id: id, ClusterId: req.ClusterId, OrgId: req.OrgId, ProjectId: req.ProjectId, } logger.Info("create import success", zap.Reflect("response", resp)) return resp, nil } func (s *importService) recordAuditLogForImport(ctx context.Context, req *dataflow.ImportCreateReq, rErr error) { logger := log.WithContextS(ctx) cluster, bizError := grpcli.GetClusterByID(ctx, req.ClusterId) if bizError != nil { logger.Warn("failed to submit audit log, query cluster error", zap.Uint64("clusterId", req.ClusterId), zap.Error(bizError)) return } details := auditlog.Details{ auditlog.KeyBucketUri: req.SourceUrl, auditlog.KeyRoleArn: req.AwsRoleArn, auditlog.KeyTargetClusterName: cluster.Name, auditlog.KeyTargetRegion: cluster.RegionName, auditlog.KeyError: rErr, } switch dataflow.ImportDataFormat(req.DataFormat) { case dataflow.ImportDataFormat_SQL: details[auditlog.KeySourceDataSource] = "SQL" case dataflow.ImportDataFormat_AURORA_SNAPSHOT: details[auditlog.KeySourceDataSource] = "Parquet" // TODO: change to Aurora case dataflow.ImportDataFormat_CSV: details[auditlog.KeySourceDataSource] = "CustomCSV" case dataflow.ImportDataFormat_PARQUET: details[auditlog.KeySourceDataSource] = "RawParquet" } var eventType auditlog.EventType switch dataflow.ImportDataSource(req.DataSourceType) { case dataflow.ImportDataSource_S3_SOURCE: details[auditlog.KeySourceDataSource] = "AWS" eventType = auditlog.EventImportDataFromAWS case dataflow.ImportDataSource_GCS_SOURCE: details[auditlog.KeySourceDataSource] = "GCP" eventType = auditlog.EventImportDataFromGCP case dataflow.ImportDataSource_AZBLOB_SOURCE: details[auditlog.KeySourceDataSource] = "Azure" eventType = auditlog.EventImportDataFromAzure default: eventType = auditlog.EventImportDataFromLocal } err := auditlog.NewEvent(ctx, eventType). SetResult(rErr == nil). SetOrgID(req.OrgId). SetProjectID(req.ProjectId). SetDetails(details.Build()). Submit() if err != nil { logger.Warn("failed to submit audit log", zap.Error(err)) return } } func getRealValidationItems(originalItems []dataflow.ImportValidationItem) []dataflow.ImportValidationItem { checkItems := []dataflow.ImportValidationItem{} for _, item := range originalItems { switch item { case dataflow.ImportValidationItem_SOURCE_ACCESS_CHECK, dataflow.ImportValidationItem_IMPORT_PRECHECK, dataflow.ImportValidationItem_TARGET_TABLE_EMPTY_CHECK, dataflow.ImportValidationItem_ESTIMATE_SIZE, dataflow.ImportValidationItem_IMPORT_DATA_ONLY_CHECK: checkItems = append(checkItems, item) } } return checkItems } // POST "/api/v1/orgs/{org_id}/projects/{project_id}/clusters/{cluster_id}/imports/validate" func (s *importService) Validate(ctx context.Context, req *dataflow.ImportValidateReq) (*emptypb.Empty, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Error("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } logger.Info("validate import", zap.Reflect("request", req)) // trim spaces for some input fields req.SourceUrl = strings.TrimSpace(req.SourceUrl) req.AwsRoleArn = strings.TrimSpace(req.AwsRoleArn) // add default value when CustomCSVFormat is nill if dataflow.ImportDataFormat(req.DataFormat) == dataflow.ImportDataFormat_CSV && req.CsvFormat == nil { req.CsvFormat = &dataflow.CustomCSVFormat{ Separator: ",", Delimiter: "\"", Header: true, BackslashEscape: true, } } if err := checkValidationRequest(ctx, req, logger); err != nil { errMsg := "Invalid field" return nil, errcode.NewInvalidInputErr(errors.Annotate(err, errMsg)) } if err := s.mergeLightningTemplConfig(ctx, req); err != nil { return nil, utils.WrapAsErrorCodeErr(err, "merge lightning template config error") } imp, err := buildImportForValidation(ctx, req) if err != nil { logger.Error("build import for validation error", zap.Error(err)) return nil, utils.WrapAsErrorCodeErr(err, "prepare validation error") } logger.Info("validate import[buildImportForValidation]", zap.Reflect("import", imp)) checkItems := getRealValidationItems(req.GetCheckItems()) if err := userCtx.ValidateImport(ctx, entity.ID(req.GetOrgId()), entity.ID(req.GetProjectId()), imp, checkItems); err != nil { translatedErr := TranslateValidateImportError(err, imp, logger) // There are some situations that after the translation, the result error is nil. // This helps to ignore some kind of error and do the rest of work. if translatedErr != nil { return nil, translatedErr } } return new(emptypb.Empty), nil } // mergeLightningTemplConfig merges the Lightning configuration template into the import validation request. // It retrieves the Lightning configuration template based on the provided cluster ID and updates the CSV format // settings in the request if they are present in the template. // // Parameters: // - ctx: The context for the operation. // - req: The import validation request containing the cluster ID and CSV format settings. // // Returns: // - error: An error if the Lightning configuration template cannot be retrieved or processed. func (s *importService) mergeLightningTemplConfig(ctx context.Context, req *dataflow.ImportValidateReq) error { if dataflow.ImportDataFormat(req.DataFormat) != dataflow.ImportDataFormat_CSV || req.CsvFormat == nil { return nil } userCtx := s.userImport(ctx) lightningConfig, err := userCtx.GetLightningConfigTemplate(ctx, req.ClusterId) if err != nil { return err } if lightningConfig == nil { return nil } impConfig := infra_config.New(lightningConfig.Config.MP) if value := impConfig.Get("mydumper.csv.trim-last-separator").Interface(); value != nil { if trimLastSep, ok := value.(bool); ok { req.CsvFormat.TrimLastSeparator = &trimLastSep } } if value := impConfig.Get("mydumper.csv.not-null").Interface(); value != nil { if notNull, ok := value.(bool); ok { req.CsvFormat.NotNull = ¬Null } } if value := impConfig.Get("mydumper.csv.null").Interface(); value != nil { if csvNull, ok := value.(string); ok { req.CsvFormat.Null = &csvNull } } return nil } func (s *importService) ValidateTargetTableEmpty(ctx context.Context, req *dataflow.ImportsValidateTargetTableEmptyReq) (*emptypb.Empty, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } fullTableName := req.GetTargetTable() isTableEmpty, err := userCtx.CheckIsTableEmpty(ctx, req.GetClusterId(), pimport.FullTableName{ DatabaseName: fullTableName.GetSchema(), TableName: fullTableName.GetTable(), }) if err != nil { logger.Error("check table is empty from agent has error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "Checking table is empty has errors.") } if !isTableEmpty { return nil, errcode.NewCodedError( errors.Errorf("The target table %s.%s is not empty", fullTableName.GetSchema(), fullTableName.GetTable()), errcode.StateCode, ) } return new(emptypb.Empty), nil } func (s *importService) Get(ctx context.Context, in *dataflow.ClusterImportId) (*dataflow.ImportGetResp, error) { logger := s.Log(ctx) logger.Info("get import", zap.Reflect("request", in)) orgID := entity.ID(in.GetOrgId()) projectID := entity.ID(in.GetProjectId()) clusterID := entity.ID(in.GetClusterId()) importID := entity.ID(in.GetId()) userCtx := s.userImport(ctx) imp, err := userCtx.GetImport(ctx, orgID, projectID, clusterID, importID) if err != nil { logger.Error("failed to get import", zap.Uint64("cluster id", in.ClusterId), zap.Uint64("import id", in.Id), zap.Error(err)) if errors.Cause(err) == gorm.ErrRecordNotFound { return nil, errcode.NewNotFoundErr(err) } return nil, wrapNonErrorCodeAsInternalError(err, "") } resp := ImportEntity2PB(ctx, logger, in.ClusterId, imp, true) if imp.LocalImportCfg != nil { return resp, nil } // if the importing is on going, fetch the imformation directly from import status from infra-API, // because some information are not persisted in meta-DB, but could be got from the infra-API's import status // also, at the importing time, the table status for each table can be got from the infra-API's import status, // but they haven't been persisted into the import_tables table ctx = common.ContextWithImport(ctx) if imp.Status == entity.ImportStatusImporting { resourceAPI, err := pimport.NewRegionalResourceAPI( ctx, imp.Cluster.GetInfraProviderType(), imp.Cluster.GetRegionName(), imp.Cluster.GetTenantId()) if err != nil { logger.Error("failed to get regional resource api", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } importCR, err := resourceAPI.GetImport(ctx, in.ProjectId, uint64(imp.ID)) if err != nil { logger.Error("get import CR from infra API error", zap.Error(err)) resp.Message += "| Failed to get import progress from infra API. " } else { saveImportStatusIntoPBResp(&(importCR.Status), resp) } } tweakImportProgress(imp, resp) return resp, nil } // tweakImportProgress tweaks the overall import progress. // In some situations, include the preparing phase as some percent func tweakImportProgress(imp *pimport.ImportTask, resp *dataflow.ImportGetResp) { // not check on import status, because for a canceling/canceled/failed import, // we cannot get whether the task is halted at prepare phase or import phase. // So here check according to the start at timestamp as well as the completed import percent isInPreparation := (imp.StartedAt == nil && imp.CompletedPercent == 0) // TODO: tweak the max estimated data size. Currently we set it to 1TB // This is because we want to set the value large enough, that it exceeds most of the actual data size. // It can ensure that when the actual data size is got, the progress still goes forward rather than backward // However, setting it too big will cause: // 1) the estimation progress at the prepare phase moves very slowly, or stops at the preparation percent threshold for quite a long time // 2) after the actual data size is got, if the actual size is far less than the max size, the progress bar moves forward dramatically resp.CompletedPercent = uint32(getOverallImportProgress(isInPreparation, resp, 10, 1*units.TiB)) } // getOverallImportProgress get the overall import progress. // It will also include the preparation phase as some progress func getOverallImportProgress(isInPreparation bool, resp *dataflow.ImportGetResp, maxPreparationPercent float64, maxEstimatedDataSize uint64) float64 { var overallPercent float64 // first, get an overall percent based on importing completed percent and elapsed time (used before importing) if isInPreparation { // we cannot know the total file size in advance when preparing ! // make some elapsed seconds increase some percents overallPercent = float64(resp.GetElapsedTimeSeconds()) / float64(expectedPreparationSeconds) * maxPreparationPercent if overallPercent > maxPreparationPercent { overallPercent = maxPreparationPercent } } else { overallPercent = maxPreparationPercent + (100-maxPreparationPercent)/100*float64(resp.GetCompletedPercent()) if overallPercent > 100 { overallPercent = 100 } } // second, get an overall percent estimation based on the elapsed time and overall data size // this estimation can make sure the progress always moves, because elapsed time constantly changes. // If it is larger than the overall percent above, use this estimation as the progress // This makes the progress not stop at a certain point for a long time var dataSize uint64 if resp.TotalSize > 0 { dataSize = resp.TotalSize } else { dataSize = maxEstimatedDataSize } estimatedPercent := estimateImportProgressByElapsedTime(resp.GetElapsedTimeSeconds(), dataSize) if estimatedPercent > overallPercent { overallPercent = estimatedPercent } return overallPercent } // estimateImportProgressByElapsedTime estimates import progress by elapsed time and file size. // It is a very tricky one: at the very beginning when we cannot get the actual file size, // we need to set estimated file size to a large enough value, so that the calculated percent is a small one. // After we get the actual data size, we can adjust the estimated file size, to let the progress move forward. func estimateImportProgressByElapsedTime(elapsedSeconds uint32, estimatedFileSize uint64) float64 { // estimate import speed is 350 KiB / s // Set it to a slow speed, so that in most cases the actual progress will be shown estimatedImportSeconds := float64(estimatedFileSize) / (350 * units.KiB) result := float64(elapsedSeconds) / (expectedPreparationSeconds + estimatedImportSeconds) * 100 if result > 99 { result = 99 } return result } func (s *importService) List(ctx context.Context, req *dataflow.ListImportReq) (*dataflow.ListImportResp, error) { userCtx := s.userImport(ctx) logger := s.Log(ctx) logger.Info("list import", zap.Reflect("req", req)) pager := utils.AdjustedPager(&req.Page, &req.PageSize) logger.Info("adjusted pager", zap.Reflect("pager", pager)) statuses := pbStatusList2ImportStatusList(req.GetStatuses()) orgID := entity.ID(req.GetOrgId()) projectID := entity.ID(req.GetProjectId()) clusterID := entity.ID(req.GetClusterId()) totalCount, importEntities, err := userCtx.ListImport( ctx, orgID, projectID, clusterID, pager, statuses, ) if err != nil { errMsg := "list import metadatas failed" logger.Error(errMsg, zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, errMsg) } return importEntitiesToListRespPB(ctx, clusterID, totalCount, importEntities, logger) } func (s *importService) Cancel(ctx context.Context, req *dataflow.ClusterImportId) (*emptypb.Empty, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Error("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } logger.Info("cancel import", zap.Reflect("request", req)) if err := userCtx.CancelImport(ctx, entity.ID(req.GetProjectId()), entity.ID(req.GetClusterId()), entity.ID(req.GetId()), false, /* waitCanceled */ ); err != nil { errMsg := "failed to cancel import" logger.Warn(errMsg, zap.Reflect("request", req), zap.Error(err)) switch realErr := err.(type) { case *pimport.CancelImportInvalidStatusError: return nil, errcode.NewCodedError(realErr, errcode.StateCode) default: return nil, wrapNonErrorCodeAsInternalError(err, errMsg) } } logger.Info("cancel import success") return &emptypb.Empty{}, nil } func (s *importService) GetS3BucketRegion(ctx context.Context, req *dataflow.GetS3BucketRegionReq) (*dataflow.GetS3BucketRegionResp, error) { logger := s.Log(ctx) u, err := s3.NewS3URL(req.GetS3Url()) if err != nil { logger.Error("construct a S3 URL object error", zap.Error(err), zap.String("s3_url", req.GetS3Url())) return nil, errcode.NewInvalidInputErr(errors.Annotatef(err, "cannot parse S3 URL '%s'", req.GetS3Url())) } region, err := u.GetBucketRegion(ctx) if err != nil { errMsg := "get bucket region error" logger.Error(errMsg, zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, errMsg) } return &dataflow.GetS3BucketRegionResp{Region: region}, nil } func (s *importService) Exit(ctx context.Context, req *dataflow.ClusterImportId) (*emptypb.Empty, error) { // To avoid the error when frontend call this api yet, we didn't delete this api. return &emptypb.Empty{}, nil } func (s *importService) userCluster(ctx context.Context) cluster.UserContext { user := s.CurrentUser(ctx) return s.clustersModel.User(ctx, user) } // todo remove cluster model // GET /api/v1/orgs/${orgld}/projects/${projectld}/clusters/${clusterld}/data-ingestion-profile func (s *importService) GetDataIngestionProfile(ctx context.Context, req *dataflow.ClusterIdReq) (*dataflow.DataIngestionProfile, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } logger.Info("GetDataIngestionProfile", zap.Reflect("request", req)) profile, err := s.userCluster(ctx).GetDataIngestionProfile(ctx, entity.ID(req.ClusterId)) if err != nil { errMsg := "get data ingestion profile error" logger.Error(errMsg, zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, errMsg) } return &dataflow.DataIngestionProfile{ Type: profile.Type, AwsExternalId: profile.AwsExternalId, AwsAccountId: profile.AwsAccountId, GcpAccountId: profile.GcpAccountId, }, nil } func (s *importService) GetSchemasMeta(ctx context.Context, req *dataflow.ImportsGetSchemasMetaReq) (*dataflow.ImportsGetSchemasMetaResp, error) { return nil, errcode.NewNotFoundErr(errors.New("get schemas meta is not implemented")) } func (s *importService) Preview(ctx context.Context, req *dataflow.ImportsPreviewReq) (*dataflow.ImportsPreviewResp, error) { return nil, errcode.NewNotFoundErr(errors.New("preview is not implemented")) } // api: "/api/v1/orgs/{org_id}/projects/{project_id}/clusters/{cluster_id}/import-target-tables" func (s *importService) GetClusterTables(ctx context.Context, req *dataflow.ImportsGetClusterTablesReq) (*dataflow.ImportsGetClusterTablesResp, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } password, err := userCtx.GenCloudAdminPassword(ctx, logger, req.OrgId, req.ProjectId, req.ClusterId) if err != nil { return nil, wrapNonErrorCodeAsInternalError(err, "Internal error, please contact support.") } ctx = common.ContextWithImport(ctx) cluster, _, agentClient, err := common.GetClusterNetworkAgentById(ctx, req.ClusterId) if err != nil { errMsg := "get cluster, network, agent client failed" logger.Error(errMsg, zap.Error(err)) return nil, utils.WrapAsErrorCodeErr(err, errMsg) } tcNameGetter := &infraapi.ClusterNameGetter{ InfraProvider: cluster.InfraProviderType, InfraAPIObjNamespace: infraapi.GetObjectNamespace(req.ProjectId), ClusterName: infraapi.ClusterName(cluster.Id), } agentReq := &agentpb.Cluster{ ClusterId: req.ClusterId, Namespace: tcNameGetter.ClusterNamespace(), User: entity.CloudAdminUser, Password: password, ProviderType: cluster.GetInfraProviderType(), } agentClient.SetHttpTimeout(50 * time.Second) agentResp, err := agentClient.GetClusterTables(ctx, agentReq) if err != nil { return nil, errors.Annotatef(err, "failed to execute sql on cluster") } resp := &dataflow.ImportsGetClusterTablesResp{ Schemas: agentResp.GetSchemas(), } return resp, nil } const ( MaxUploadSize = 50 * units.MiB DefaultRegion = "us-east-1" ) func (s *importService) GenerateUploadUrl(ctx context.Context, req *dataflow.ImportsGenerateUploadUrlReq) (*dataflow.ImportsGenerateUploadUrlResp, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } if req.ContentLength > MaxUploadSize { logger.Warn("file size is bigger than 50MiB", zap.Int64("content_length(MiB)", req.ContentLength/units.MiB)) return nil, errcode.NewInvalidInputErr(errors.New("Your file size is bigger than 50MiB")) } region := DefaultRegion if req.Region != nil { region = *req.Region } else { logger.Info("the region is empty in request, so use the default region(us-eas-1)") } // todo: check region valid svc, err := s3.CreateClientWithAssumeRole(region, "", "") if err != nil { logger.Error("failed to create client with assume role", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "Internal error, please contact support.") } now := time.Now() fileName := fmt.Sprintf("%s-%s", now.Format(time.RFC3339), filepath.Base(req.FileName)) bucketNaming := pimport.NewLocalImportBucketNaming(region, fileName, userCtx.UserID()) objReq, _ := svc.PutObjectRequest(&awss3.PutObjectInput{ Bucket: aws.String(bucketNaming.GenBucketName()), Key: aws.String(bucketNaming.GenObjectKey()), ContentLength: aws.Int64(req.ContentLength), }) // the front end will use the url immediately after fetching it. url, err := objReq.Presign(3 * time.Minute) if err != nil { logger.Error("failed to presign url", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "Internal error, please contact support.") } return &dataflow.ImportsGenerateUploadUrlResp{ Url: url, FileName: fileName, }, nil } func (s *importService) ListDataTypes(ctx context.Context, req *emptypb.Empty) (*dataflow.ImportsListDataTypesResp, error) { resp := &dataflow.ImportsListDataTypesResp{ DataTypes: pimport.GetDefaultDataType(), } return resp, nil } func (s *importService) PreviewLocal(ctx context.Context, req *dataflow.ImportsPreviewLocalReq) (*dataflow.ImportsPreviewLocalResp, error) { logger := s.Log(ctx) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } if req.DataFormat == dataflow.ImportDataFormat_CSV { if err := validateCustomCSV(req.CsvFormat); err != nil { return nil, errcode.NewInvalidInputErr(err) } } adjustCSVConfig(ctx, req.CsvFormat) firstN := req.N if firstN == 0 { // default return first 10 firstN = 10 } bucketNaming := pimport.NewLocalImportBucketNaming(req.GetRegion(), req.GetFileName(), userCtx.UserID()) var previewer pimport.ImportPreviewer = &pimport.LocalImportPreviewer{ CSVFormat: req.GetCsvFormat(), SourceURL: bucketNaming.GenBucketUrl(), DataFormat: pbDataFormat2Entity(req.DataFormat), PreviewFileMappings: []*pimport.PreviewFileMapping{ { DatabaseName: "target_db", // this is a fake name. TableName: "target_table", // this is a fake name. FileName: req.GetFileName(), }, }, } dbPreviewResults, err := previewer.Preview(ctx, uint(firstN)) if err != nil { logger.Error("preview table data has errors", zap.Error(err)) switch concreteErr := errors.Cause(err).(type) { case *pimport.ErrCSVSyntax: return nil, errcode.NewInvalidInputErr(errors.New("There is a syntax error. Please edit the 'CSV Configuration' to enable the program to parse and preview the source file content correctly.")) case *pimport.ErrCSVRowOverflow: return nil, errcode.NewInvalidInputErr(errors.Errorf("The CSV file is invalid, your %s row data has overflowed.", humanize.Ordinal(concreteErr.RowNum+1))) case *pimport.ErrHeaderDataColumnMismatch: return nil, errcode.NewInvalidInputErr(errors.Errorf("Invalid CSV file. Your row data has %d columns, but your headers have %d columns. Please modify the CSV files to ensure them match.", concreteErr.DataColumnsCount, concreteErr.HeaderColumnsCount)) default: return nil, wrapNonErrorCodeAsInternalError(err, "Internal error when previewing table data. Please seek support.") } } if len(dbPreviewResults) < 1 { return nil, errcode.NewCodedError(errors.New("no preview result for the database"), errcode.StateCode) } dbPreviewResult := dbPreviewResults[0] if len(dbPreviewResult.GetTables()) < 1 { return nil, errcode.NewCodedError(errors.New("no preview result for the table"), errcode.StateCode) } return &dataflow.ImportsPreviewLocalResp{ Table: dbPreviewResult.GetTables()[0], }, nil } func (s *importService) GetTableColumns(ctx context.Context, req *dataflow.ImportsGetTableColumnsReq) (*dataflow.ImportsGetTableColumnsResp, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } table, err := userCtx.GetTableColumns(ctx, req) if err != nil { return nil, wrapNonErrorCodeAsInternalError(err, "Internal error when getting table columns. Please contact support.") } return &dataflow.ImportsGetTableColumnsResp{ Table: table, }, nil } // POST /api/v1/orgs/${orgld}/projects/${projectld}/clusters/${clusterld}/imports/search-import-files func (s *importService) SearchImportFiles(ctx context.Context, req *dataflow.ImportsSearchImportFilesReq) (*dataflow.ImportsSearchImportFilesResp, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { errMsg := "check cluster ID and cluster status error" logger.Error(errMsg, zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, errMsg) } if err := validateCustomFile(req.SourceUri, req.TargetTableInfos); err != nil { return nil, errcode.NewInvalidInputErr(err) } resp, err := userCtx.SearchImportFiles(ctx, req) if err != nil { logger.Error("failed to SearchImportFiles", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "Internal error when validate target source match. Please contact support.") } return resp, nil } func (s *importService) PreviewCsvFile(ctx context.Context, req *dataflow.ImportsPreviewCsvFileReq) (*dataflow.ImportsPreviewCsvFileResp, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { errMsg := "check cluster ID and cluster status error" logger.Error(errMsg, zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, errMsg) } resp, err := userCtx.PreviewCsvFile(ctx, req) if err != nil { return nil, wrapNonErrorCodeAsInternalError(err, "Internal error when previewing file, please contact support.") } return resp, nil } func (s *importService) Precheck(ctx context.Context, req *dataflow.ImportPrecheckReq) (*dataflow.ImportPrecheckResp, error) { logger := s.Log(ctx). With(zap.Uint64("project_id", req.GetProjectId())). With(zap.Uint64("cluster_id", req.GetClusterId())) userCtx := s.userImport(ctx) if err := userCtx.CheckIdAndCluster(ctx, req.OrgId, req.ProjectId, req.ClusterId); err != nil { logger.Warn("check cluster ID and cluster status error", zap.Error(err)) return nil, wrapNonErrorCodeAsInternalError(err, "") } password, err := userCtx.GenCloudAdminPassword(ctx, logger, req.OrgId, req.ProjectId, req.ClusterId) if err != nil { return nil, wrapNonErrorCodeAsInternalError(err, "Internal error, please contact support.") } ctx = common.ContextWithImport(ctx) cluster, _, agentClient, err := common.GetClusterNetworkAgentById(ctx, req.ClusterId) if err != nil { errMsg := "get cluster, network, agent client failed" logger.Error(errMsg, zap.Error(err)) return nil, utils.WrapAsErrorCodeErr(err, errMsg) } tcNameGetter := &infraapi.ClusterNameGetter{ InfraProvider: cluster.InfraProviderType, InfraAPIObjNamespace: infraapi.GetObjectNamespace(req.ProjectId), ClusterName: infraapi.ClusterName(cluster.Id), } // Build agent request with auto-filled target cluster info agentReq := &agentpb.PrecheckImportJobReq{ OrgId: strconv.FormatUint(req.OrgId, 10), ProjectId: strconv.FormatUint(req.ProjectId, 10), ClusterId: strconv.FormatUint(req.ClusterId, 10), ImportOptions: req.ImportOptions, Source: req.Source, TargetCluster: &agentpb.Cluster{ ClusterId: req.ClusterId, Namespace: tcNameGetter.ClusterNamespace(), User: entity.CloudAdminUser, Password: password, ProviderType: cluster.GetInfraProviderType(), }, } agentResp, err := agentClient.PrecheckImportJob(ctx, agentReq) if err != nil { logger.Error("agent precheck failed", zap.Error(err)) return nil, utils.WrapAsErrorCodeErr(err, "precheck failed") } resp := &dataflow.ImportPrecheckResp{ TargetTables: agentResp.GetTargetTables(), ErrorMessage: agentResp.GetErrorMessage(), TotalDataFilesCount: agentResp.GetTotalDataFilesCount(), TotalTablesCount: agentResp.GetTotalTablesCount(), IsTruncated: agentResp.GetIsTruncated(), Status: agentResp.GetStatus(), } return resp, nil }