pull tables form settings and update test data

src/main/java/bio/terra/service/snapshot/SnapshotService.java

@@ -34,6 +34,7 @@
 import bio.terra.model.SamPolicyModel;
 import bio.terra.model.SnapshotAccessRequestResponse;
 import bio.terra.model.SnapshotAccessRequestStatus;
+import bio.terra.model.SnapshotBuilderDatasetConceptSet;
 import bio.terra.model.SnapshotBuilderOutputTable;
 import bio.terra.model.SnapshotBuilderSettings;
 import bio.terra.model.SnapshotBuilderTable;
@@ -688,7 +689,7 @@ public SnapshotAccessRequestResponse getSnapshotAccessRequestById(UUID accessReq
   }
 
   public AssetSpecification buildAssetFromSnapshotAccessRequest(
-      Dataset dataset, SnapshotAccessRequestResponse snapshotRequestModel) {
+      Dataset dataset, SnapshotAccessRequestResponse snapshotAccessRequest) {
     // build asset model from snapshot request
     AssetModel assetModel =
         new AssetModel()
@@ -699,7 +700,7 @@ public AssetSpecification buildAssetFromSnapshotAccessRequest(
     assetModel.addTablesItem(new AssetTableModel().name("person"));
 
     // Build asset tables and columns based on the concept sets included in the snapshot request
-    List<SnapshotBuilderTable> tables = pullTables(snapshotRequestModel);
+    List<SnapshotBuilderTable> tables = pullTables(snapshotAccessRequest);
     tables.forEach(
         table -> {
           assetModel.addTablesItem(
@@ -721,116 +722,29 @@ public AssetSpecification buildAssetFromSnapshotAccessRequest(
   }
 
   @VisibleForTesting
-  List<SnapshotBuilderTable> pullTables(SnapshotAccessRequestResponse snapshotRequestModel) {
-    var tables = snapshotRequestModel.getSnapshotSpecification().getOutputTables();
-    var tableNames = tables.stream().map(SnapshotBuilderOutputTable::getName).toList();
-
-    Map<String, SnapshotBuilderTable> tableMap = populateManualTableMap();
-
-    Set<String> missing = new HashSet<>(tableNames);
-    missing.removeAll(tableMap.keySet());
+  List<SnapshotBuilderTable> pullTables(SnapshotAccessRequestResponse snapshotAccessRequest) {
+    List<String> includedTableNames =
+        snapshotAccessRequest.getSnapshotSpecification().getOutputTables().stream()
+            .map(SnapshotBuilderOutputTable::getName)
+            .toList();
+    List<SnapshotBuilderDatasetConceptSet> allTables =
+        snapshotBuilderSettingsDao
+            .getBySnapshotId(snapshotAccessRequest.getSourceSnapshotId())
+            .getDatasetConceptSets();
+
+    Set<String> missing = new HashSet<>(includedTableNames);
+    allTables.stream()
+        .map(SnapshotBuilderDatasetConceptSet::getName)
+        .toList()
+        .forEach(missing::remove);
     if (!missing.isEmpty()) {
       throw new IllegalArgumentException("Unknown value set names: " + missing);
     }
 
-    return tableNames.stream().map(tableMap::get).toList();
-  }
-
-  private Map<String, SnapshotBuilderTable> populateManualTableMap() {
-    // manual definition of domain names -> dataset table
-    Map<String, SnapshotBuilderTable> tableMap = new HashMap<>();
-    tableMap.put(
-        "Demographics",
-        new SnapshotBuilderTable()
-            .datasetTableName("person")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_person_gender_concept",
-                    "fpk_person_race_concept",
-                    "fpk_person_ethnicity_concept",
-                    "fpk_person_gender_concept_s",
-                    "fpk_person_race_concept_s",
-                    "fpk_person_ethnicity_concept_s")));
-    tableMap.put(
-        "Drug",
-        new SnapshotBuilderTable()
-            .datasetTableName("drug_exposure")
-            .primaryTableRelationship("fpk_person_drug")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_drug_concept",
-                    "fpk_drug_type_concept",
-                    "fpk_drug_route_concept",
-                    "fpk_drug_concept_s")));
-    tableMap.put(
-        "Measurement",
-        new SnapshotBuilderTable()
-            .datasetTableName("measurement")
-            .primaryTableRelationship("fpk_person_measurement")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_measurement_concept",
-                    "fpk_measurement_unit",
-                    "fpk_measurement_concept_s",
-                    "fpk_measurement_value",
-                    "fpk_measurement_type_concept",
-                    "fpk_measurement_operator")));
-    tableMap.put(
-        "Visit",
-        new SnapshotBuilderTable()
-            .datasetTableName("visit_occurrence")
-            .primaryTableRelationship("fpk_person_visit")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_visit_preceding",
-                    "fpk_visit_concept_s",
-                    "fpk_visit_type_concept",
-                    "fpk_visit_concept",
-                    "fpk_visit_discharge")));
-    tableMap.put(
-        "Device",
-        new SnapshotBuilderTable()
-            .datasetTableName("device_exposure")
-            .primaryTableRelationship("fpk_person_device")
-            .secondaryTableRelationships(
-                List.of("fpk_device_concept", "fpk_device_concept_s", "fpk_device_type_concept")));
-    tableMap.put(
-        "Condition",
-        new SnapshotBuilderTable()
-            .datasetTableName("condition_occurrence")
-            .primaryTableRelationship("fpk_person_condition")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_condition_concept",
-                    "fpk_condition_type_concept",
-                    "fpk_condition_status_concept",
-                    "fpk_condition_concept_s")));
-    tableMap.put(
-        "Procedure",
-        new SnapshotBuilderTable()
-            .datasetTableName("procedure_occurrence")
-            .primaryTableRelationship("fpk_person_procedure")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_procedure_concept",
-                    "fpk_procedure_concept_s",
-                    "fpk_procedure_type_concept",
-                    "fpk_procedure_modifier")));
-    tableMap.put(
-        "Observation",
-        new SnapshotBuilderTable()
-            .datasetTableName("observation")
-            .primaryTableRelationship("fpk_person_observation")
-            .secondaryTableRelationships(
-                List.of(
-                    "fpk_observation_concept",
-                    "fpk_observation_concept_s",
-                    "fpk_observation_unit",
-                    "fpk_observation_qualifier",
-                    "fpk_observation_type_concept",
-                    "fpk_observation_value")));
-
-    return tableMap;
+    return allTables.stream()
+        .filter((table) -> includedTableNames.contains(table.getName()))
+        .map(SnapshotBuilderDatasetConceptSet::getTable)
+        .toList();
   }
 
   public static AssetSpecification getAssetByNameFromDataset(Dataset dataset, String assetName) {

src/test/java/bio/terra/service/snapshot/SnapshotServiceTest.java

@@ -160,6 +160,7 @@ class SnapshotServiceTest {
   @Mock private EcmService ecmService;
   @Mock private RawlsService rawlsService;
   @Mock private DuosClient duosClient;
+  @Mock private SnapshotBuilderSettingsDao settingsDao;
   private final UUID snapshotId = UUID.randomUUID();
   private final UUID datasetId = UUID.randomUUID();
   private final UUID snapshotTableId = UUID.randomUUID();
@@ -187,7 +188,7 @@ void beforeEach() {
             azureSynapsePdao,
             rawlsService,
             duosClient,
-            mock(SnapshotBuilderSettingsDao.class));
+            settingsDao);
   }
 
   @Test
