emissary/internal/datastore/sqlite/tenant_repository.go

package sqlite

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"forge.cadoles.com/Cadoles/emissary/internal/datastore"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

type TenantRepository struct {
	repository
}

// Query implements datastore.TenantRepository.
func (r *TenantRepository) Query(ctx context.Context, opts ...datastore.TenantQueryOptionFunc) ([]*datastore.Tenant, int, error) {
	options := &datastore.TenantQueryOptions{}
	for _, fn := range opts {
		fn(options)
	}

	tenants := make([]*datastore.Tenant, 0)
	count := 0

	err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
		query := `SELECT id, label, created_at, updated_at FROM tenants`

		limit := 10
		if options.Limit != nil {
			limit = *options.Limit
		}

		offset := 0
		if options.Offset != nil {
			offset = *options.Offset
		}

		filters := ""
		paramIndex := 3
		args := []any{offset, limit}

		if options.IDs != nil && len(options.IDs) > 0 {
			filter, newArgs, newParamIndex := inFilter("id", paramIndex, options.IDs)
			filters += filter
			paramIndex = newParamIndex
			args = append(args, newArgs...)
		}

		if filters != "" {
			filters = ` WHERE ` + filters
		}

		query += filters + ` LIMIT $2 OFFSET $1`

		logger.Debug(ctx, "executing query", logger.F("query", query), logger.F("args", args))

		rows, err := tx.QueryContext(ctx, query, args...)
		if err != nil {
			return errors.WithStack(err)
		}

		defer func() {
			if err := rows.Close(); err != nil {
				err = errors.WithStack(err)
				logger.Error(ctx, "could not close rows", logger.CapturedE(err))
			}
		}()

		for rows.Next() {
			tenant := &datastore.Tenant{}

			if err := rows.Scan(&tenant.ID, &tenant.Label, &tenant.CreatedAt, &tenant.UpdatedAt); err != nil {
				return errors.WithStack(err)
			}

			tenants = append(tenants, tenant)
		}

		if err := rows.Err(); err != nil {
			return errors.WithStack(err)
		}

		row := tx.QueryRowContext(ctx, `SELECT count(id) FROM tenants `+filters, args...)
		if err := row.Scan(&count); err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return nil, 0, errors.WithStack(err)
	}

	return tenants, count, nil
}

// Create implements datastore.TenantRepository.
func (r *TenantRepository) Create(ctx context.Context, label string) (*datastore.Tenant, error) {
	var tenant datastore.Tenant

	err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
		now := time.Now().UTC()

		query := `
		INSERT INTO tenants (id, label, created_at, updated_at) 
		VALUES($1, $2, $3, $3)
		RETURNING "id", "label", "created_at", "updated_at"
		`

		tenantID := datastore.NewTenantID()

		row := tx.QueryRowContext(
			ctx, query,
			tenantID, label, now,
		)

		if err := row.Scan(&tenant.ID, &tenant.Label, &tenant.CreatedAt, &tenant.UpdatedAt); err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &tenant, nil
}

// Delete implements datastore.TenantRepository.
func (r *TenantRepository) Delete(ctx context.Context, id datastore.TenantID) error {
	err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
		if exists, err := r.tenantExists(ctx, tx, id); !exists {
			return errors.WithStack(err)
		}

		query := `DELETE FROM tenants WHERE id = $1`
		_, err := tx.ExecContext(ctx, query, id)
		if err != nil {
			return errors.WithStack(err)
		}

		query = `DELETE FROM agents WHERE tenant_id = $1`
		_, err = tx.ExecContext(ctx, query, id)
		if err != nil {
			return errors.WithStack(err)
		}

		query = `DELETE FROM specs WHERE tenant_id = $1`
		_, err = tx.ExecContext(ctx, query, id)
		if err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Get implements datastore.TenantRepository.
func (r *TenantRepository) Get(ctx context.Context, id datastore.TenantID) (*datastore.Tenant, error) {
	var tenant datastore.Tenant

	err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
		query := `
		SELECT "id", "label", "created_at", "updated_at"
		FROM tenants
		WHERE id = $1
		`

		row := tx.QueryRowContext(ctx, query, id)

		if err := row.Scan(&tenant.ID, &tenant.Label, &tenant.CreatedAt, &tenant.UpdatedAt); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return errors.WithStack(datastore.ErrNotFound)
			}

			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &tenant, nil
}

// Update implements datastore.TenantRepository.
func (r *TenantRepository) Update(ctx context.Context, id datastore.TenantID, updates ...datastore.TenantUpdateOptionFunc) (*datastore.Tenant, error) {
	options := &datastore.TenantUpdateOptions{}
	for _, fn := range updates {
		fn(options)
	}

	var tenant datastore.Tenant

	err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
		if exists, err := r.tenantExists(ctx, tx, id); !exists {
			return errors.WithStack(err)
		}

		query := `
			UPDATE tenants SET updated_at = $1
		`

		args := []any{id}
		index := 2

		if options.Label != nil {
			query += fmt.Sprintf(`, label = $%d`, index)
			args = append(args, *options.Label)
			index++
		}

		updated := options.Label != nil
		if updated {
			now := time.Now().UTC()
			query += fmt.Sprintf(`, updated_at = $%d`, index)
			args = append(args, now)
			index++
		}

		query += `
			WHERE id = $1 
			RETURNING "id", "label", "created_at", "updated_at"
		`

		logger.Debug(ctx, "executing query", logger.F("query", query), logger.F("args", args))

		row := tx.QueryRowContext(ctx, query, args...)

		if err := row.Scan(&tenant.ID, &tenant.Label, &tenant.CreatedAt, &tenant.UpdatedAt); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return errors.WithStack(datastore.ErrNotFound)
			}

			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &tenant, nil
}

func (r *TenantRepository) tenantExists(ctx context.Context, tx *sql.Tx, tenantID datastore.TenantID) (bool, error) {
	row := tx.QueryRowContext(ctx, `SELECT count(id) FROM tenants WHERE id = $1`, tenantID)

	var count int

	if err := row.Scan(&count); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return false, errors.WithStack(datastore.ErrNotFound)
		}

		return false, errors.WithStack(err)
	}

	if count == 0 {
		return false, errors.WithStack(datastore.ErrNotFound)
	}

	return true, nil
}

func NewTenantRepository(db *sql.DB, sqliteBusyRetryMaxAttempts int) *TenantRepository {
	return &TenantRepository{
		repository: repository{db, sqliteBusyRetryMaxAttempts},
	}
}

var _ datastore.TenantRepository = &TenantRepository{}