feat(proto): Add CreateLocationEnergySource RPC

internal/server/dummy/dataserverimpl.go

@@ -299,6 +299,18 @@ func (d *DataPlatformDataServiceServerImpl) CreateLocation(
 	}, nil
 }
 
+// CreateLocationEnergySource implements dp.DataPlatformDataServiceServer.
+func (d *DataPlatformDataServiceServerImpl) CreateLocationEnergySource(
+	ctx context.Context,
+	req *pb.CreateLocationEnergySourceRequest,
+) (*pb.CreateLocationEnergySourceResponse, error) {
+	return &pb.CreateLocationEnergySourceResponse{
+		LocationUuid:           req.LocationUuid,
+		EnergySource:           req.EnergySource,
+		EffectiveCapacityWatts: req.EffectiveCapacityWatts,
+	}, nil
+}
+
 // UpdateLocation implements dp.DataPlatformDataServiceServer.
 func (d *DataPlatformDataServiceServerImpl) UpdateLocation(
 	ctx context.Context,

internal/server/postgres/dataserverimpl.go

@@ -24,6 +24,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	pb "github.com/openclimatefix/data-platform/internal/gen/ocf/dp"
@@ -963,6 +964,40 @@ func (s *DataPlatformDataServiceServerImpl) GetLocationAsTimeseries(
 	}, nil
 }
 