@@ -1094,15 +1095,15 @@ void testCreateSnapshotWithByRequestId() {
     SnapshotAccessRequestResponse snapshotAccessRequestResponse =
         new SnapshotAccessRequestResponse()
             .status(SnapshotAccessRequestStatus.APPROVED)
-            .id(snapshotAccessRequestId);
+            .id(snapshotAccessRequestId)
+            .sourceSnapshotId(UUID.randomUUID());
     when(snapshotRequestDao.getById(snapshotAccessRequestId))
         .thenReturn(snapshotAccessRequestResponse);
     when(snapshotDao.retrieveSnapshot(snapshotAccessRequestResponse.getSourceSnapshotId()))
         .thenReturn(
             new Snapshot()
                 .snapshotSources(
                     List.of(new SnapshotSource().dataset(new Dataset().id(UUID.randomUUID())))));
-
     String result =
         service.createSnapshot(
             request, service.getSourceDatasetFromSnapshotRequest(request), TEST_USER);
@@ -1333,13 +1334,17 @@ void makeSnapshotFromSnapshotRequestByRequestId() {
 
     Dataset dataset = SnapshotBuilderTestData.DATASET;
     dataset.name("datasetName");
+    UUID sourceSnapshotId = UUID.randomUUID();
     when(snapshotRequestDao.getById(snapshotAccessRequestId))
         .thenReturn(
             new SnapshotAccessRequestResponse()
+                .sourceSnapshotId(sourceSnapshotId)
                 .snapshotSpecification(
                     new SnapshotBuilderRequest()
                         .addOutputTablesItem(new SnapshotBuilderOutputTable().name("Drug"))
                         .addOutputTablesItem(new SnapshotBuilderOutputTable().name("Condition"))));
+    when(settingsDao.getBySnapshotId(sourceSnapshotId))
+        .thenReturn(SnapshotBuilderTestData.SETTINGS);
 
     Snapshot actual = service.makeSnapshotFromSnapshotRequest(snapshotRequestModel, dataset);
     SnapshotSource snapshotSource = new SnapshotSource().dataset(dataset);
@@ -1453,6 +1458,8 @@ void pullTables() {
     var accessRequestResponse =
         SnapshotBuilderTestData.createSnapshotAccessRequestResponse(sourceSnapshotId);
     accessRequestResponse.id(snapshotAccessRequestId);
+    when(settingsDao.getBySnapshotId(sourceSnapshotId))
+        .thenReturn(SnapshotBuilderTestData.SETTINGS);
     var firstTable = service.pullTables(accessRequestResponse).get(0);
     assertThat(firstTable.getDatasetTableName(), is("drug_exposure"));
     // Must preserve relationship order
@@ -1477,6 +1484,7 @@ void buildAssetFromSnapshotAccessRequest() {
     var accessRequestResponse =
         SnapshotBuilderTestData.createSnapshotAccessRequestResponse(snapshotId);
     accessRequestResponse.id(snapshotAccessRequestId);
+    when(settingsDao.getBySnapshotId(snapshotId)).thenReturn(SnapshotBuilderTestData.SETTINGS);
 
     var actualAssetSpec =
         service.buildAssetFromSnapshotAccessRequest(sourceDataset, accessRequestResponse);

src/test/java/bio/terra/service/snapshotbuilder/SnapshotBuilderTestData.java

@@ -22,6 +22,7 @@
 import bio.terra.model.SnapshotBuilderProgramDataRangeOption;
 import bio.terra.model.SnapshotBuilderRequest;
 import bio.terra.model.SnapshotBuilderSettings;
+import bio.terra.model.SnapshotBuilderTable;
 import bio.terra.model.SnapshotRequestContentsModel;
 import bio.terra.model.SnapshotRequestIdModel;
 import bio.terra.model.SnapshotRequestModel;
@@ -33,7 +34,8 @@
 import bio.terra.service.snapshotbuilder.query.table.Person;
 import bio.terra.service.snapshotbuilder.utils.constants.ConditionOccurrence;
 import bio.terra.service.snapshotbuilder.utils.constants.DrugExposure;
-import bio.terra.service.snapshotbuilder.utils.constants.ObservationOccurrence;
+import bio.terra.service.snapshotbuilder.utils.constants.Measurement;
+import bio.terra.service.snapshotbuilder.utils.constants.Observation;
 import bio.terra.service.snapshotbuilder.utils.constants.ProcedureOccurrence;
 import java.util.List;
 import java.util.UUID;
@@ -122,8 +124,8 @@ private static SnapshotBuilderProgramDataListOption generateSnapshotBuilderProgr
                           .hasChildren(true)),
                   generateSnapshotBuilderDomainOption(
                       OBSERVATION_DOMAIN_ID,
-                      ObservationOccurrence.TABLE_NAME,
-                      ObservationOccurrence.OBSERVATION_CONCEPT_ID,
+                      Observation.TABLE_NAME,
+                      Observation.OBSERVATION_CONCEPT_ID,
                       "Observation",
                       new SnapshotBuilderConcept()
                           .id(300)
@@ -170,8 +172,53 @@ private static SnapshotBuilderProgramDataListOption generateSnapshotBuilderProgr
                       List.of(new SnapshotBuilderProgramDataListItem().id(43).name("unused 3")))))
           .datasetConceptSets(
               List.of(
-                  new SnapshotBuilderDatasetConceptSet().name("Demographics"),
-                  new SnapshotBuilderDatasetConceptSet().name("All surveys")));
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Drug")
+                      .table(
+                          new SnapshotBuilderTable()
+                              .datasetTableName(DrugExposure.TABLE_NAME)
+                              .primaryTableRelationship("fpk_person_drug")
+                              .secondaryTableRelationships(
+                                  List.of(
+                                      "fpk_drug_concept",
+                                      "fpk_drug_type_concept",
+                                      "fpk_drug_route_concept",
+                                      "fpk_drug_concept_s"))),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Condition")
+                      .table(
+                          new SnapshotBuilderTable()
+                              .datasetTableName(ConditionOccurrence.TABLE_NAME)
+                              .primaryTableRelationship("fpk_person_condition")
+                              .secondaryTableRelationships(
+                                  List.of(
+                                      "fpk_condition_concept",
+                                      "fpk_condition_type_concept",
+                                      "fpk_condition_status_concept",
+                                      "fpk_condition_concept_s"))),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Procedure")
+                      .table(
+                          new SnapshotBuilderTable()
+                              .datasetTableName(ProcedureOccurrence.TABLE_NAME)),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Observation")
+                      .table(new SnapshotBuilderTable().datasetTableName(Observation.TABLE_NAME)),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Measurement")
+                      .table(new SnapshotBuilderTable().datasetTableName(Measurement.TABLE_NAME)),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Visit")
+                      .table(new SnapshotBuilderTable().datasetTableName("visit_occurrence")),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Device")
+                      .table(new SnapshotBuilderTable().datasetTableName("device_exposure")),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Demographics")
+                      .table(new SnapshotBuilderTable().datasetTableName(Person.TABLE_NAME)),
+                  new SnapshotBuilderDatasetConceptSet()
+                      .name("Genomics")
+                      .table(new SnapshotBuilderTable().datasetTableName("sample"))));
 
   public static final Column PERSON_ID_COLUMN =
       new Column().name("person_id").type(TableDataType.INTEGER);

