diff --git a/pkg/probo/service.go b/pkg/probo/service.go index 903e17dd1..0ddf566dc 100644 --- a/pkg/probo/service.go +++ b/pkg/probo/service.go @@ -192,15 +192,38 @@ func (s *Service) WithTenant(tenantID gid.TenantID) *TenantService { } func (s *Service) ExportFrameworkJob(ctx context.Context) error { - return s.pg.WithTx( - ctx, + fe, scope, err := s.lockExport(ctx) + if err != nil { + return fmt.Errorf("cannot lock framework export: %w", err) + } + + fe, buildErr := s.buildAndUploadExport(ctx, scope, fe) + if buildErr != nil { + if err := s.commitFailedExport(ctx, scope, fe); err != nil { + return fmt.Errorf( + "cannot build and upload framework export: %w, and cannot commit failed export: %w", + buildErr, + err, + ) + } + + return fmt.Errorf("cannot build and upload framework export: %w", buildErr) + } + + return nil +} + +func (s *Service) lockExport(ctx context.Context) (*coredata.FrameworkExport, coredata.Scoper, error) { + fe := &coredata.FrameworkExport{} + var scope coredata.Scoper + + err := s.pg.WithTx(ctx, func(tx pg.Conn) error { - fe := &coredata.FrameworkExport{} if err := fe.LoadNextPendingForUpdateSkipLocked(ctx, tx); err != nil { - return err + return fmt.Errorf("cannot load next pending framework export: %w", err) } - scope := coredata.NewScope(fe.ID.TenantID()) + scope = coredata.NewScope(fe.ID.TenantID()) fe.Status = coredata.FrameworkExportStatusProcessing fe.StartedAt = ref.Ref(time.Now()) @@ -208,73 +231,50 @@ func (s *Service) ExportFrameworkJob(ctx context.Context) error { return fmt.Errorf("cannot update framework export: %w", err) } + return nil + }, + ) + if err != nil { + return nil, nil, fmt.Errorf("cannot lock framework export: %w", err) + } + + return fe, scope, nil +} + +func (s *Service) buildAndUploadExport(ctx context.Context, scope coredata.Scoper, fe *coredata.FrameworkExport) (*coredata.FrameworkExport, error) { + err := s.pg.WithTx( + ctx, + func(tx pg.Conn) error { framework := &coredata.Framework{} if err := framework.LoadByID(ctx, tx, scope, fe.FrameworkID); err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - return fmt.Errorf("cannot load framework: %w", err) } - tenantService := s.WithTenant(fe.ID.TenantID()) - tempDir := os.TempDir() tempFile, err := os.CreateTemp(tempDir, "probo-framework-export-*.zip") if err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - return fmt.Errorf("cannot create temp file: %w", err) } defer tempFile.Close() defer os.Remove(tempFile.Name()) + tenantService := s.WithTenant(scope.GetTenantID()) err = tenantService.Frameworks.Export(ctx, fe.FrameworkID, tempFile) if err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - return fmt.Errorf("cannot export framework: %w", err) } uuid, err := uuid.NewV4() if err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - - return fmt.Errorf("cannot update framework export: %w", err) + return fmt.Errorf("cannot generate uuid: %w", err) } if _, err := tempFile.Seek(0, 0); err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - return fmt.Errorf("cannot seek temp file: %w", err) } fileInfo, err := tempFile.Stat() if err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - return fmt.Errorf("cannot stat temp file: %w", err) } @@ -293,12 +293,6 @@ func (s *Service) ExportFrameworkJob(ctx context.Context) error { }, ) if err != nil { - fe.Status = coredata.FrameworkExportStatusFailed - fe.CompletedAt = ref.Ref(time.Now()) - if err := fe.Update(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update framework export: %w", err) - } - return fmt.Errorf("cannot upload file to S3: %w", err) } @@ -308,7 +302,7 @@ func (s *Service) ExportFrameworkJob(ctx context.Context) error { ID: gid.New(fe.ID.TenantID(), coredata.FileEntityType), BucketName: s.bucket, MimeType: "application/zip", - FileName: fmt.Sprintf("%s Archive %s.zip", framework.Name, time.Now().Format("2006-01-02")), + FileName: fmt.Sprintf("%s Archive %s.zip", framework.Name, now.Format("2006-01-02")), FileKey: uuid.String(), FileSize: int(fileInfo.Size()), CreatedAt: now, @@ -320,7 +314,7 @@ func (s *Service) ExportFrameworkJob(ctx context.Context) error { } fe.FileID = &file.ID - fe.CompletedAt = &now + fe.CompletedAt = ref.Ref(time.Now()) fe.Status = coredata.FrameworkExportStatusCompleted if err := fe.Update(ctx, tx, scope); err != nil { return fmt.Errorf("cannot update framework export: %w", err) @@ -329,4 +323,26 @@ func (s *Service) ExportFrameworkJob(ctx context.Context) error { return nil }, ) + + if err != nil { + return fe, fmt.Errorf("cannot build and upload export: %w", err) + } + + return fe, nil +} + +func (s *Service) commitFailedExport(ctx context.Context, scope coredata.Scoper, fe *coredata.FrameworkExport) error { + fe.CompletedAt = ref.Ref(time.Now()) + fe.Status = coredata.FrameworkExportStatusFailed + + return s.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := fe.Update(ctx, conn, scope); err != nil { + return fmt.Errorf("cannot update framework export: %w", err) + } + + return nil + }, + ) }