HMS-5571: add upload cleanup (#1087)

Author: rverdile

cmd/external-repos/main.go

@@ -103,6 +103,11 @@ func main() {
 		} else {
 			log.Warn().Msg("Snapshotting disabled")
 		}
+	} else if args[1] == "upload-cleanup" {
+		err = uploadCleanup(ctx, db.DB)
+		if err != nil {
+			log.Error().Err(err).Msgf("error starting upload cleanup tasks")
+		}
 	} else if args[1] == "process-repos" {
 		err = enqueueIntrospectAllRepos(ctx)
 		if err != nil {
@@ -128,6 +133,10 @@ func main() {
 				log.Error().Err(err).Msg("error queueing delete snapshot tasks for snapshot cleanup")
 			}
 		}
+		err = uploadCleanup(ctx, db.DB)
+		if err != nil {
+			log.Error().Err(err).Msgf("error starting upload cleanup tasks")
+		}
 	} else if args[1] == "pulp-orphan-cleanup" {
 		batchSize := 5
 		if len(args) > 2 {
@@ -140,7 +149,7 @@ func main() {
 		if !config.PulpConfigured() {
 			log.Error().Msg("cannot run orphan cleanup if pulp is not configured")
 		}
-		err := pulpOrphanCleanup(ctx, db.DB, batchSize)
+		err = pulpOrphanCleanup(ctx, db.DB, batchSize)
 		if err != nil {
 			log.Error().Err(err).Msg("error starting pulp orphan cleanup tasks")
 		}
@@ -520,3 +529,44 @@ func pulpOrphanCleanup(ctx context.Context, db *gorm.DB, batchSize int) error {
 	}
 	return nil
 }
+
+func uploadCleanup(ctx context.Context, db *gorm.DB) error {
+	var err error
+	daoReg := dao.GetDaoRegistry(db)
+
+	log.Info().Msg("===== Starting upload cleanup =====")
+
+	uploads, err := daoReg.Uploads.ListUploadsForCleanup(ctx)
+	if err != nil {
+		return fmt.Errorf("error listing uploads for cleanup: %w", err)
+	}
+
+	var cleanupCounter int
+	for _, upload := range uploads {
+		orgID := upload.OrgID
+
+		domainName, err := daoReg.Domain.Fetch(ctx, orgID)
+		if err != nil {
+			log.Error().Err(err).Msgf("error fetching domain name for org %v", orgID)
+			continue
+		}
+		logger := log.Logger.With().Str("org_id", orgID).Str("pulp_domain_name", domainName).Logger()
+		pulpClient := pulp_client.GetPulpClientWithDomain(domainName)
+
+		uploadHref := "/api/pulp/" + domainName + "/api/v3/uploads/" + upload.UploadUUID + "/"
+		_, err = pulpClient.DeleteUpload(ctx, uploadHref)
+		if err != nil {
+			logger.Error().Err(err).Msgf("error deleting pulp upload with uuid: %v", upload.UploadUUID)
+			continue
+		}
+
+		err = daoReg.Uploads.DeleteUpload(ctx, upload.UploadUUID)
+		if err != nil {
+			logger.Error().Err(err).Msgf("error deleting upload")
+			continue
+		}
+		cleanupCounter++
+	}
+	log.Info().Msgf("Cleaned up %v uploads", cleanupCounter)
+	return nil
+}

pkg/clients/pulp_client/interfaces.go

@@ -95,4 +95,5 @@ type PulpClient interface {
 	CreateUpload(ctx context.Context, size int64) (*zest.UploadResponse, int, error)
 	UploadChunk(ctx context.Context, uploadHref string, contentRange string, file *os.File, sha256 string) (*zest.UploadResponse, int, error)
 	FinishUpload(ctx context.Context, uploadHref string, sha256 string) (*zest.AsyncOperationResponse, int, error)
+	DeleteUpload(ctx context.Context, uploadHref string) (int, error)
 }

pkg/clients/pulp_client/pulp_client_mock.go

@@ -2,10 +2,13 @@ package pulp_client
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"os"
+	"strings"
 
 	zest "github.com/content-services/zest/release/v2025"
+	"github.com/rs/zerolog/log"
 )
 
 // CreateUpload Creates an upload
@@ -63,3 +66,29 @@ func (r *pulpDaoImpl) FinishUpload(ctx context.Context, uploadHref string, sha25
 	}
 	return readResp, statusCode, nil
 }
