Implement deletion of S3 datasets

Author: frcroth

app/models/dataset/DatasetService.scala

@@ -671,16 +671,16 @@ class DatasetService @Inject()(organizationDAO: OrganizationDAO,
 
   def deleteVirtualOrDiskDataset(dataset: Dataset)(implicit ctx: DBAccessContext): Fox[Unit] =
     for {
-      _ <- if (dataset.isVirtual) {
-        // At this point, we should also free space in S3 once implemented.
-        // Right now, we can just mark the dataset as deleted in the database.
-        datasetDAO.deleteDataset(dataset._id, onlyMarkAsDeleted = true)
-      } else {
-        for {
-          datastoreClient <- clientFor(dataset)
-          _ <- datastoreClient.deleteOnDisk(dataset._id)
-        } yield ()
-      } ?~> "dataset.delete.failed"
+      //_ <- if (dataset.isVirtual) {
+      // At this point, we should also free space in S3 once implemented.
+      // Right now, we can just mark the dataset as deleted in the database.
+      //  datasetDAO.deleteDataset(dataset._id, onlyMarkAsDeleted = true)
+      //} else {
+      //for {
+      datastoreClient <- clientFor(dataset)
+      _ <- datastoreClient.deleteOnDisk(dataset._id) ?~> "dataset.delete.failed"
+      //  } yield ()
+      //} ?~> "dataset.delete.failed"
     } yield ()
 
   def publicWrites(dataset: Dataset,

webknossos-datastore/app/com/scalableminds/webknossos/datastore/controllers/DataSourceController.scala

@@ -448,20 +448,22 @@ class DataSourceController @Inject()(
         for {
           dataSource <- datasetCache.getById(datasetId) ~> NOT_FOUND
           dataSourceId = dataSource.id
-          _ <- if (dataSourceService.existsOnDisk(dataSourceId.organizationId, dataSourceId.directoryName)) {
+          existsOnDisk = dataSourceService.existsOnDisk(dataSourceId.organizationId, dataSourceId.directoryName)
+          _ <- if (existsOnDisk) {
             for {
               _ <- dataSourceService.deleteOnDisk(
                 dataSourceId.organizationId,
                 dataSourceId.directoryName,
                 Some(datasetId),
                 reason = Some("the user wants to delete the dataset")) ?~> "dataset.delete.failed"
-              _ <- dsRemoteWebknossosClient.deleteDataSource(dataSourceId)
             } yield ()
           } else
             for {
+              _ <- Fox.runIf(dataSourceService.datasetInControlledS3(dataSource))(dataSourceService.deleteFromControlledS3(dataSource))
               _ <- dsRemoteWebknossosClient.deleteDataSource(dataSourceId)
               _ = logger.warn(s"Tried to delete dataset ${dataSource.id} that is not on disk.")
             } yield ()
+          _ <- dsRemoteWebknossosClient.deleteDataSource(dataSourceId)
         } yield Ok
       }
     }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datavault/S3DataVault.scala

@@ -159,6 +159,7 @@ object S3DataVault {
     new S3DataVault(credential, remoteSourceDescriptor.uri, ws, ec)
   }
 
+  // TODO: Move non private methods to trait?
   def hostBucketFromUri(uri: URI): Option[String] = {
     val host = uri.getHost
     if (isShortStyle(uri)) { // assume host is omitted from uri, shortcut form s3://bucket/key
@@ -185,7 +186,7 @@ object S3DataVault {
   private def isShortStyle(uri: URI): Boolean =
     !uri.getHost.contains(".")
 
-  private def objectKeyFromUri(uri: URI): Box[String] =
+  def objectKeyFromUri(uri: URI): Box[String] =
     if (isVirtualHostedStyle(uri)) {
       Full(uri.getPath)
     } else if (isPathStyle(uri)) {

webknossos-datastore/app/com/scalableminds/webknossos/datastore/services/DataSourceService.scala

@@ -13,19 +13,38 @@ import com.scalableminds.webknossos.datastore.dataformats.{MagLocator, MappingPr
 import com.scalableminds.webknossos.datastore.helpers.{DatasetDeleter, IntervalScheduler}
 import com.scalableminds.webknossos.datastore.models.datasource._
 import com.scalableminds.webknossos.datastore.models.datasource.inbox.{InboxDataSource, UnusableDataSource}
-import com.scalableminds.webknossos.datastore.storage.{DataVaultService, RemoteSourceDescriptorService}
+import com.scalableminds.webknossos.datastore.storage.{
+  CredentialConfigReader,
+  DataVaultService,
+  RemoteSourceDescriptorService,
+  S3AccessKeyCredential
+}
 import com.typesafe.scalalogging.LazyLogging
 import com.scalableminds.util.tools.Box.tryo
 import com.scalableminds.util.tools._
+import com.scalableminds.webknossos.datastore.datavault.S3DataVault
 import play.api.inject.ApplicationLifecycle
 import play.api.libs.json.Json
+import software.amazon.awssdk.auth.credentials.{AwsBasicCredentials, StaticCredentialsProvider}
+import software.amazon.awssdk.core.checksums.RequestChecksumCalculation
+import software.amazon.awssdk.regions.Region
+import software.amazon.awssdk.services.s3.S3AsyncClient
+import software.amazon.awssdk.services.s3.model.{
+  Delete,
+  DeleteObjectsRequest,
+  DeleteObjectsResponse,
+  ListObjectsV2Request,
+  ObjectIdentifier
+}
 
 import java.io.{File, FileWriter}
 import java.net.URI
 import java.nio.file.{Files, Path}
 import scala.concurrent.ExecutionContext
 import scala.concurrent.duration._
 import scala.io.Source
+import scala.jdk.CollectionConverters._
+import scala.jdk.FutureConverters._
 
 class DataSourceService @Inject()(
     config: DataStoreConfig,
@@ -446,4 +465,115 @@ class DataSourceService @Inject()(
           remoteSourceDescriptorService.removeVaultFromCache(attachment)))
       } yield dataLayer.mags.length
     } yield removedEntriesList.sum
+
+  private lazy val globalCredentials = {
+    val res = config.Datastore.DataVaults.credentials.flatMap { credentialConfig =>
+      new CredentialConfigReader(credentialConfig).getCredential
+    }
+    logger.info(s"Parsed ${res.length} global data vault credentials from datastore config.")
+    res
+  }
+
+  def datasetInControlledS3(dataSource: DataSource) = {
+    def commonPrefix(strings: Seq[String]): String = {
+      if (strings.isEmpty) return ""
+
+      strings.reduce { (a, b) =>
+        a.zip(b).takeWhile { case (c1, c2) => c1 == c2 }.map(_._1).mkString
+      }
+    }
+
+    val allPaths = dataSource.allExplicitPaths
+    val sharedPath = commonPrefix(allPaths)
+    val matchingCredentials = globalCredentials.filter(c => sharedPath.startsWith(c.name))
+    matchingCredentials.nonEmpty && sharedPath.startsWith("s3")
+  }
+
+  private lazy val s3UploadCredentialsOpt: Option[(String, String)] =
+    config.Datastore.DataVaults.credentials.flatMap { credentialConfig =>
+      new CredentialConfigReader(credentialConfig).getCredential
+    }.collectFirst {
+      case S3AccessKeyCredential(credentialName, accessKeyId, secretAccessKey, _, _)
+          if config.Datastore.S3Upload.credentialName == credentialName =>
+        (accessKeyId, secretAccessKey)
+    }
+  private lazy val s3Client: S3AsyncClient = S3AsyncClient
+    .builder()
+    .credentialsProvider(
+      StaticCredentialsProvider.create(
+        AwsBasicCredentials.builder
+          .accessKeyId(s3UploadCredentialsOpt.getOrElse(("", ""))._1)
+          .secretAccessKey(s3UploadCredentialsOpt.getOrElse(("", ""))._2)
+          .build()
+      ))
+    .crossRegionAccessEnabled(true)
+    .forcePathStyle(true)
+    .endpointOverride(new URI(config.Datastore.S3Upload.endpoint))
+    .region(Region.US_EAST_1)
+    // Disabling checksum calculation prevents files being stored with Content Encoding "aws-chunked".
+    .requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED)
+    .build()
+
+  def deleteFromControlledS3(dataSource: DataSource): Fox[Unit] = {
+    // TODO: Do we handle other datasets using the same layers?
+
+    def deleteBatch(bucket: String, keys: Seq[String]): Fox[DeleteObjectsResponse] =
+      if (keys.isEmpty) Fox.empty
+      else {
+        Fox.fromFuture(
+          s3Client
+            .deleteObjects(
+              DeleteObjectsRequest
+                .builder()
+                .bucket(bucket)
+                .delete(
+                  Delete
+                    .builder()
+                    .objects(
+                      keys.map(k => ObjectIdentifier.builder().key(k).build()).asJava
+                    )
+                    .build()
+                )
+                .build()
+            )
+            .asScala)
+      }
+
+    def listKeysAtPrefix(bucket: String, prefix: String): Fox[Seq[String]] = {
+      def listRec(continuationToken: Option[String], acc: Seq[String]): Fox[Seq[String]] = {
+        val builder = ListObjectsV2Request.builder().bucket(bucket).prefix(prefix).maxKeys(1000)
+        val request = continuationToken match {
+          case Some(token) => builder.continuationToken(token).build()
+          case None        => builder.build()
+        }
+        for {
+          response <- Fox.fromFuture(s3Client.listObjectsV2(request).asScala)
+          keys = response.contents().asScala.map(_.key())
+          allKeys = acc ++ keys
+          result <- if (response.isTruncated) {
+            listRec(Option(response.nextContinuationToken()), allKeys)
+          } else {
+            Fox.successful(allKeys)
+          }
+        } yield result
+      }
+      listRec(None, Seq())
+    }
+
+    for {
+      _ <- Fox.successful(())
+      paths = dataSource.allExplicitPaths
+      // Assume everything is in the same bucket
+      firstPath <- paths.headOption.toFox ?~> "No explicit paths found for dataset in controlled S3"
+      bucket <- S3DataVault
+        .hostBucketFromUri(new URI(firstPath))
+        .toFox ?~> s"Could not determine S3 bucket from path $firstPath"
+      prefixes <- Fox.combined(paths.map(path => S3DataVault.objectKeyFromUri(new URI(path)).toFox))
+      keys: Seq[String] <- Fox.serialCombined(prefixes)(listKeysAtPrefix(bucket, _)).map(_.flatten)
+      uniqueKeys = keys.distinct
+      _ = logger.info(
+        s"Deleting ${uniqueKeys.length} objects from controlled S3 bucket $bucket for dataset ${dataSource.id}")
+      _ <- Fox.serialCombined(uniqueKeys.grouped(1000).toSeq)(deleteBatch(bucket, _)).map(_ => ())
+    } yield ()
+  }
 }