query-orchestration/internal/document/clean/clean.go

package documentclean

import (
	"context"
	"fmt"
	"log/slog"
	"queryorchestration/internal/database"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
)

type CleanParams struct {
	ID       uuid.UUID
	Hash     string
	Location document.Location
}

type InvalidDocumentReason string

const (
	InvalidDocumentMimeType        InvalidDocumentReason = "invalid_mimetype"
	InvalidDocumentRead            InvalidDocumentReason = "invalid_read"
	InvalidDocumentReadPages       InvalidDocumentReason = "invalid_read_pages"
	InvalidDocumentZeroPageCount   InvalidDocumentReason = "zero_page_count"
	InvalidDocumentLargePageCount  InvalidDocumentReason = "large_page_count"
	InvalidDocumentLargeFile       InvalidDocumentReason = "large_file"
	InvalidDocumentSmallDimensions InvalidDocumentReason = "small_dimensions"
	InvalidDocumentLargeDimensions InvalidDocumentReason = "large_dimensions"
)

type ExecuteCleanResponse struct {
	location   *document.Location
	mimetype   *MimeType
	failReason *InvalidDocumentReason
}

func (s *Service) executeCleanTasks(ctx context.Context, params *CleanParams) (*ExecuteCleanResponse, error) {
	out, err := s.cfg.GetStoreClient().HeadObject(ctx, &s3.HeadObjectInput{
		Bucket:  &params.Location.Bucket,
		Key:     &params.Location.Key,
		IfMatch: &params.Hash,
	})
	if err != nil {
		return nil, err
	}

	length := int64(1)
	if out.ContentLength != nil {
		length = *out.ContentLength
	}

	mimeType, err := s.getAcceptedMimeType(ctx, params, out.ContentType, length)
	if err != nil {
		return nil, err
	} else if mimeType == MimeTypeInvalid {
		reason := InvalidDocumentMimeType
		return &ExecuteCleanResponse{
			failReason: &reason,
		}, nil
	}

	content, err := s.getContent(ctx, params, mimeType)
	if err != nil {
		return nil, err
	}

	corruptReason := content.isCorrupt(ctx)
	if corruptReason != nil {
		return &ExecuteCleanResponse{
			failReason: corruptReason,
		}, nil
	}

	return &ExecuteCleanResponse{
		location: &params.Location,
		mimetype: &mimeType,
	}, nil
}

func (s *Service) clean(ctx context.Context, id uuid.UUID) error {
	slog.Debug("cleaning document", "id", id.String())
	docId := database.MustToDBUUID(id)

	doc, err := s.cfg.GetDBQueries().GetDocument(ctx, docId)
	if err != nil {
		return err
	}

	entry, err := s.cfg.GetDBQueries().GetDocumentEntry(ctx, docId)
	if err != nil {
		return err
	}

	out, err := s.executeCleanTasks(ctx, &CleanParams{
		ID:   id,
		Hash: doc.Hash,
		Location: document.Location{
			Bucket: entry.Bucket,
			Key:    entry.Key,
		},
	})
	if err != nil {
		return err
	}

	version := s.svc.Version.GetVersion()

	params := &repository.AddDocumentCleanEntryParams{
		Documentid: docId,
		Version:    version,
	}
	if out.failReason != nil {
		slog.Info("Failed document", "id", id, "reason", *out.failReason)
		params.Fail = repository.NullCleanfailtype{
			Cleanfailtype: repository.Cleanfailtype(*out.failReason),
			Valid:         true,
		}
	} else {
		params.Bucket = &out.location.Bucket
		params.Key = &out.location.Key
		params.Mimetype = repository.NullCleanmimetypes{
			Cleanmimetypes: repository.Cleanmimetypes(*out.mimetype),
			Valid:          true,
		}
	}

	err = s.cfg.GetDBQueries().AddDocumentCleanEntry(ctx, params)
	if err != nil {
		return err
	}

	if out.failReason != nil {
		return fmt.Errorf("%s", *out.failReason)
	}

	return nil
}