[SCD] Populate uss_availability field of operational intents (#971)

* [SCD] Populate uss_availability field of operational intents * handle availability unknown
interuss · Nov 14, 2023 · e712cac · e712cac
1 parent 2d8027a
commit e712cac
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 13 deletions.
diff --git a/pkg/scd/models/operational_intents.go b/pkg/scd/models/operational_intents.go
@@ -53,18 +53,19 @@ func (s OperationalIntentState) IsValidInDSS() bool {
 // OperationalIntent models an operational intent.
 type OperationalIntent struct {
 	// Reference
-	ID             dssmodels.ID
-	Manager        dssmodels.Manager
-	Version        VersionNumber
-	State          OperationalIntentState
-	OVN            OVN
-	StartTime      *time.Time
-	EndTime        *time.Time
-	USSBaseURL     string
-	SubscriptionID dssmodels.ID
-	AltitudeLower  *float32
-	AltitudeUpper  *float32
-	Cells          s2.CellUnion
+	ID              dssmodels.ID
+	Manager         dssmodels.Manager
+	UssAvailability UssAvailabilityState
+	Version         VersionNumber
+	State           OperationalIntentState
+	OVN             OVN
+	StartTime       *time.Time
+	EndTime         *time.Time
+	USSBaseURL      string
+	SubscriptionID  dssmodels.ID
+	AltitudeLower   *float32
+	AltitudeUpper   *float32
+	Cells           s2.CellUnion
 }
 
 func (s OperationalIntentState) String() string {
@@ -86,7 +87,7 @@ func (o *OperationalIntent) ToRest() *restapi.OperationalIntentReference {
 		UssBaseUrl:      restapi.OperationalIntentUssBaseURL(o.USSBaseURL),
 		SubscriptionId:  restapi.SubscriptionID(o.SubscriptionID.String()),
 		State:           o.State.ToRest(),
-		UssAvailability: UssAvailabilityStateUnknown.ToRest(),
+		UssAvailability: o.UssAvailability.ToRest(),
 	}
 
 	if o.StartTime != nil {

diff --git a/pkg/scd/store/cockroach/operational_intents.go b/pkg/scd/store/cockroach/operational_intents.go
@@ -13,6 +13,8 @@ import (
 	dsssql "github.com/interuss/dss/pkg/sql"
 	"github.com/interuss/stacktrace"
 	"github.com/jackc/pgtype"
+	"github.com/jackc/pgx/v4"
+	"github.com/pkg/errors"
 )
 
 var (
@@ -59,6 +61,7 @@ func (s *repo) fetchOperationalIntents(ctx context.Context, q dsssql.Queryable,
 
 	var payload []*scdmodels.OperationalIntent
 	pgCids := pgtype.Int8Array{}
+	ussAvailabilities := map[dssmodels.Manager]scdmodels.UssAvailabilityState{}
 	for rows.Next() {
 		var (
 			o         = &scdmodels.OperationalIntent{}
@@ -87,16 +90,30 @@ func (s *repo) fetchOperationalIntents(ctx context.Context, q dsssql.Queryable,
 		}
 		o.OVN = scdmodels.NewOVNFromTime(updatedAt, o.ID.String())
 		o.SetCells(cids)
+		ussAvailabilities[o.Manager] = scdmodels.UssAvailabilityStateUnknown
 		payload = append(payload, o)
 	}
 	if err := rows.Err(); err != nil {
 		return nil, stacktrace.Propagate(err, "Error in rows query result")
 	}
 
+	for manager := range ussAvailabilities {
+		ussAvailability, err := s.GetUssAvailability(ctx, manager)
+		if err != nil && !errors.Is(err, pgx.ErrNoRows) {
+			return nil, stacktrace.Propagate(err, "Error getting USS availability of %s", manager)
+		}
+
+		if ussAvailability != nil {
+			ussAvailabilities[manager] = ussAvailability.Availability
+		}
+	}
+
 	for _, op := range payload {
 		if err := s.populateOperationalIntentCells(ctx, q, op); err != nil {
 			return nil, stacktrace.Propagate(err, "Error populating cells for Operation %s", op.ID)
 		}
+
+		op.UssAvailability = ussAvailabilities[op.Manager]
 	}
 
 	return payload, nil