[CTM-248] Defer/skip adding auth domains for (certain) PFBs (#1075)

Author: calypsomatic

service/src/main/java/org/databiosphere/workspacedataservice/config/DataImportProperties.java

@@ -153,7 +153,8 @@ public record ImportSourceConfig(
       List<Pattern> urls,
       boolean requirePrivateWorkspace,
       boolean requireProtectedDataPolicy,
-      List<String> requiredAuthDomainGroups) {
+      List<String> requiredAuthDomainGroups,
+      boolean alwaysApplyAuthDomains) {
     public boolean matchesUri(URI uri) {
       String uriString = uri.toString();
       return urls.stream().anyMatch(urlPattern -> urlPattern.matcher(uriString).find());

service/src/main/java/org/databiosphere/workspacedataservice/dataimport/DefaultImportValidator.java

@@ -152,7 +152,9 @@ private void validateDestinationWorkspace(
           "Data from this source cannot be imported into a public workspace.");
     }
 
-    if (!requirements.requiredAuthDomainGroups().isEmpty()) {
+    // Apply auth domain groups if present and always applied, else defer
+    if (requirements.alwaysApplyAuthDomains()
+        && !requirements.requiredAuthDomainGroups().isEmpty()) {
       protectedDataSupport.addAuthDomainGroupsToWorkspace(
           destinationWorkspaceId, requirements.requiredAuthDomainGroups());
     }

service/src/main/java/org/databiosphere/workspacedataservice/dataimport/ImportRequirements.java

@@ -3,8 +3,11 @@
 import java.util.List;
 
 public record ImportRequirements(
-    boolean privateWorkspace, boolean protectedDataPolicy, List<String> requiredAuthDomainGroups) {
+    boolean privateWorkspace,
+    boolean protectedDataPolicy,
+    List<String> requiredAuthDomainGroups,
+    boolean alwaysApplyAuthDomains) {
   public ImportRequirements() {
-    this(false, false, List.of());
+    this(false, false, List.of(), true);
   }
 }

service/src/main/java/org/databiosphere/workspacedataservice/dataimport/ImportRequirementsFactory.java

@@ -32,7 +32,16 @@ public ImportRequirements getRequirementsForImport(URI importUri) {
                 .flatMap(source -> source.requiredAuthDomainGroups().stream())
                 .toList();
 
+    boolean alwaysApplyAuthDomains =
+        sources != null
+            && sources.stream()
+                .filter(source -> source.matchesUri(importUri))
+                .anyMatch(ImportSourceConfig::alwaysApplyAuthDomains);
+
     return new ImportRequirements(
-        requiresPrivateWorkspace, requiresProtectedDataPolicy, requiredAuthDomainGroups);
+        requiresPrivateWorkspace,
+        requiresProtectedDataPolicy,
+        requiredAuthDomainGroups,
+        alwaysApplyAuthDomains);
   }
 }

service/src/main/java/org/databiosphere/workspacedataservice/dataimport/pfb/PfbQuartzJob.java

@@ -11,6 +11,7 @@
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import java.net.URI;
+import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 import java.util.Spliterator;
@@ -26,6 +27,9 @@
 import org.databiosphere.workspacedataservice.dao.JobDao;
 import org.databiosphere.workspacedataservice.dataimport.ImportDetails;
 import org.databiosphere.workspacedataservice.dataimport.ImportDetailsRetriever;
+import org.databiosphere.workspacedataservice.dataimport.ImportRequirements;
+import org.databiosphere.workspacedataservice.dataimport.ImportRequirementsFactory;
+import org.databiosphere.workspacedataservice.dataimport.protecteddatasupport.ProtectedDataSupport;
 import org.databiosphere.workspacedataservice.dataimport.snapshotsupport.MultiCloudSnapshotSupportFactory;
 import org.databiosphere.workspacedataservice.dataimport.snapshotsupport.SnapshotSupport;
 import org.databiosphere.workspacedataservice.jobexec.JobDataMapReader;
@@ -66,6 +70,8 @@ public class PfbQuartzJob extends QuartzJob {
   private final ImportDetailsRetriever importDetailsRetriever;
   private final ImportMetrics importMetrics;
   private final DrsService drsService;
+  private final ProtectedDataSupport protectedDataSupport;
+  private final ImportRequirementsFactory importRequirementsFactory;
 
   public PfbQuartzJob(
       JobDao jobDao,
@@ -78,7 +84,8 @@ public PfbQuartzJob(
       MultiCloudSnapshotSupportFactory snapshotSupportFactory,
       DataImportProperties dataImportProperties,
       ImportDetailsRetriever importDetailsRetriever,
-      DrsService drsService) {
+      DrsService drsService,
+      ProtectedDataSupport protectedDataSupport) {
     super(jobDao, observationRegistry, dataImportProperties);
     this.recordSourceFactory = recordSourceFactory;
     this.recordSinkFactory = recordSinkFactory;
@@ -88,6 +95,9 @@ public PfbQuartzJob(
     this.importDetailsRetriever = importDetailsRetriever;
     this.importMetrics = importMetrics;
     this.drsService = drsService;
+    this.protectedDataSupport = protectedDataSupport;
+    this.importRequirementsFactory =
+        new ImportRequirementsFactory(dataImportProperties.getSources());
   }
 
   @Override
@@ -109,15 +119,36 @@ protected void executeInternal(UUID jobId, JobExecutionContext context) {
 
     ImportDetails details = importDetailsRetriever.fetch(jobId, jobData, PrefixStrategy.PFB);
 
+    // Early PFB content inspection to determine if auth domains should be applied
+    // This is HTTP connection #0 to the PFB.
+    logger.info("Inspecting PFB content for auth domain requirements...");
+    Set<UUID> snapshotIds = withPfbStream(uri, this::findSnapshots);
+    logger.info("Found {} unique snapshot IDs in PFB", snapshotIds.size());
+    boolean hasNresConsent = false;
+    if (snapshotIds.isEmpty()) {
+      hasNresConsent = withPfbStream(uri, this::hasNresConsentGroup);
+      logger.info("NRES consent group present in PFB: {}", hasNresConsent);
+    }
+    ImportRequirements requirements = importRequirementsFactory.getRequirementsForImport(uri);
+
+    // If we skipped auth domains earlier, but now there is neither an NRES consent group nor any
+    // snapshots,
+    // apply the configured auth domains now.
+    if (!requirements.requiredAuthDomainGroups().isEmpty()
+        && !requirements.alwaysApplyAuthDomains()
+        && snapshotIds.isEmpty()
+        && !hasNresConsent) {
+      logger.info("Applying auth domain groups based on PFB content analysis...");
+      protectedDataSupport.addAuthDomainGroupsToWorkspace(
+          details.workspaceId(), requirements.requiredAuthDomainGroups());
+    }
+
     // Find all the snapshot ids in the PFB, then create or verify references from the
     // workspace to the snapshot for each of those snapshot ids.
     // This will throw an exception if there are policy conflicts between the workspace
     // and the snapshots.
     //
-    // This is HTTP connection #1 to the PFB.
-    logger.info("Finding snapshots in this PFB...");
-    Set<UUID> snapshotIds = withPfbStream(uri, this::findSnapshots);
-
+    // This is HTTP connection #1 to the PFB (reusing snapshotIds from earlier).
     logger.info("Linking snapshots...");
     linkSnapshots(snapshotIds, details.workspaceId());
 
@@ -261,4 +292,49 @@ private UUID maybeUuid(String input) {
       return null;
     }
   }
+
+  /** Check if the PFB contains any anvil_dataset records with consent_group set to NRES */
+  boolean hasNresConsentGroup(DataFileStream<GenericRecord> dataStream) {
+    Stream<GenericRecord> recordStream =
+        StreamSupport.stream(
+            Spliterators.spliteratorUnknownSize(dataStream.iterator(), Spliterator.ORDERED), false);
+
+    // Collect all consent groups from anvil_dataset records
+    List<Object> consentGroups =
+        recordStream
+            .filter(rec -> "anvil_dataset".equals(rec.get("name").toString()))
+            .map(rec -> rec.get("object"))
+            .filter(GenericRecord.class::isInstance)
+            .map(GenericRecord.class::cast)
+            .filter(anvilDataset -> anvilDataset.hasField("consent_group"))
+            .map(anvilDataset -> anvilDataset.get("consent_group"))
+            .filter(Objects::nonNull)
+            .filter(
+                consentGroup -> {
+                  if (consentGroup instanceof java.util.Collection) {
+                    return !((java.util.Collection<?>) consentGroup).isEmpty();
+                  } else {
+                    return !consentGroup.toString().isEmpty();
+                  }
+                })
+            .toList();
+
+    // Must have at least one consent group AND all must be NRES
+    if (consentGroups.isEmpty()) {
+      return false;
+    }
+
+    return consentGroups.stream()
+        .allMatch(
+            consentGroup -> {
+              if (consentGroup instanceof java.util.Collection) {
+                return ((java.util.Collection<?>) consentGroup)
+                    .stream()
+                        .filter(Objects::nonNull)
+                        .allMatch(item -> "NRES".equals(String.valueOf(item)));
+              } else {
+                return "NRES".equals(String.valueOf(consentGroup));
+              }
+            });
+  }
 }

service/src/main/resources/application-prod.yml

@@ -19,7 +19,6 @@ twds:
           - ^https:\/\/s3\.amazonaws.com\/gen3-biodatacatalyst-nhlbi-nih-gov-pfb-export\/
           - ^https:\/\/pic-sure-auth-prod-data-export\.s3\.amazonaws\.com\/
           - ^https:\/\/s3\.amazonaws\.com\/pic-sure-auth-prod-data-export\/
-          - ^https:\/\/s3\.amazonaws\.com\/edu-ucsc-gi-platform-anvil-prod-storage-anvilprod\.us-east-1\/
           - ^https:\/\/nih-nhlbi-.*\.amazonaws\.com\/
           # anvil data is in TDR now, these are the old URLs
           - ^https:\/\/gen3-theanvil-io-pfb-export\.s3\.amazonaws\.com\/
@@ -28,12 +27,20 @@ twds:
         requireProtectedDataPolicy: true
         requiredAuthDomainGroups:
           - "federal_data_lockdown"
+        alwaysApplyAuthDomains: true
       - urls:
           - ^https:\/\/storage\.googleapis\.com\/datarepo-.*-snapshot-export-bucket
           - ^https:\/\/s3\.amazonaws\.com\/edu-ucsc-gi-platform-hca-prod-storage-prod\.us-east-1
         requirePrivateWorkspace: false
         requireProtectedDataPolicy: false
         requiredAuthDomainGroups:
+      - urls:
+          - ^https:\/\/s3\.amazonaws\.com\/edu-ucsc-gi-platform-anvil-prod-storage-anvilprod\.us-east-1\/
+        requirePrivateWorkspace: true
+        requireProtectedDataPolicy: true
+        requiredAuthDomainGroups:
+          - "federal_data_lockdown"
+        alwaysApplyAuthDomains: false
 
 # Set the allowed hosts for drs pfb imports
 drs:

service/src/main/resources/application.yml

@@ -160,6 +160,7 @@ twds:
         requirePrivateWorkspace: false
         requireProtectedDataPolicy: false
         requiredAuthDomainGroups:
+        alwaysApplyAuthDomains: true
 
   pg_dump:
     path: ${PGDUMP_PATH:/usr/bin/pg_dump}

service/src/test/java/org/databiosphere/workspacedataservice/dataimport/DefaultImportValidatorTest.java

@@ -8,6 +8,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -70,24 +71,34 @@ DefaultImportValidator getDefaultImportValidatorForTest(
                   /* urls */ List.of(Pattern.compile("authdomain\\.pfb")),
                   /* requirePrivateWorkspace */ false,
                   /* requireProtectedDataPolicy */ false,
-                  /* requiredAuthDomainGroups */ List.of(authDomain)),
+                  /* requiredAuthDomainGroups */ List.of(authDomain),
+                  /* alwaysApplyAuthDomains */ true),
               new ImportSourceConfig(
                   /* urls */ List.of(Pattern.compile("protected\\.pfb")),
                   /* requirePrivateWorkspace */ false,
                   /* requireProtectedDataPolicy */ true,
-                  /* requiredAuthDomainGroups */ List.of()),
+                  /* requiredAuthDomainGroups */ List.of(),
+                  /* alwaysApplyAuthDomains */ true),
               new ImportSourceConfig(
                   /* urls */ List.of(Pattern.compile("private\\.pfb")),
                   /* requirePrivateWorkspace */ true,
                   /* requireProtectedDataPolicy */ false,
-                  /* requiredAuthDomainGroups */ List.of()),
+                  /* requiredAuthDomainGroups */ List.of(),
+                  /* alwaysApplyAuthDomains */ true),
               new ImportSourceConfig(
                   /* urls */ List.of(
                       Pattern.compile(
                           "^https:\\/\\/storage\\.googleapis\\.com/datarepo-.*-snapshot-export-bucket")),
                   /* requirePrivateWorkspace */ false,
                   /* requireProtectedDataPolicy */ false,
-                  /* requiredAuthDomainGroups */ List.of())),
+                  /* requiredAuthDomainGroups */ List.of(),
+                  /* alwaysApplyAuthDomains */ true),
+              new ImportSourceConfig(
+                  /* urls */ List.of(Pattern.compile("special-case\\.pfb")),
+                  /* requirePrivateWorkspace */ false,
+                  /* requireProtectedDataPolicy */ false,
+                  /* requiredAuthDomainGroups */ List.of("an-auth-domain"),
+                  /* alwaysApplyAuthDomains */ false)),
           /* allowedRawlsBucket */ "test-bucket",
           new NoopConnectivityChecker(),
           drsImportProperties);
