query-orchestration/internal/document/init/create.go

package documentinit

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"queryorchestration/internal/database"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document"
	documentclean "queryorchestration/internal/document/clean"
	"queryorchestration/internal/job"
	"queryorchestration/internal/serviceconfig/queue"
	"regexp"

	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgtype"
)

type Create struct {
	JobID    uuid.UUID
	Location document.Location
	Bucket   string
	Hash     string
}

func (s *Service) Create(ctx context.Context, doc *Create) (uuid.UUID, error) {
	j, err := s.svc.Job.Get(ctx, doc.JobID)
	if err != nil {
		return uuid.Nil, err
	}

	params, err := s.getCreateParams(ctx, doc)
	if err != nil {
		return uuid.Nil, err
	}

	id, err := s.submitCreate(ctx, params)
	if err != nil {
		return uuid.Nil, err
	}

	if params.ID == nil {
		err = s.informCreate(ctx, id, j)
		if err != nil {
			return uuid.Nil, err
		}
	}

	return id, nil
}

type createDocumentParams struct {
	ID       *pgtype.UUID
	JobID    pgtype.UUID
	Hash     string
	Bucket   string
	Location string
}

func (s *Service) getCreateParams(ctx context.Context, doc *Create) (*createDocumentParams, error) {
	idbyhash, err := s.cfg.GetDBQueries().GetDocumentIDByHash(ctx, &repository.GetDocumentIDByHashParams{
		Jobid: database.MustToDBUUID(doc.JobID),
		Hash:  doc.Hash,
	})
	var docID *pgtype.UUID
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, err
	} else if err == nil {
		docID = &idbyhash
	}

	return &createDocumentParams{
		ID:       docID,
		JobID:    database.MustToDBUUID(doc.JobID),
		Hash:     doc.Hash,
		Bucket:   doc.Bucket,
		Location: doc.Location,
	}, nil
}

func (s *Service) submitCreate(ctx context.Context, params *createDocumentParams) (uuid.UUID, error) {
	var id uuid.UUID

	err := s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
		var dbid pgtype.UUID
		if params.ID == nil {
			createid, err := s.cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
				Jobid: params.JobID,
				Hash:  params.Hash,
			})
			if err != nil {
				return err
			}

			dbid = createid
		} else {
			dbid = *params.ID
		}

		err := s.cfg.GetDBQueries().AddDocumentEntry(ctx, &repository.AddDocumentEntryParams{
			Documentid: dbid,
			Bucket:     params.Bucket,
			Location:   params.Location,
		})
		if err != nil {
			return err
		}

		id = database.MustToUUID(dbid)

		return nil
	})
	if err != nil {
		return uuid.Nil, err
	}

	return id, nil
}

func (s *Service) informCreate(ctx context.Context, id uuid.UUID, j *job.Job) error {
	if !j.CanSync {
		return nil
	}

	err := s.cfg.SendToQueue(ctx, &queue.SendParams{
		QueueURL: s.cfg.GetDocumentCleanURL(),
		Body: documentclean.Create{
			ID: id,
		},
		Attributes: map[string]types.MessageAttributeValue{},
	})
	if err != nil {
		return err
	}

	return nil
}

func (s *Service) GetJobIDFromKey(key string) (uuid.UUID, error) {
	re := regexp.MustCompile(`^.+/(.+)/.+$`)
	match := re.FindStringSubmatch(key)
	if match == nil || len(match) != 2 {
		return uuid.Nil, fmt.Errorf("no job id match found in key")
	}

	id, err := uuid.Parse(match[1])
	if err != nil {
		return uuid.Nil, err
	}

	return id, nil
}