o11y/pkg/query-service/agentConf/db.go

package agentConf

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

	"log/slog"

	"golang.org/x/exp/slices"

	"github.com/hanzoai/o11y/pkg/errors"
	"github.com/hanzoai/o11y/pkg/query-service/model"
	"github.com/hanzoai/o11y/pkg/sqlstore"
	"github.com/hanzoai/o11y/pkg/types"
	"github.com/hanzoai/o11y/pkg/types/opamptypes"
	"github.com/hanzoai/o11y/pkg/valuer"
)

var (
	CodeConfigVersionNotFound          = errors.MustNewCode("config_version_not_found")
	CodeElementTypeRequired            = errors.MustNewCode("element_type_required")
	CodeConfigElementsRequired         = errors.MustNewCode("config_elements_required")
	CodeConfigVersionInsertFailed      = errors.MustNewCode("config_version_insert_failed")
	CodeConfigElementInsertFailed      = errors.MustNewCode("config_element_insert_failed")
	CodeConfigDeployStatusUpdateFailed = errors.MustNewCode("config_deploy_status_update_failed")
	CodeConfigHistoryGetFailed         = errors.MustNewCode("config_history_get_failed")
)

// Repo handles DDL and DML ops on ingestion rules
type Repo struct {
	store sqlstore.SQLStore
}

func (r *Repo) GetConfigHistory(
	ctx context.Context, orgId valuer.UUID, typ opamptypes.ElementType, limit int,
) ([]opamptypes.AgentConfigVersion, error) {
	var c []opamptypes.AgentConfigVersion
	err := r.store.BunDB().NewSelect().
		Model(&c).
		ColumnExpr("id, version, element_type, deploy_status, deploy_result, created_at").
		ColumnExpr("COALESCE(created_by, '') as created_by").
		ColumnExpr(`COALESCE((SELECT display_name FROM users WHERE users.id = acv.created_by), 'unknown') as created_by_name`).
		ColumnExpr("COALESCE(hash, '') as hash, COALESCE(config, '{}') as config").
		Where("acv.element_type = ?", typ).
		Where("acv.org_id = ?", orgId).
		OrderExpr("acv.created_at DESC, acv.version DESC").
		Limit(limit).
		Scan(ctx)

	if err != nil {
		return nil, errors.WrapInternalf(err, CodeConfigHistoryGetFailed, "failed to get config history")
	}

	incompleteStatuses := []opamptypes.DeployStatus{opamptypes.DeployInitiated, opamptypes.Deploying}
	for idx := 1; idx < len(c); idx++ {
		if slices.Contains(incompleteStatuses, c[idx].DeployStatus) {
			c[idx].DeployStatus = opamptypes.DeployStatusUnknown
		}
	}

	return c, nil
}

func (r *Repo) GetConfigVersion(
	ctx context.Context, orgId valuer.UUID, typ opamptypes.ElementType, v int,
) (*opamptypes.AgentConfigVersion, error) {
	var c opamptypes.AgentConfigVersion
	err := r.store.BunDB().NewSelect().
		Model(&c).
		ColumnExpr("id, version, element_type, deploy_status, deploy_result, created_at").
		ColumnExpr("COALESCE(created_by, '') as created_by").
		ColumnExpr(`COALESCE((SELECT display_name FROM users WHERE users.id = acv.created_by), 'unknown') as created_by_name`).
		ColumnExpr("COALESCE(hash, '') as hash, COALESCE(config, '{}') as config").
		Where("acv.element_type = ?", typ).
		Where("acv.version = ?", v).
		Where("acv.org_id = ?", orgId).
		Scan(ctx)

	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, errors.WrapNotFoundf(err, CodeConfigVersionNotFound, "config version not found")
		}
		return nil, errors.WrapInternalf(err, errors.CodeInternal, "failed to get config version")
	}

	return &c, nil
}

func (r *Repo) GetLatestVersion(
	ctx context.Context, orgId valuer.UUID, typ opamptypes.ElementType,
) (*opamptypes.AgentConfigVersion, error) {
	var c opamptypes.AgentConfigVersion
	err := r.store.BunDB().NewSelect().
		Model(&c).
		ColumnExpr("id, version, element_type, deploy_status, deploy_result, created_at").
		ColumnExpr("COALESCE(created_by, '') as created_by").
		ColumnExpr(`COALESCE((SELECT display_name FROM users WHERE users.id = acv.created_by), 'unknown') as created_by_name`).
		Where("acv.element_type = ?", typ).
		Where("acv.org_id = ?", orgId).
		Where("version = (SELECT MAX(version) FROM agent_config_version WHERE acv.element_type = ?)", typ).
		Scan(ctx)

	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, errors.WrapNotFoundf(err, CodeConfigVersionNotFound, "config latest version not found")
		}
		return nil, errors.WrapInternalf(err, errors.CodeInternal, "failed to get latest config version")
	}

	return &c, nil
}

