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

package documentclean

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"

	"queryorchestration/internal/database"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document"
	"queryorchestration/internal/serviceconfig/build"

	"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"
	InvalidDocumentSmallDPI        InvalidDocumentReason = "small_dpi"
	InvalidDocumentLargeDPI        InvalidDocumentReason = "large_dpi"
)

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
	}

	err = s.storeClean(ctx, id, out)
	if err != nil {
		return err
	}

	return nil
}

func (s *Service) storeClean(ctx context.Context, id uuid.UUID, out *ExecuteCleanResponse) error {
	docId := database.MustToDBUUID(id)
	version := build.GetVersionUnixTimestamp()

	params := &repository.AddDocumentCleanParams{
		Documentid: docId,
	}
	if out.failReason != nil {
		slog.Info("Failed document", "id", id, "reason", *out.failReason)
		params.Fail = ToDBNullFailType(*out.failReason)
	} else {
		params.Bucket = &out.location.Bucket
		params.Key = &out.location.Key
		params.Mimetype = ToDBNullMimeType(*out.mimetype)
	}

	err := s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
		lastClean, err := q.GetMostRecentDocumentCleanEntry(ctx, docId)
		if err != nil && !errors.Is(err, sql.ErrNoRows) {
			return err
		}

		newID := lastClean.ID
		if s.isNewClean(lastClean, out) {
			id, err := q.AddDocumentClean(ctx, params)
			if err != nil {
				return err
			}

			newID = id
		}

		err = q.AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{
			Cleanid: newID,
			Version: version,
		})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

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

	return nil
}

func (s *Service) isNewClean(lastEntry *repository.GetMostRecentDocumentCleanEntryRow, out *ExecuteCleanResponse) bool {
	if lastEntry == nil || !lastEntry.ID.Valid {
		return true
	}

	if out.failReason != nil {
		return *out.failReason != ParseDBNullFailType(lastEntry.Fail)
	}

	if lastEntry.Fail.Valid {
		return true
	}

	return !(out.location.Bucket == *lastEntry.Bucket &&
		out.location.Key == *lastEntry.Key &&
		*out.mimetype == ParseDBNullMimeType(lastEntry.Mimetype))
}