diff --git a/internal/db/models/training_dataset.go b/internal/db/models/training_dataset.go index fb95001..e0423fd 100644 --- a/internal/db/models/training_dataset.go +++ b/internal/db/models/training_dataset.go @@ -11,6 +11,7 @@ type TrainingDataset struct { gorm.Model Name string `gorm:"type:varchar(255);not null;uniqueIndex"` AODFiles []jalien.AODFile `gorm:"serializer:json"` + IsONe bool AnchorProdTimestampFrom time.Time AnchorProdTimestampTo time.Time UserId uint diff --git a/internal/handler/queue_handler.go b/internal/handler/queue_handler.go index 978432f..12e402b 100644 --- a/internal/handler/queue_handler.go +++ b/internal/handler/queue_handler.go @@ -103,10 +103,14 @@ func (qh *QueueHandler) QueryTask(w http.ResponseWriter, r *http.Request) { response := struct { ID uint AODFiles []jalien.AODFile + IsONe bool + IsData bool Configuration interface{} }{ ID: tt.ID, AODFiles: tt.TrainingDataset.AODFiles, + IsONe: tt.TrainingDataset.IsONe, + IsData: false, // For now we don't use any real experimental data (simulated data only). Configuration: tt.Configuration, } diff --git a/internal/service/training_dataset_service.go b/internal/service/training_dataset_service.go index 97568a1..147479e 100644 --- a/internal/service/training_dataset_service.go +++ b/internal/service/training_dataset_service.go @@ -94,6 +94,16 @@ func (s *TrainingDatasetService) Create(td *models.TrainingDataset) error { td.AnchorProdTimestampFrom = timestampFrom td.AnchorProdTimestampTo = timestampTo + isONe, err := resolveDatasetIsONe(td.AODFiles, s.MonalisaService) + if err != nil { + var external *ErrExternalServiceTimeout + if errors.As(err, &external) { + return err + } + return InternalError("resolve training dataset IsONe", err) + } + td.IsONe = isONe + err = s.TrainingDataset.Create(td) if err != nil { @@ -154,3 +164,22 @@ func (s *TrainingDatasetService) FindAods(path string) ([]jalien.AODFile, error) return aodFiles, nil } + +func resolveDatasetIsONe(aodFiles []jalien.AODFile, monalisaService IMonalisaService) (bool, error) { + runList, err := monalisaService.GetRunList() + if err != nil { + return false, mapMonalisaError("GetRunList", err) + } + + for _, aod := range aodFiles { + if aod.LHCPeriod != "" { + desc := runList.TagToDesc[aod.LHCPeriod] + if strings.Contains(desc, "O-O") || strings.Contains(desc, "OO") || + strings.Contains(desc, "Ne-Ne") || strings.Contains(desc, "NeNe") { + return true, nil + } + } + } + + return false, nil +} diff --git a/test/integration/training_dataset_handler_test.go b/test/integration/training_dataset_handler_test.go index d3e48c3..52e5057 100644 --- a/test/integration/training_dataset_handler_test.go +++ b/test/integration/training_dataset_handler_test.go @@ -173,7 +173,7 @@ func mockDatasetTimestampRange(ut *IntegrationTestUtils, period string, minRun, TagToRuns: map[string][]uint64{ period: {minRun, maxRun}, }, - }, nil).Once() + }, nil) ut.CCDB.On("GetRunInformation", minRun).Return(&ccdb.RunInformation{ RunNumber: minRun, SOR: uint64(time.Now().Add(-time.Hour).UnixMilli()), diff --git a/test/repository/training_dataset_repository_test.go b/test/repository/training_dataset_repository_test.go index a8edb13..264de58 100644 --- a/test/repository/training_dataset_repository_test.go +++ b/test/repository/training_dataset_repository_test.go @@ -40,7 +40,7 @@ func TestTrainingDatasetRepository_Create(t *testing.T) { mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "training_datasets" (.+) RETURNING "id"`). - WithArgs(AnyTime(), AnyTime(), AnyTime(), trainingDataset.Name, marshalAODFiles(t, trainingDataset), AnyTime(), AnyTime(), 1). + WithArgs(AnyTime(), AnyTime(), AnyTime(), trainingDataset.Name, marshalAODFiles(t, trainingDataset), false, AnyTime(), AnyTime(), 1). WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1)) mock.ExpectCommit() diff --git a/test/repository/training_task_repository_test.go b/test/repository/training_task_repository_test.go index 6776241..017667a 100644 --- a/test/repository/training_task_repository_test.go +++ b/test/repository/training_task_repository_test.go @@ -48,7 +48,7 @@ func TestTrainingTaskRepository_Create(t *testing.T) { mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "training_datasets" (.+) RETURNING "id"`). - WithArgs(AnyTime(), AnyTime(), AnyTime(), trainingDataset.Name, marshalAODFiles(t, trainingDataset), AnyTime(), AnyTime(), 1). + WithArgs(AnyTime(), AnyTime(), AnyTime(), trainingDataset.Name, marshalAODFiles(t, trainingDataset), false, AnyTime(), AnyTime(), 1). WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1)) mock.ExpectQuery(`INSERT INTO "training_tasks" (.+) RETURNING "id"`). WithArgs(AnyTime(), AnyTime(), AnyTime(), trainingTask.Name, trainingTask.Status, 1, 1, nil, marshalTrainingTaskConfig(t, trainingTask)).