modelRT/database/query_component_measurement.go

// Package database define database operation functions
package database

import (
	"context"
	"fmt"
	"strings"
	"time"

	"modelRT/common"
	"modelRT/orm"
	"modelRT/sql"

	"golang.org/x/sync/errgroup"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// QueryMeasurementByID return the result of query circuit diagram component measurement info by id from postgresDB
func QueryMeasurementByID(ctx context.Context, tx *gorm.DB, id int64) (orm.Measurement, error) {
	var measurement orm.Measurement
	// ctx超时判断
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	result := tx.WithContext(cancelCtx).
		Where(sql.MeasurementIDWhere, id).
		Clauses(clause.Locking{Strength: "UPDATE"}).
		First(&measurement)

	if result.Error != nil {
		return orm.Measurement{}, result.Error
	}
	return measurement, nil
}

// QueryMeasurementByToken define function query circuit diagram component measurement info by token from postgresDB
func QueryMeasurementByToken(ctx context.Context, tx *gorm.DB, token string) (orm.Measurement, error) {
	measurement, _, err := QueryMeasurementByDataObjectToken(ctx, tx, token)
	if err != nil {
		return orm.Measurement{}, err
	}
	return *measurement, nil
}

// ValidateMeasurementToken checks whether token uniquely identifies an existing
// measurement through the measurement, component, bay, station, zone, and grid
// relationships. Supported formats are token1-token7, token4-token7, and
// token4.token7.
func ValidateMeasurementToken(ctx context.Context, db *gorm.DB, token string) error {
	query, args, err := buildMeasurementTokenValidationQuery(token)
	if err != nil {
		return err
	}

	var count int64
	if err := db.WithContext(ctx).Raw(query, args...).Scan(&count).Error; err != nil {
		return fmt.Errorf("query measurement token %q: %w", token, err)
	}

	switch {
	case count == 0:
		return fmt.Errorf("%w: %q", common.ErrMeasurementTokenNotFound, token)
	case count > 1:
		return fmt.Errorf("%w: %q matched %d records", common.ErrAmbiguousMeasurementToken, token, count)
	default:
		return nil
	}
}

// QueryMeasurementByDataObjectToken validates token and returns the existing
// measurement and its owning component for attribute response construction.
func QueryMeasurementByDataObjectToken(ctx context.Context, db *gorm.DB, token string) (*orm.Measurement, *orm.Component, error) {
	validationQuery, args, err := buildMeasurementTokenValidationQuery(token)
	if err != nil {
		return nil, nil, err
	}

	query := buildMeasurementRowsQuery(validationQuery)
	var rows []orm.Measurement
	if err := db.WithContext(ctx).Raw(query, args...).Scan(&rows).Error; err != nil {
		return nil, nil, fmt.Errorf("query measurement token %q: %w", token, err)
	}

	switch len(rows) {
	case 0:
		return nil, nil, fmt.Errorf("%w: %q", common.ErrMeasurementTokenNotFound, token)
	case 1:
		// Continue by loading the owning component.
	default:
		return nil, nil, fmt.Errorf("%w: %q matched more than one record", common.ErrAmbiguousMeasurementToken, token)
	}

	var component orm.Component
	result := db.WithContext(ctx).
		Raw(compactMeasurementSQL(sql.MeasurementComponentByUUID), rows[0].ComponentUUID).
		Scan(&component)
	if result.Error != nil {
		return nil, nil, fmt.Errorf("query component for measurement token %q: %w", token, result.Error)
	}
	if result.RowsAffected == 0 {
		return nil, nil, fmt.Errorf("%w: component for %q", common.ErrMeasurementTokenNotFound, token)
	}

	return &rows[0], &component, nil
}

func buildMeasurementRowsQuery(validationQuery string) string {
	measurementQuery := strings.Replace(
		validationQuery,
		sql.MeasurementCountSelect,
		sql.MeasurementRowsSelect,
		1,
	)
	return compactMeasurementSQL(strings.Join([]string{measurementQuery, sql.MeasurementLimitTwo}, "\n"))
}

func compactMeasurementSQL(statement string) string {
	return strings.Join(strings.Fields(statement), " ")
}

func buildMeasurementTokenValidationQuery(token string) (string, []any, error) {
	parts := strings.Split(token, ".")
	for _, part := range parts {
		if part == "" {
			return "", nil, fmt.Errorf("%w %q: token segment cannot be empty", common.ErrInvalidMeasurementToken, token)
		}
	}

	switch len(parts) {
	case 7:
		if parts[5] != "bay" {
			return "", nil, fmt.Errorf("%w %q: token6 must be bay", common.ErrInvalidMeasurementToken, token)
		}
		query := compactMeasurementSQL(strings.Join([]string{
			sql.MeasurementTokenValidationQueryBase,
			sql.MeasurementSevenPartTokenWhere,
		}, "\n"))
		return query, []any{parts[0], parts[1], parts[2], parts[3], parts[4], parts[6]}, nil
	case 4:
		if parts[2] != "bay" {
			return "", nil, fmt.Errorf("%w %q: token6 must be bay", common.ErrInvalidMeasurementToken, token)
		}
		query := compactMeasurementSQL(strings.Join([]string{
			sql.MeasurementTokenValidationQueryBase,
			sql.MeasurementFourPartTokenWhere,
		}, "\n"))
		return query, []any{parts[0], parts[1], parts[3]}, nil
	case 2:
		query := compactMeasurementSQL(strings.Join([]string{
			sql.MeasurementTokenValidationQueryBase,
			sql.MeasurementTwoPartTokenWhere,
		}, "\n"))
		return query, []any{parts[0], parts[1]}, nil
	default:
		return "", nil, fmt.Errorf("%w %q: expected 2, 4, or 7 segments, got %d", common.ErrInvalidMeasurementToken, token, len(parts))
	}
}

// GetAllMeasurements define func to query all measurement info from postgresDB
func GetAllMeasurements(ctx context.Context, tx *gorm.DB) ([]orm.Measurement, error) {
	var measurements []orm.Measurement

	// ctx超时判断
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	result := tx.WithContext(cancelCtx).Clauses(clause.Locking{Strength: "UPDATE"}).Find(&measurements)
	if result.Error != nil {
		return nil, result.Error
	}
	return measurements, nil
}

// GetFullMeasurementSet queries all hierarchy tags required to build
// measurement recommendations.
func GetFullMeasurementSet(ctx context.Context, db *gorm.DB) (*orm.MeasurementSet, error) {
	mSet := &orm.MeasurementSet{
		GridToZoneTags:       make(map[string][]string),
		ZoneToStationTags:    make(map[string][]string),
		StationToCompNSPaths: make(map[string][]string),
		CompNSPathToCompTags: make(map[string][]string),
		CompTagToMeasTags:    make(map[string][]string),
		CompNSPathToMeasTags: make(map[string][]string),
	}

	g, gctx := errgroup.WithContext(ctx)
	db = db.WithContext(gctx)

	g.Go(func() error {
		var grids []orm.Grid
		if err := db.Raw(compactMeasurementSQL(sql.MeasurementGridTags)).Scan(&grids).Error; err != nil {
			return fmt.Errorf("query grids: %w", err)
		}
		for _, grid := range grids {
			if grid.TAGNAME != "" {
				mSet.AllGridTags = append(mSet.AllGridTags, grid.TAGNAME)
			}
		}
		return nil
	})

	g.Go(func() error {
		var zones []struct {
			orm.Zone
			GridTag string `gorm:"column:grid_tag"`
		}
		if err := db.Raw(compactMeasurementSQL(sql.MeasurementZoneHierarchy)).Scan(&zones).Error; err != nil {
			return fmt.Errorf("query zones: %w", err)
		}
		for _, zone := range zones {
			mSet.AllZoneTags = append(mSet.AllZoneTags, zone.TAGNAME)
			if zone.GridTag != "" {
				mSet.GridToZoneTags[zone.GridTag] = append(mSet.GridToZoneTags[zone.GridTag], zone.TAGNAME)
			}
		}
		return nil
	})

	g.Go(func() error {
		var stations []struct {
			orm.Station
			ZoneTag string `gorm:"column:zone_tag"`
		}
		if err := db.Raw(compactMeasurementSQL(sql.MeasurementStationHierarchy)).Scan(&stations).Error; err != nil {
			return fmt.Errorf("query stations: %w", err)
		}
		for _, station := range stations {
			mSet.AllStationTags = append(mSet.AllStationTags, station.TAGNAME)
			if station.ZoneTag != "" {
				mSet.ZoneToStationTags[station.ZoneTag] = append(mSet.ZoneToStationTags[station.ZoneTag], station.TAGNAME)
			}
		}
		return nil
	})

	g.Go(func() error {
		var components []struct {
			orm.Component
			StationTag string `gorm:"column:station_tag"`
		}
		if err := db.Raw(compactMeasurementSQL(sql.MeasurementComponentHierarchy)).Scan(&components).Error; err != nil {
			return fmt.Errorf("query components: %w", err)
		}
		for _, component := range components {
			mSet.AllCompNSPaths = append(mSet.AllCompNSPaths, component.NSPath)
			mSet.AllCompTags = append(mSet.AllCompTags, component.Tag)
			if component.StationTag != "" {
				mSet.StationToCompNSPaths[component.StationTag] = append(
					mSet.StationToCompNSPaths[component.StationTag],
					component.NSPath,
				)
			}
			if component.NSPath != "" {
				mSet.CompNSPathToCompTags[component.NSPath] = append(
					mSet.CompNSPathToCompTags[component.NSPath],
					component.Tag,
				)
			}
		}
		return nil
	})

	g.Go(func() error {
		var measurements []struct {
			orm.Measurement
			CompTag    string `gorm:"column:comp_tag"`
			CompNSPath string `gorm:"column:comp_nspath"`
			BayTag     string `gorm:"column:bay_tag"`
		}
		if err := db.Raw(compactMeasurementSQL(sql.MeasurementTagHierarchy)).Scan(&measurements).Error; err != nil {
			return fmt.Errorf("query measurements: %w", err)
		}
		for _, measurement := range measurements {
			mSet.AllMeasTags = append(mSet.AllMeasTags, measurement.Tag)
			if measurement.CompTag != "" {
				mSet.CompTagToMeasTags[measurement.CompTag] = append(
					mSet.CompTagToMeasTags[measurement.CompTag],
					measurement.Tag,
				)
			}
			if m.CompNSPath != "" && m.CompNSPath == m.BayTag {
				mSet.CompNSPathToMeasTags[m.CompNSPath] = append(mSet.CompNSPathToMeasTags[m.CompNSPath], m.Tag)
			}
		}
		return nil
	})

	if err := g.Wait(); err != nil {
		return nil, err
	}

	mSet.AllConfigTags = append(mSet.AllConfigTags, "bay")
	return mSet, nil
}