src/test/java/bio/terra/service/snapshotbuilder/utils/constants/Measurement.java

@@ -0,0 +1,8 @@
+package bio.terra.service.snapshotbuilder.utils.constants;
+
+public class Measurement {
+  private Measurement() {}
+
+  public static final String TABLE_NAME = "measurement";
+  public static final String MEASUREMENT_CONCEPT_ID = "measurement_concept_id";
+}

src/test/java/bio/terra/service/snapshotbuilder/utils/constants/Observation.java

@@ -0,0 +1,8 @@
+package bio.terra.service.snapshotbuilder.utils.constants;
+
+public final class Observation {
+  private Observation() {}
+
+  public static final String TABLE_NAME = "observation";
+  public static final String OBSERVATION_CONCEPT_ID = "observation_concept_id";
+}

src/test/resources/omop/settings.json

@@ -241,6 +241,7 @@
         "primaryTableRelationship": "fpk_condition_person",
         "secondaryTableRelationships": [
           "fpk_condition_concept",
+          "fpk_condition_type_concept",
           "fpk_condition_status_concept",
           "fpk_condition_concept_s"]
       }

tools/setupResourceScripts/files/OMOPDataset/snapshot_builder_settings.json

@@ -241,6 +241,7 @@
         "primaryTableRelationship": "fpk_condition_person",
         "secondaryTableRelationships": [
           "fpk_condition_concept",
+          "fpk_condition_type_concept",
           "fpk_condition_status_concept",
           "fpk_condition_concept_s"]
       }