Foundry-Scrapping-API/internal/foundry/temp_models/db/utils.go

package db

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/jmoiron/sqlx"
)

var (
	ErrNoQuery = errors.New("Query has not been set")
)

type AllowedIds interface {
	~uint | ~string
}

type InsertId[T AllowedIds] struct {
	id        T
	fieldName string
	tableName string
	query     string
}

type SyncDB struct {
	errChan chan error
	wg      sync.WaitGroup
}

func NewSyncDB() *SyncDB {
	return &SyncDB{
		wg:      sync.WaitGroup{},
		errChan: make(chan error),
	}
}

func (s *SyncDB) Wait() error {
	return WaitSync(&s.wg, s.errChan)
}

func WaitSync(wg *sync.WaitGroup, errChan chan error) error {
	wgDone := make(chan struct{})

	go func() {
		wg.Wait()
		close(wgDone)
	}()

	select {
	case <-wgDone:
		return nil
	case err := <-errChan:
		close(wgDone)
		return err
	}
}

type Insertable[T AllowedIds] interface {
	Query(data *InsertId[T])
	Insert(tx *sqlx.Tx, relId *InsertId[T]) error
	InsertCtx(ctx context.Context, tx *sqlx.Tx, relId *InsertId[T]) error
}

func InsertWithCtx[T AllowedIds, I Insertable[T]](tx *sqlx.Tx, data I, relId *InsertId[T]) error {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	data.Query(relId)
	return data.InsertCtx(ctx, tx, relId)
}

func InsertWithCtxParallel[T AllowedIds, I Insertable[T]](syncDb *SyncDB, tx *sqlx.Tx, data I, relId *InsertId[T]) {
	syncDb.wg.Go(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()

		data.Query(relId)
		err := data.InsertCtx(ctx, tx, relId)
		if err != nil {
			syncDb.errChan <- err
		}
	})
}

func InsertSlice[T AllowedIds, I Insertable[T]](tx *sqlx.Tx, data []I, relId *InsertId[T]) error {
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	if len(data) > 0 {
		data[0].Query(relId)
	}
	for i := range data {
		err := data[i].InsertCtx(ctx, tx, relId)
		if err != nil {
			return err
		}
	}
	return nil
}

func InsertSliceParallel[T AllowedIds, I Insertable[T]](syncDb *SyncDB, tx *sqlx.Tx, data []I, relId *InsertId[T]) {
	syncDb.wg.Go(func() {
		wg := sync.WaitGroup{}
		errChan := make(chan error)
		defer close(errChan)

		var err error
		if len(data) > 0 {
			data[0].Query(relId)
		}
		for i := range data {
			wg.Go(func() {
				ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
				defer cancel()

				err = data[i].InsertCtx(ctx, tx, relId)
				if err != nil {
					errChan <- err
				}
			})
		}

		err = WaitSync(&wg, errChan)
		if err != nil {
			syncDb.errChan <- err
		}
	})
}

func InsertSimpleSlice[T AllowedIds, I any](tx *sqlx.Tx, data []I, relId *InsertId[T]) error {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	query := fmt.Sprintf(`
		INSERT INTO %s (%s, value)
		VALUES ($1, $2)`, relId.tableName, relId.fieldName)
	for i := range data {
		_, err := tx.ExecContext(ctx, query, relId.id, data[i])
		if err != nil {
			return err
		}
	}

	return nil
}

func InsertSimpleSliceParallel[T AllowedIds, I any](syncDb *SyncDB, tx *sqlx.Tx, data []I, relId *InsertId[T]) {
	syncDb.wg.Go(func() {
		wg := sync.WaitGroup{}
		errChan := make(chan error)
		defer close(errChan)

		query := fmt.Sprintf(`
			INSERT INTO %s (%s, value)
			VALUES ($1, $2)`, relId.tableName, relId.fieldName)
		for i := range data {
			wg.Go(func() {
				ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
				defer cancel()

				_, err := tx.ExecContext(ctx, query, relId.id, data[i])
				if err != nil {
					errChan <- err
				}
			})
		}

		err := WaitSync(&wg, errChan)
		if err != nil {
			syncDb.errChan <- err
		}
	})
}