package models import ( "context" "database/sql" "errors" "fmt" "screenmark/screenmark/.gen/haystack/haystack/model" . "screenmark/screenmark/.gen/haystack/haystack/table" . "github.com/go-jet/jet/v2/postgres" "github.com/google/uuid" ) type ImageModel struct { dbPool *sql.DB } type ImageData struct { model.UserImages Image model.Image } type ProcessingImageData struct { model.UserImagesToProcess Image model.Image } func (m ImageModel) Process(ctx context.Context, userId uuid.UUID, image model.Image) (model.UserImagesToProcess, error) { tx, err := m.dbPool.BeginTx(ctx, nil) if err != nil { return model.UserImagesToProcess{}, err } insertImageStmt := Image. INSERT(Image.ImageName, Image.Image). VALUES(image.ImageName, image.Image). RETURNING(Image.ID) insertedImage := model.Image{} err = insertImageStmt.QueryContext(ctx, tx, &insertedImage) if err != nil { return model.UserImagesToProcess{}, err } stmt := UserImagesToProcess. INSERT(UserImagesToProcess.UserID, UserImagesToProcess.ImageID). VALUES(userId, insertedImage.ID). RETURNING(UserImagesToProcess.AllColumns) userImage := model.UserImagesToProcess{} err = stmt.QueryContext(ctx, tx, &userImage) if err != nil { return model.UserImagesToProcess{}, err } err = tx.Commit() return userImage, err } func (m ImageModel) GetToProcess(ctx context.Context, imageId uuid.UUID) (model.UserImagesToProcess, error) { getToProcessStmt := UserImagesToProcess. SELECT(UserImagesToProcess.AllColumns). WHERE(UserImagesToProcess.ID.EQ(UUID(imageId))) images := []model.UserImagesToProcess{} err := getToProcessStmt.QueryContext(ctx, m.dbPool, &images) if len(images) != 1 { return model.UserImagesToProcess{}, errors.New(fmt.Sprintf("Expected 1, got %d\n", len(images))) } return images[0], err } func (m ImageModel) GetToProcessWithData(ctx context.Context, imageId uuid.UUID) (ProcessingImageData, error) { stmt := SELECT(UserImagesToProcess.AllColumns, Image.AllColumns). FROM( UserImagesToProcess.INNER_JOIN( Image, Image.ID.EQ(UserImagesToProcess.ImageID), ), ).WHERE(UserImagesToProcess.ID.EQ(UUID(imageId))) images := []ProcessingImageData{} err := stmt.QueryContext(ctx, m.dbPool, &images) if len(images) != 1 { return ProcessingImageData{}, errors.New(fmt.Sprintf("Expected 1, got %d\n", len(images))) } return images[0], err } func (m ImageModel) FinishProcessing(ctx context.Context, imageId uuid.UUID) (model.UserImages, error) { imageToProcess, err := m.GetToProcess(ctx, imageId) if err != nil { return model.UserImages{}, err } tx, err := m.dbPool.Begin() if err != nil { return model.UserImages{}, err } insertImageStmt := UserImages. INSERT(UserImages.UserID, UserImages.ImageID). VALUES(imageToProcess.UserID, imageToProcess.ImageID). RETURNING(UserImages.ID, UserImages.UserID, UserImages.ImageID) userImage := model.UserImages{} err = insertImageStmt.QueryContext(ctx, tx, &userImage) if err != nil { return model.UserImages{}, err } removeProcessingStmt := UserImagesToProcess. DELETE(). WHERE(UserImagesToProcess.ID.EQ(UUID(imageToProcess.ID))) _, err = removeProcessingStmt.ExecContext(ctx, tx) if err != nil { return model.UserImages{}, err } err = tx.Commit() return userImage, err } func (m ImageModel) Get(ctx context.Context, imageId uuid.UUID) (ImageData, error) { getImageStmt := SELECT(UserImages.AllColumns, Image.AllColumns). FROM( UserImages.INNER_JOIN(Image, Image.ID.EQ(UserImages.ImageID)), ). WHERE(UserImages.ID.EQ(UUID(imageId))) images := []ImageData{} err := getImageStmt.QueryContext(ctx, m.dbPool, &images) if len(images) != 1 { return ImageData{}, errors.New(fmt.Sprintf("Expected 1, got %d\n", len(images))) } return images[0], err } func NewImageModel(db *sql.DB) ImageModel { return ImageModel{dbPool: db} }