func (r *Repo) insertConfig(
	ctx context.Context, orgId valuer.UUID, userId valuer.UUID, c *opamptypes.AgentConfigVersion, elements []string,
) error {
	if c.ElementType.StringValue() == "" {
		return errors.NewInvalidInputf(CodeElementTypeRequired, "element type is required for creating agent config version")
	}

	// allowing empty elements for logs - use case is deleting all pipelines
	if len(elements) == 0 && c.ElementType != opamptypes.ElementTypeLogPipelines {
		slog.ErrorContext(ctx, "insert config called with no elements", "element_type", c.ElementType.StringValue())
		return errors.NewInvalidInputf(CodeConfigElementsRequired, "config must have atleast one element")
	}

	if c.Version != 0 {
		// the version can not be set by the user, we want to auto-assign the versions
		// in a monotonically increasing order starting with 1. hence, we reject insert
		// requests with version anything other than 0. here, 0 indicates un-assigned
		slog.ErrorContext(ctx, "invalid version assignment while inserting agent config", "version", c.Version, "element_type", c.ElementType.StringValue())
		return errors.NewInvalidInputf(errors.CodeInvalidInput, "user defined versions are not supported in the agent config")
	}

	configVersion, err := r.GetLatestVersion(ctx, orgId, c.ElementType)
	if err != nil && !errors.Ast(err, errors.TypeNotFound) {
		slog.ErrorContext(ctx, "failed to fetch latest config version", errors.Attr(err))
		return err
	}

	if configVersion != nil {
		c.IncrementVersion(configVersion.Version)
	} else {
		// first version
		c.Version = 1
	}

	// Track whether we've successfully finished the insert operation
	success := false

	defer func() {
		if !success {
			// remove all the damage (invalid rows from db)
			// Delete elements first, then version (to respect potential foreign key constraints)
			_, delErr := r.store.BunDB().NewDelete().Model(new(opamptypes.AgentConfigElement)).Where("version_id = ?", c.ID).Exec(ctx)
			if delErr != nil {
				slog.ErrorContext(ctx, "failed to delete config elements during cleanup", errors.Attr(delErr), "version_id", c.ID.String())
			}
			_, delErr = r.store.BunDB().NewDelete().Model(new(opamptypes.AgentConfigVersion)).Where("id = ?", c.ID).Where("org_id = ?", orgId).Exec(ctx)
			if delErr != nil {
				slog.ErrorContext(ctx, "failed to delete config version during cleanup", errors.Attr(delErr), "version_id", c.ID.String())
			}
		}
	}()

	_, dbErr := r.store.
		BunDB().
		NewInsert().
		Model(c).
		Exec(ctx)
	if dbErr != nil {
		slog.ErrorContext(ctx, "error in inserting config version", errors.Attr(dbErr))
		return errors.WrapInternalf(dbErr, CodeConfigVersionInsertFailed, "failed to insert config version")
	}

	for _, e := range elements {
		agentConfigElement := &opamptypes.AgentConfigElement{
			Identifiable: types.Identifiable{ID: valuer.GenerateUUID()},
			TimeAuditable: types.TimeAuditable{
				CreatedAt: time.Now(),
				UpdatedAt: time.Now(),
			},
			VersionID:   c.ID,
			ElementType: c.ElementType.StringValue(),
			ElementID:   e,
		}
		_, dbErr = r.store.BunDB().NewInsert().Model(agentConfigElement).Exec(ctx)
		if dbErr != nil {
			return errors.WrapInternalf(dbErr, CodeConfigElementInsertFailed, "failed to insert config element")
		}
	}

	success = true
	return nil
}

func (r *Repo) updateDeployStatus(ctx context.Context,
	orgId valuer.UUID,
	elementType opamptypes.ElementType,
	version int,
	status string,
	result string,
	lastHash string,
	lastconf string) *model.ApiError {

	// check if it has org orgID prefix
	// ensuring it here and also ensuring in coordinator.go
	if !strings.HasPrefix(lastHash, orgId.String()) {
		lastHash = orgId.String() + lastHash
	}

	_, err := r.store.BunDB().NewUpdate().
		Model(new(opamptypes.AgentConfigVersion)).
		Set("deploy_status = ?", status).
		Set("deploy_result = ?", result).
		Set("hash = COALESCE(?, hash)", lastHash).
		Set("config = ?", lastconf).
		Where("version = ?", version).
		Where("element_type = ?", elementType).
		Where("org_id = ?", orgId).
		Exec(ctx)
	if err != nil {
		slog.ErrorContext(ctx, "failed to update deploy status", errors.Attr(err))
		return model.BadRequest(fmt.Errorf("failed to update deploy status"))
	}

	return nil
}

// GetDeployStatusByHash returns the DeployStatus for the given config hash
// (stored with orgId prefix). Returns DeployStatusUnknown when no matching row exists.
func (r *Repo) GetDeployStatusByHash(ctx context.Context, orgId valuer.UUID, configHash string) (opamptypes.DeployStatus, error) {
	var version opamptypes.AgentConfigVersion
	err := r.store.BunDB().NewSelect().
		Model(&version).
		ColumnExpr("deploy_status").
		Where("hash = ?", configHash).
		Where("org_id = ?", orgId).
		Scan(ctx)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return opamptypes.DeployStatusUnknown, nil
		}
		return opamptypes.DeployStatusUnknown, errors.WrapInternalf(err, errors.CodeInternal, "failed to query deploy status by hash")
	}
	return version.DeployStatus, nil
}

func (r *Repo) updateDeployStatusByHash(
	ctx context.Context, orgId valuer.UUID, confighash string, status string, result string,
) error {
	_, err := r.store.BunDB().NewUpdate().
		Model(new(opamptypes.AgentConfigVersion)).
		Set("deploy_status = ?", status).
		Set("deploy_result = ?", result).
		Where("hash = ?", confighash).
		Where("org_id = ?", orgId).
		Exec(ctx)
	if err != nil {
		slog.ErrorContext(ctx, "failed to update deploy status", errors.Attr(err))
		return errors.WrapInternalf(err, CodeConfigDeployStatusUpdateFailed, "failed to update deploy status")
	}

	return nil
}