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

package documentclean

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

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document"
	documenttypes "queryorchestration/internal/document/types"
	"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 ExecuteCleanResponse struct {
	location   *document.Location
	mimetype   *documenttypes.MimeType
	hash       *string
	failReason *documenttypes.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 == documenttypes.MimeTypeInvalid {
		reason := documenttypes.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,
		hash:     &params.Hash,
	}, nil
}

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

	doc, err := s.cfg.GetDBQueries().GetDocumentSummary(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 := id
	version := build.GetVersionUnixTimestamp()

	params := &repository.AddDocumentCleanParams{
		Documentid: docId,
	}
	if out.failReason != nil {
		slog.Info("Failed document", "id", id, "reason", *out.failReason)
		params.Fail = documenttypes.ToDBNullFailType(*out.failReason)
	} else {
		params.Bucket = &out.location.Bucket
		params.Key = &out.location.Key
		params.Hash = out.hash
		params.Mimetype = documenttypes.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 {
		return true
	}

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

	if lastEntry.Fail.Valid {
		return true
	}

	if out.location == nil || out.mimetype == nil || lastEntry.Bucket == nil || lastEntry.Key == nil {
		return true
	}

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