signoz/pkg/modules/cloudintegration/implcloudintegration/store.go

package implcloudintegration

import (
	"context"
	"time"

	"github.com/SigNoz/signoz/pkg/errors"
	"github.com/SigNoz/signoz/pkg/sqlstore"
	"github.com/SigNoz/signoz/pkg/types/cloudintegrationtypes"
	"github.com/SigNoz/signoz/pkg/valuer"
)

type store struct {
	store sqlstore.SQLStore
}

func NewStore(sqlStore sqlstore.SQLStore) cloudintegrationtypes.Store {
	return &store{store: sqlStore}
}

func (s *store) GetAccountByID(ctx context.Context, orgID, id valuer.UUID, provider cloudintegrationtypes.CloudProviderType) (*cloudintegrationtypes.StorableCloudIntegration, error) {
	account := new(cloudintegrationtypes.StorableCloudIntegration)
	err := s.store.BunDB().NewSelect().Model(account).
		Where("id = ?", id).
		Where("org_id = ?", orgID).
		Where("provider = ?", provider).
		Scan(ctx)
	if err != nil {
		return nil, s.store.WrapNotFoundErrf(err, cloudintegrationtypes.ErrCodeCloudIntegrationNotFound, "cloud integration account with id %s not found", id)
	}
	return account, nil
}

func (s *store) CreateAccount(ctx context.Context, orgID valuer.UUID, account *cloudintegrationtypes.StorableCloudIntegration) (*cloudintegrationtypes.StorableCloudIntegration, error) {
	now := time.Now()
	if account.ID.IsZero() {
		account.ID = valuer.GenerateUUID()
	}
	account.OrgID = orgID
	account.CreatedAt = now
	account.UpdatedAt = now

	_, err := s.store.BunDBCtx(ctx).NewInsert().Model(account).Exec(ctx)
	if err != nil {
		return nil, s.store.WrapAlreadyExistsErrf(err, errors.CodeAlreadyExists, "cloud integration account with id %s already exists", account.ID)
	}

	return account, nil
}

func (s *store) UpdateAccount(ctx context.Context, account *cloudintegrationtypes.StorableCloudIntegration) error {
	account.UpdatedAt = time.Now()
	_, err := s.store.BunDBCtx(ctx).NewUpdate().Model(account).
		Where("id = ?", account.ID).
		Where("org_id = ?", account.OrgID).
		Where("provider = ?", account.Provider).
		Exec(ctx)
	return err
}

func (s *store) RemoveAccount(ctx context.Context, orgID, id valuer.UUID, provider cloudintegrationtypes.CloudProviderType) error {
	_, err := s.store.BunDBCtx(ctx).NewUpdate().Model((*cloudintegrationtypes.StorableCloudIntegration)(nil)).
		Set("removed_at = ?", time.Now()).
		Where("id = ?", id).
		Where("org_id = ?", orgID).
		Where("provider = ?", provider).
		Exec(ctx)
	return err
}

func (s *store) GetConnectedAccounts(ctx context.Context, orgID valuer.UUID, provider cloudintegrationtypes.CloudProviderType) ([]*cloudintegrationtypes.StorableCloudIntegration, error) {
	var accounts []*cloudintegrationtypes.StorableCloudIntegration
	err := s.store.BunDB().NewSelect().Model(&accounts).
		Where("org_id = ?", orgID).
		Where("provider = ?", provider).
		Where("removed_at IS NULL").
		Where("account_id IS NOT NULL").
		Where("last_agent_report IS NOT NULL").
		Order("created_at ASC").
		Scan(ctx)
	if err != nil {
		return nil, err
	}
	return accounts, nil
}

func (s *store) GetConnectedAccount(ctx context.Context, orgID valuer.UUID, provider cloudintegrationtypes.CloudProviderType, providerAccountID string) (*cloudintegrationtypes.StorableCloudIntegration, error) {
	account := new(cloudintegrationtypes.StorableCloudIntegration)
	err := s.store.BunDB().NewSelect().Model(account).
		Where("org_id = ?", orgID).
		Where("provider = ?", provider).
		Where("account_id = ?", providerAccountID).
		Where("last_agent_report IS NOT NULL").
		Where("removed_at IS NULL").
		Scan(ctx)
	if err != nil {
		return nil, s.store.WrapNotFoundErrf(err, cloudintegrationtypes.ErrCodeCloudIntegrationNotFound, "connected account with provider account id %s not found", providerAccountID)
	}
	return account, nil
}

func (s *store) GetServiceByType(ctx context.Context, cloudIntegrationID valuer.UUID, serviceType string) (*cloudintegrationtypes.StorableCloudIntegrationService, error) {
	service := new(cloudintegrationtypes.StorableCloudIntegrationService)
	err := s.store.BunDB().NewSelect().Model(service).
		Where("cloud_integration_id = ?", cloudIntegrationID).
		Where("type = ?", serviceType).
		Scan(ctx)
	if err != nil {
		return nil, s.store.WrapNotFoundErrf(err, cloudintegrationtypes.ErrCodeCloudIntegrationNotFound, "cloud integration service with type %s not found", serviceType)
	}
	return service, nil
}

func (s *store) CreateService(ctx context.Context, cloudIntegrationID valuer.UUID, service *cloudintegrationtypes.StorableCloudIntegrationService) (*cloudintegrationtypes.StorableCloudIntegrationService, error) {
	now := time.Now()
	if service.ID.IsZero() {
		service.ID = valuer.GenerateUUID()
	}
	service.CloudIntegrationID = cloudIntegrationID
	if service.CreatedAt.IsZero() {
		service.CreatedAt = now
	}
	service.UpdatedAt = now

	_, err := s.store.BunDBCtx(ctx).NewInsert().Model(service).Exec(ctx)
	if err != nil {
		return nil, s.store.WrapAlreadyExistsErrf(err, errors.CodeAlreadyExists, "cloud integration service with type %s already exists", service.Type)
	}

	return service, nil
}

func (s *store) UpdateService(ctx context.Context, cloudIntegrationID valuer.UUID, service *cloudintegrationtypes.StorableCloudIntegrationService) error {
	service.CloudIntegrationID = cloudIntegrationID
	service.UpdatedAt = time.Now()
	_, err := s.store.BunDBCtx(ctx).NewUpdate().Model(service).
		Where("cloud_integration_id = ?", cloudIntegrationID).
		Where("type = ?", service.Type).
		Exec(ctx)
	return err
}

func (s *store) GetServices(ctx context.Context, cloudIntegrationID valuer.UUID) ([]*cloudintegrationtypes.StorableCloudIntegrationService, error) {
	var services []*cloudintegrationtypes.StorableCloudIntegrationService
	err := s.store.BunDB().NewSelect().Model(&services).
		Where("cloud_integration_id = ?", cloudIntegrationID).
		Scan(ctx)
	if err != nil {
		return nil, err
	}
	return services, nil
}