+// createLocationSource inserts a source entry for a geometry.
+// Importantly, this also refreshes the sources materialised view.
+func createLocationSource(
+	ctx context.Context,
+	querier *db.Queries,
+	geometryUuid uuid.UUID,
+	sourceTypeID int16,
+	capacityWatts uint64,
+	metadata *structpb.Struct,
+	validFrom time.Time,
+) (db.CreateSourceEntryRow, error) {
+	csprms := db.CreateSourceEntryParams{
+		GeometryUuid:  geometryUuid,
+		SourceTypeID:  sourceTypeID,
+		CapacityWatts: int64(capacityWatts),
+		Metadata:      metadata,
+		ValidFromUtc:  pgtype.Timestamp{Time: validFrom, Valid: true},
+	}
+
+	dbSource, err := querier.CreateSourceEntry(ctx, csprms)
+	if err != nil {
+		return db.CreateSourceEntryRow{}, fmt.Errorf("invalid location source: %w", err)
+	}
+
+	if err := querier.RefreshSourcesMaterializedView(ctx); err != nil {
+		return db.CreateSourceEntryRow{}, fmt.Errorf(
+			"failed to update sources materialised view: %w",
+			err,
+		)
+	}
+
+	return dbSource, nil
+}
+
 func (s *DataPlatformDataServiceServerImpl) CreateLocation(
 	ctx context.Context,
 	req *pb.CreateLocationRequest,
@@ -1007,17 +1042,12 @@ func (s *DataPlatformDataServiceServerImpl) CreateLocation(
 		req.ValidFromUtc = timestamppb.New(time.Now().UTC().Truncate(time.Minute))
 	}
 
-	csprms := db.CreateSourceEntryParams{
-		GeometryUuid:  dbLocation.GeometryUuid,
-		SourceTypeID:  int16(req.EnergySource),
-		CapacityWatts: int64(req.EffectiveCapacityWatts),
-		Metadata:      req.Metadata,
-		ValidFromUtc:  pgtype.Timestamp{Time: req.ValidFromUtc.AsTime(), Valid: true},
-	}
-
-	dbSource, err := querier.CreateSourceEntry(ctx, csprms)
+	dbSource, err := createLocationSource(
+		ctx, querier, dbLocation.GeometryUuid, int16(req.EnergySource.Number()),
+		req.EffectiveCapacityWatts, req.Metadata, req.ValidFromUtc.AsTime(),
+	)
 	if err != nil {
-		return nil, fmt.Errorf("invalid location: %w", err)
+		return nil, err
 	}
 
 	l.Debug().
@@ -1027,16 +1057,63 @@ func (s *DataPlatformDataServiceServerImpl) CreateLocation(
 		Str("dp.source.valid_from_utc", dbSource.ValidFromUtc.Time.String()).
 		Msg("created source entry for location")
 
-	err = querier.RefreshSourcesMaterializedView(ctx)
+	return &pb.CreateLocationResponse{
+		LocationUuid:           dbLocation.GeometryUuid.String(),
+		LocationName:           dbLocation.GeometryName,
+		EffectiveCapacityWatts: uint64(dbSource.CapacityWatts),
+	}, nil
+}
+
+func (s *DataPlatformDataServiceServerImpl) CreateLocationEnergySource(
+	ctx context.Context,
+	req *pb.CreateLocationEnergySourceRequest,
+) (*pb.CreateLocationEnergySourceResponse, error) {
+	l := zerolog.Ctx(ctx)
+	querier := db.New(ix.GetTxFromContext(ctx))
+
+	locationUuid := uuid.MustParse(req.LocationUuid)
+
+	if req.ValidFromUtc == nil {
+		req.ValidFromUtc = timestamppb.New(time.Now().UTC().Truncate(time.Minute))
+	}
+
+	// Reject if this energy source already exists for the location.
+	cseprms := db.CheckSourceExistsParams{
+		GeometryUuid: locationUuid,
+		SourceTypeID: int16(req.EnergySource.Number()),
+	}
+
+	exists, err := querier.CheckSourceExists(ctx, cseprms)
 	if err != nil {
-		return nil, fmt.Errorf("failed to update sources materialised view: %w", err)
+		return nil, fmt.Errorf("failed to check source existence: %w", err)
 	}
 
-	l.Debug().Msg("refreshed sources materialised view")
+	if exists {
+		return nil, status.Errorf(
+			codes.AlreadyExists,
+			"energy source '%s' has already been created for location '%s'. New inserts to an existing energy source must go through the UpdateLocation RPC.",
+			req.EnergySource,
+			req.LocationUuid,
+		)
+	}
 
-	return &pb.CreateLocationResponse{
-		LocationUuid:           dbLocation.GeometryUuid.String(),
-		LocationName:           dbLocation.GeometryName,
+	dbSource, err := createLocationSource(
+		ctx, querier, locationUuid, int16(req.EnergySource.Number()),
+		req.EffectiveCapacityWatts, req.Metadata, req.ValidFromUtc.AsTime(),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	l.Debug().
+		Str("dp.source.geometry_uuid", locationUuid.String()).
+		Int16("dp.source.type_id", int16(req.EnergySource.Number())).
+		Int64("dp.source.capacity", int64(dbSource.CapacityWatts)).
+		Msg("created new energy source for location")
+
+	return &pb.CreateLocationEnergySourceResponse{
+		LocationUuid:           locationUuid.String(),
+		EnergySource:           req.EnergySource,
 		EffectiveCapacityWatts: uint64(dbSource.CapacityWatts),
 	}, nil
 }

internal/server/postgres/dataserverimpl_test.go

@@ -331,6 +331,108 @@ func TestCreateLocation(t *testing.T) {
 	})
 }
 
+func TestCreateLocationEnergySource(t *testing.T) {
+	metadata := createTestMetadata(t, map[string]any{"source": "test_energy_source"})
+
+	// Create a location with solar valid from today
+	createResp := createTestLocation(
+		t,
+		"scotland_region",
+		"POINT(-0.127 51.507)",
+		1230,
+		time.Now().UTC().Truncate(time.Minute),
+		metadata,
+	)
+
+	testcases := []struct {
+		name         string
+		req          *pb.CreateLocationEnergySourceRequest
+		shouldErr    bool
+		errCodeCheck codes.Code
+	}{
+		{
+			name: "Should add wind energy source to existing location",
+			req: &pb.CreateLocationEnergySourceRequest{
+				LocationUuid:           createResp.LocationUuid,
+				EnergySource:           pb.EnergySource_ENERGY_SOURCE_WIND,
+				EffectiveCapacityWatts: 5000,
+				Metadata:               metadata,
+			},
+			shouldErr: false,
+		},
+		{
+			name: "Should fail to add already existing energy source",
+			req: &pb.CreateLocationEnergySourceRequest{
+				LocationUuid:           createResp.LocationUuid,
+				EnergySource:           pb.EnergySource_ENERGY_SOURCE_SOLAR, // Already exists
+				EffectiveCapacityWatts: 2000,
+				Metadata:               metadata,
+			},
+			shouldErr:    true,
+			errCodeCheck: codes.AlreadyExists,
+		},
+		{
+			name: "Should fail if the energy source already exists but we query a past valid_from",
+			req: &pb.CreateLocationEnergySourceRequest{
+				LocationUuid:           createResp.LocationUuid,
+				EnergySource:           pb.EnergySource_ENERGY_SOURCE_SOLAR, // Created today
+				EffectiveCapacityWatts: 2000,
+				ValidFromUtc: timestamppb.New(
+					time.Now().UTC().Add(-24 * time.Hour).Truncate(time.Minute),
+				),
+				Metadata: metadata,
+			},
+			shouldErr:    true,
+			errCodeCheck: codes.AlreadyExists,
+		},
+		{
+			name: "Should fail for non-existent location",
+			req: &pb.CreateLocationEnergySourceRequest{
+				LocationUuid:           uuid.New().String(),
+				EnergySource:           pb.EnergySource_ENERGY_SOURCE_WIND,
+				EffectiveCapacityWatts: 5000,
+				Metadata:               metadata,
+			},
+			shouldErr: true,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			resp, err := dc.CreateLocationEnergySource(t.Context(), tc.req)
+
+			if tc.shouldErr {
+				require.Error(t, err, "Expected an error")
+
+				if tc.errCodeCheck != codes.OK {
+					require.Equal(t, tc.errCodeCheck, status.Code(err))
+				}
+			} else {
+				require.NoError(t, err, "Expected to be able to create the location energy source")
+				require.Equal(t, tc.req.LocationUuid, resp.LocationUuid)
+				require.Equal(t, tc.req.EnergySource, resp.EnergySource)
+				require.Equal(t, tc.req.EffectiveCapacityWatts, resp.EffectiveCapacityWatts)
+
+				// Verify it can be fetched
+				resp2, err := dc.GetLocation(
+					t.Context(),
+					&pb.GetLocationRequest{
+						LocationUuid:    tc.req.LocationUuid,
+						EnergySource:    tc.req.EnergySource,
+						IncludeGeometry: false,
+					},
+				)
+				require.NoError(
+					t,
+					err,
+					"Expected to be able to fetch the newly created energy source",
+				)
+				require.Equal(t, tc.req.EffectiveCapacityWatts, resp2.EffectiveCapacityWatts)
+			}
+		})
+	}
+}
+
 func TestUpdateLocation(t *testing.T) {
 	metadata := createTestMetadata(t, map[string]any{"source": "test"})
 

internal/server/postgres/sql/queries/locations.sql

@@ -111,6 +111,15 @@ WHERE
     AND s.source_type_id = $2
     AND s.sys_period @> sqlc.arg(at_timestamp_utc)::TIMESTAMP;
 
+-- name: CheckSourceExists :one
+/* CheckSourceExists returns true if the given geometry and source type has ever had an entry. */
+SELECT EXISTS(
+    SELECT 1
+    FROM loc.sources_history
+    WHERE geometry_uuid = $1
+        AND source_type_id = $2
+);
+
 -- name: CreateSourceEntry :one
 /* CreateSourceEntry creates a new source entry for a given geometry and source type.
  * It fetches the state prior to the input valid time, and only inserts the new row if it differs

proto/ocf/dp/dp-data.messages.proto

@@ -469,6 +469,32 @@ message CreateLocationResponse {
 }
 
 
+message CreateLocationEnergySourceRequest {
+  string location_uuid = 1 [
+    (buf.validate.field).required = true,
+    (buf.validate.field).string.uuid = true
+  ];
+  EnergySource energy_source = 2 [
+    (buf.validate.field).required = true
+  ];
+  // The effective capacity of this source in watts.
+  uint64 effective_capacity_watts = 3 [
+    (buf.validate.field).required = true
+  ];
+  optional google.protobuf.Struct metadata = 4;
+  // The UTC time from which this source is considered valid. Leave empty to use current time.
+  optional google.protobuf.Timestamp valid_from_utc = 5 [
+    (buf.validate.field).timestamp = { gt: { seconds: 112000000}, lt_now: true }
+  ];
+}
+
+message CreateLocationEnergySourceResponse {
+  string location_uuid = 1;
+  EnergySource energy_source = 2;
+  uint64 effective_capacity_watts = 3;
+}
+
+
 message GetLocationRequest {
   string location_uuid = 1 [
     (buf.validate.field).required = true,

proto/ocf/dp/dp-data.service.proto

@@ -34,6 +34,10 @@ service DataPlatformDataService {
   rpc GetLocationAsTimeseries(GetLocationAsTimeseriesRequest) returns (GetLocationAsTimeseriesResponse) {}
   /* CreateLocation registers a new location in which to log or forecast generation. */
   rpc CreateLocation(CreateLocationRequest) returns (CreateLocationResponse) {}
+  /* CreateLocationEnergySource attaches a new energy source to an existing location.
+   * Errors if the location already has a source of this type.
+   */
+  rpc CreateLocationEnergySource(CreateLocationEnergySourceRequest) returns (CreateLocationEnergySourceResponse) {}
   /* UpdateLocation modifies various attributes associated with a given location. */
   rpc UpdateLocation(UpdateLocationRequest) returns (UpdateLocationResponse) {}
   /* UpdateLocationOwner changes the ownership of a location. */