Haystack/backend/models/image.go

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
}

type ImageWithProperties struct {
	ID uuid.UUID

	Image model.UserImages

	Tags  []model.ImageTags
	Links []model.ImageLinks
	Text  []model.ImageText
}

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)

	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
}

// TODO: move this to `user.go` model file
func (m ImageModel) ListWithProperties(ctx context.Context, userId uuid.UUID) ([]ImageWithProperties, error) {
	listWithPropertiesStmt := SELECT(
		UserImages.ID.AS("UserImagesWithInfo.ID"),
		Image.ID,
		Image.ImageName,
		ImageTags.AllColumns,
		ImageText.AllColumns,
		ImageLinks.AllColumns).
		FROM(
			UserImages.INNER_JOIN(Image, Image.ID.EQ(UserImages.ImageID)).
				LEFT_JOIN(ImageTags, ImageTags.ImageID.EQ(UserImages.ID)).
				LEFT_JOIN(ImageText, ImageText.ImageID.EQ(UserImages.ID)).
				LEFT_JOIN(ImageLinks, ImageLinks.ImageID.EQ(UserImages.ID))).
		WHERE(UserImages.UserID.EQ(UUID(userId)))

	images := []ImageWithProperties{}
	err := listWithPropertiesStmt.QueryContext(ctx, m.dbPool, &images)

	return images, err
}

func NewImageModel(db *sql.DB) ImageModel {
	return ImageModel{dbPool: db}
}