Refacto framework exports

Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>

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(
+	fe, scope, err := s.lockExport(ctx)
-		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
+				return fmt.Errorf("cannot generate uuid: %w", err)
-				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)
 			}
 
 			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
+		},
+	)
 }