+
+func (r *pulpDaoImpl) DeleteUpload(ctx context.Context, uploadHref string) (int, error) {
+	ctx, client := getZestClient(ctx)
+	statusCode := http.StatusInternalServerError
+	var body []byte
+	var readErr error
+
+	httpResp, err := client.UploadsAPI.UploadsDelete(ctx, uploadHref).Execute()
+	if httpResp != nil {
+		statusCode = httpResp.StatusCode
+		defer httpResp.Body.Close()
+
+		body, readErr = io.ReadAll(httpResp.Body)
+		if readErr != nil {
+			log.Logger.Error().Err(readErr).Msg("DeleteUpload: could not read http body")
+		}
+	}
+	if err != nil {
+		// want to differentiate between resource not found and page not found
+		if statusCode == http.StatusNotFound && strings.Contains(string(body), "No Upload matches") {
+			return statusCode, nil
+		}
+		return statusCode, errorWithResponseBody("deleting upload", httpResp, err)
+	}
+	return statusCode, nil
+}

pkg/clients/pulp_client/uploads.go

@@ -226,4 +226,5 @@ type UploadDao interface {
 	StoreChunkUpload(ctx context.Context, orgID string, uploadUUID string, sha256 string) error
 	GetExistingUploadIDAndCompletedChunks(ctx context.Context, orgID string, sha256 string, chunkSize int64) (string, []string, error)
 	DeleteUpload(ctx context.Context, uploadUUID string) error
+	ListUploadsForCleanup(ctx context.Context) ([]models.Upload, error)
 }

pkg/dao/interfaces.go

@@ -27,7 +27,6 @@ func (t uploadDaoImpl) StoreFileUpload(ctx context.Context, orgID string, upload
 	upload.UploadUUID = uploadUUID
 	upload.Sha256 = sha256
 	upload.ChunkSize = chunkSize
-
 	upload.ChunkList = []string{}
 
 	db := t.db.Model(models.Upload{}).WithContext(ctx).Create(&upload)
@@ -77,3 +76,14 @@ func (t uploadDaoImpl) DeleteUpload(ctx context.Context, uploadUUID string) erro
 	}
 	return nil
 }
+
+func (t uploadDaoImpl) ListUploadsForCleanup(ctx context.Context) ([]models.Upload, error) {
+	var uploads []models.Upload
+	err := t.db.WithContext(ctx).
+		Where("created_at < current_date - INTERVAL '1' day").
+		Find(&uploads).Error
+	if err != nil {
+		return nil, err
+	}
+	return uploads, nil
+}

pkg/dao/uploads.go

@@ -3,6 +3,7 @@ package dao
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/content-services/content-sources-backend/pkg/clients/pulp_client"
 	"github.com/content-services/content-sources-backend/pkg/models"
@@ -77,3 +78,39 @@ func (s *UploadsSuite) TestDeleteUpload() {
 	require.Error(s.T(), err)
 	assert.Equal(s.T(), "record not found", err.Error())
 }
+
+func (s *UploadsSuite) TestListUploadsForCleanup() {
+	uploadDao := s.uploadsDao()
+	ctx := context.Background()
+
+	// Insert an upload with a timestamp older than 1 day
+	oldUpload := models.Upload{
+		CreatedAt:  time.Now().AddDate(0, 0, -2), // 2 days ago
+		UploadUUID: uuid.NewString(),
+		OrgID:      orgIDTest,
+		ChunkSize:  int64(1),
+		Sha256:     uuid.NewString(),
+		ChunkList:  []string{uuid.NewString()},
+	}
+	err := s.tx.Create(&oldUpload).Error
+	require.NoError(s.T(), err)
+
+	// Insert an upload with a recent timestamp
+	recentUpload := models.Upload{
+		CreatedAt:  time.Now(),
+		UploadUUID: uuid.NewString(),
+		OrgID:      orgIDTest,
+		ChunkSize:  int64(1),
+		Sha256:     uuid.NewString(),
+		ChunkList:  []string{uuid.NewString()},
+	}
+	err = s.tx.Create(&recentUpload).Error
+	require.NoError(s.T(), err)
+
+	uploads, err := uploadDao.ListUploadsForCleanup(ctx)
+	require.NoError(s.T(), err)
+
+	// Assert that only the old upload was found
+	assert.Equal(s.T(), 1, len(uploads))
+	assert.Equal(s.T(), oldUpload.UploadUUID, uploads[0].UploadUUID)
+}

pkg/dao/uploads_test.go

@@ -38,9 +38,9 @@ def generate_checksum(rpm_file):
             sha256_hash.update(byte_block)
     return sha256_hash.hexdigest()
 
-def create_upload(size):
+def create_upload(size, rpm_sha256):
     # create the upload
-    data = {'size': size}
+    data = {'size': size, 'chunk_size': size, 'sha256': rpm_sha256}
     headers = {
         'x-rh-identity': IDENTITY_HEADER,
         'Content-Type': 'application/json'
@@ -187,7 +187,8 @@ def main():
         print(f'sha256 for rpm: {rpm_sha256}')
 
         # create the upload
-        upload_id = create_upload(rpm_size)
+        print(f'rpm_size: {rpm_size}')
+        upload_id = create_upload(rpm_size, rpm_sha256)
         print(f'upload href or uuid: {upload_id}')
         upload_ids += [(rpm_sha256, upload_id)]