aicam commented on code in PR #7764: URL: https://github.com/apache/texera/pull/7764#discussion_r3845603249
########## file-service/src/main/scala/org/apache/texera/service/resource/ResourceUploadService.scala: ########## @@ -0,0 +1,1066 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.service.resource + +import jakarta.ws.rs._ +import jakarta.ws.rs.core.{HttpHeaders, Response} +import org.apache.texera.amber.core.storage.ResourceType +import org.apache.texera.amber.core.storage.util.LakeFSStorageClient +import org.apache.texera.common.config.StorageConfig +import org.apache.texera.dao.{SiteSettings, SqlServer} +import org.apache.texera.dao.SqlServer.withTransaction +import org.apache.texera.dao.jooq.generated.tables.Dataset.DATASET +import org.apache.texera.dao.jooq.generated.tables.DatasetUploadSession.DATASET_UPLOAD_SESSION +import org.apache.texera.dao.jooq.generated.tables.DatasetUploadSessionPart.DATASET_UPLOAD_SESSION_PART +import org.apache.texera.dao.jooq.generated.tables.records.{ + DatasetRecord, + DatasetUploadSessionPartRecord, + DatasetUploadSessionRecord, + DatasetUserAccessRecord +} +import org.apache.texera.service.util.LakeFSExceptionHandler.withLakeFSErrorHandling +import org.apache.texera.service.util.S3StorageClient +import org.apache.texera.service.util.S3StorageClient.{ + MAXIMUM_NUM_OF_MULTIPART_S3_PARTS, + MINIMUM_NUM_OF_MULTIPART_S3_PART, + PHYSICAL_ADDRESS_EXPIRATION_TIME_HRS +} +import org.jooq.exception.DataAccessException +import org.jooq.impl.DSL +import org.jooq.impl.DSL.{inline => inl} +import org.jooq.{DSLContext, Record, Record2, Result, Table, TableField} +import software.amazon.awssdk.services.s3.model.UploadPartResponse + +import java.io.InputStream +import java.net.URLDecoder +import java.nio.charset.StandardCharsets +import java.sql.SQLException +import java.time.OffsetDateTime +import java.util.Optional +import scala.collection.mutable.ListBuffer +import scala.jdk.CollectionConverters._ +import scala.util.Try + +/** + * Describes where a resource's files live and how its in-progress uploads are tracked. + * + * [[ResourceTables]] names the columns that carry identity, ownership and grants; this adds + * the storage side — the LakeFS repository column plus the `*_upload_session` and + * `*_upload_session_part` tables that back resumable multipart uploads. Naming the columns + * keeps one implementation of the upload engine serving every resource type, so adding the + * next one costs a descriptor rather than another copy of the locking, part-size and + * resume rules. + * + * @tparam R record type of the resource table + * @tparam A record type of the companion user-access table + * @tparam S record type of the upload-session table + * @tparam P record type of the upload-session-part table + */ +case class ResourceStorage[R <: Record, A <: Record, S <: Record, P <: Record]( + resource: ResourceTables[R, A], + resourceType: ResourceType.Value, + repositoryNameField: TableField[R, String], + sessionResourceId: TableField[S, Integer], + sessionUid: TableField[S, Integer], + sessionFilePath: TableField[S, String], + sessionUploadId: TableField[S, String], + sessionPhysicalAddress: TableField[S, String], + sessionNumParts: TableField[S, Integer], + sessionFileSize: TableField[S, java.lang.Long], + sessionPartSize: TableField[S, java.lang.Long], + sessionCreatedAt: TableField[S, OffsetDateTime], + partUploadId: TableField[P, String], + partNumber: TableField[P, Integer], + partEtag: TableField[P, String] +) { + def sessionTable: Table[S] = sessionResourceId.getTable + def partTable: Table[P] = partUploadId.getTable +} + +object ResourceStorage { + + val Dataset: ResourceStorage[ + DatasetRecord, + DatasetUserAccessRecord, + DatasetUploadSessionRecord, + DatasetUploadSessionPartRecord + ] = + ResourceStorage( + resource = ResourceTables.Dataset, + resourceType = ResourceType.Dataset, + repositoryNameField = DATASET.REPOSITORY_NAME, + sessionResourceId = DATASET_UPLOAD_SESSION.DID, + sessionUid = DATASET_UPLOAD_SESSION.UID, + sessionFilePath = DATASET_UPLOAD_SESSION.FILE_PATH, + sessionUploadId = DATASET_UPLOAD_SESSION.UPLOAD_ID, + sessionPhysicalAddress = DATASET_UPLOAD_SESSION.PHYSICAL_ADDRESS, + sessionNumParts = DATASET_UPLOAD_SESSION.NUM_PARTS_REQUESTED, + sessionFileSize = DATASET_UPLOAD_SESSION.FILE_SIZE_BYTES, + sessionPartSize = DATASET_UPLOAD_SESSION.PART_SIZE_BYTES, + sessionCreatedAt = DATASET_UPLOAD_SESSION.CREATED_AT, + partUploadId = DATASET_UPLOAD_SESSION_PART.UPLOAD_ID, + partNumber = DATASET_UPLOAD_SESSION_PART.PART_NUMBER, + partEtag = DATASET_UPLOAD_SESSION_PART.ETAG + ) +} + +/** + * The upload and version-file machinery shared by every versioned file resource. + * + * Every method here was previously duplicated per resource; the two copies differed only in + * which jOOQ columns they named. Keeping one implementation means the locking protocol + * (`FOR UPDATE NOWAIT`, SQLState 55P03 to 409), the part-size arithmetic and its overflow + * guards, ETag idempotency and the resume/restart rules are defined once. + */ +object ResourceUploadService { + + private def context: DSLContext = + SqlServer + .getInstance() + .createDSLContext() + + private def singleFileUploadMaxBytes(defaultMiB: Long = 20L): Long = + SiteSettings.getLong("single_file_upload_max_size_mib", defaultMiB) * 1024L * 1024L + + private def noAccessMessage[R <: Record, A <: Record, S <: Record, P <: Record]( + s: ResourceStorage[R, A, S, P] + ): String = s"User has no access to this ${s.resource.label}" + + /** Reads the LakeFS repository backing a resource, or 404s if the resource is gone. */ + private def repositoryNameOf[R <: Record, A <: Record, S <: Record, P <: Record]( + ctx: DSLContext, + s: ResourceStorage[R, A, S, P], + resourceId: Integer + ): String = + Option( + ctx + .select(s.repositoryNameField) + .from(s.resource.table) + .where(s.resource.idField.eq(resourceId)) + .fetchOne(s.repositoryNameField) + ).getOrElse( + throw new NotFoundException(s"${s.resource.label.capitalize} $resourceId not found") + ) + + /** Removes one staged (uncommitted) file from a resource's repository. */ + def deleteStagedFile[R <: Record, A <: Record, S <: Record, P <: Record]( + s: ResourceStorage[R, A, S, P], + resourceId: Integer, + encodedFilePath: String, + uid: Integer + ): Response = { + withTransaction(context) { ctx => + if (!ResourceAccess.userHasWriteAccess(ctx, s.resource, resourceId, uid)) { + throw new ForbiddenException(noAccessMessage(s)) + } + val repositoryName = repositoryNameOf(ctx, s, resourceId) + + val filePath = URLDecoder.decode(encodedFilePath, StandardCharsets.UTF_8.name()) + withLakeFSErrorHandling( + s"deleting file '$filePath' from the ${s.resource.label} repository" + ) { + LakeFSStorageClient.deleteObject(repositoryName, filePath) + } + + Response.ok().build() + } + } + + def uploadOneFile[R <: Record, A <: Record, S <: Record, P <: Record]( + s: ResourceStorage[R, A, S, P], + resourceId: Integer, + encodedFilePath: String, + fileStream: InputStream, + headers: HttpHeaders, + uid: Integer + ): Response = { + // These variables are defined at the top so catch block can access them + var repoName: String = null + var filePath: String = null + var uploadId: String = null + var physicalAddress: String = null + + try { + withTransaction(context) { ctx => + if (!ResourceAccess.userHasWriteAccess(ctx, s.resource, resourceId, uid)) + throw new ForbiddenException(noAccessMessage(s)) + + repoName = repositoryNameOf(ctx, s, resourceId) + filePath = URLDecoder.decode(encodedFilePath, StandardCharsets.UTF_8.name) + + // ---------- decide part-size & number-of-parts ---------- + val declaredLen = Option(headers.getHeaderString(HttpHeaders.CONTENT_LENGTH)).map(_.toLong) + var partSize = StorageConfig.s3MultipartUploadPartSize + + declaredLen.foreach { ln => + val needed = ((ln + partSize - 1) / partSize).toInt + if (needed > MAXIMUM_NUM_OF_MULTIPART_S3_PARTS) + partSize = math.max( + MINIMUM_NUM_OF_MULTIPART_S3_PART, + ln / (MAXIMUM_NUM_OF_MULTIPART_S3_PARTS - 1) + ) + } + + val expectedParts = declaredLen + .map(ln => + ((ln + partSize - 1) / partSize).toInt + 1 + ) // “+1” for the last (possibly small) part + .getOrElse(MAXIMUM_NUM_OF_MULTIPART_S3_PARTS) + + // ---------- ask LakeFS for presigned URLs ---------- + val presign = LakeFSStorageClient + .initiatePresignedMultipartUploads(repoName, filePath, expectedParts) + uploadId = presign.getUploadId + val presignedUrls = presign.getPresignedUrls.asScala.iterator + physicalAddress = presign.getPhysicalAddress + + // ---------- stream & upload parts ---------- + /* + 1. Reads the input stream in chunks of 'partSize' bytes by stacking them in a buffer + 2. Uploads each chunk (part) using a presigned URL + 3. Tracks each part number and ETag returned from S3 + 4. After all parts are uploaded, completes the multipart upload + */ + val buf = new Array[Byte](partSize.toInt) + var buffered = 0 + var partNumber = 1 + val completedParts = ListBuffer[(Int, String)]() + + @inline def flush(): Unit = { + if (buffered == 0) return + if (!presignedUrls.hasNext) + throw new WebApplicationException("Ran out of presigned part URLs – ask for more parts") + + val etag = LakeFSStorageClient.put(buf, buffered, presignedUrls.next(), partNumber) + completedParts += ((partNumber, etag)) + partNumber += 1 + buffered = 0 + } + + var read = fileStream.read(buf, buffered, buf.length - buffered) + while (read != -1) { + buffered += read + if (buffered == buf.length) flush() // buffer full Review Comment: change to >= so it still flush if overflow happen -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