@@ -272,6 +283,21 @@ void addsAuthDomains() {
         .addAuthDomainGroupsToWorkspace(destinationWorkspaceId, List.of(authDomain));
   }
 
+  // If specified to defer, do not add auth domains at validation time
+  @Test
+  void doesNotAddAuthDomainsIfDeferred() {
+    // Arrange
+    ImportRequestServerModel importRequest =
+        new ImportRequestServerModel(
+            TypeEnum.PFB, URI.create("https://files.terra.bio/special-case.pfb"));
+
+    // Act
+    importValidator.validateImport(importRequest, destinationWorkspaceId);
+
+    // Assert
+    verify(protectedDataSupport, never()).addAuthDomainGroupsToWorkspace(any(), any());
+  }
+
   @ParameterizedTest
   @MethodSource("requireProtectedWorkspacesForImportsFromConfiguredSourcesTestCases")
   void requireProtectedWorkspacesForImportsFromConfiguredSources(
@@ -325,7 +351,8 @@ void connectionFailureInvalidates(Exception ex) throws IOException {
                             "^https:\\/\\/storage\\.googleapis\\.com/datarepo-.*-snapshot-export-bucket")),
                     /* requirePrivateWorkspace */ false,
                     /* requireProtectedDataPolicy */ false,
-                    /* requiredAuthDomainGroups */ List.of())),
+                    /* requiredAuthDomainGroups */ List.of(),
+                    /* alwaysApplyAuthDomains */ true)),
             /* allowedRawlsBucket */ "test-bucket",
             mockConnectivityChecker,
             drsImportProperties);

service/src/test/java/org/databiosphere/workspacedataservice/dataimport/ImportRequirementsFactoryTest.java

@@ -45,7 +45,8 @@ void setup() {
                         "^https:\\/\\/s3\\.amazonaws\\.com\\/gen3-theanvil-io-pfb-export\\/")),
                 true, // requirePrivateWorkspace
                 true, // requireProtectedDataPolicy
-                List.of("mock-auth-group")));
+                List.of("mock-auth-group"),
+                true));
     when(dataImportProperties.getSources()).thenReturn(mockSources);
   }