-
Notifications
You must be signed in to change notification settings - Fork 3.2k
feat(uploadFiles): add endpoint to get presigned upload url #14943
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
v-tarasevich-blitz-brain
wants to merge
18
commits into
master
Choose a base branch
from
vt--ch-820-upload-files--add-upload-method
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+1,486
−3
Open
Changes from all commits
Commits
Show all changes
18 commits
Select commit
Hold shift + click to select a range
629229b
feat(uploadFiles): add endpoint to get presigned upload url
v-tarasevich-blitz-brain 1810481
fix tests
v-tarasevich-blitz-brain 26fb897
fix tests
v-tarasevich-blitz-brain c698217
add fileName, validate files extension
v-tarasevich-blitz-brain a97f364
add some keys to docker and add fileId to response
v-tarasevich-blitz-brain 0d3df07
move new graphql things into separated file
v-tarasevich-blitz-brain 7e3e042
add s3utilfactory
v-tarasevich-blitz-brain e4993ee
fix tests
v-tarasevich-blitz-brain 618c378
pass bucket name from config
v-tarasevich-blitz-brain cd5a976
fix and bump tests
v-tarasevich-blitz-brain d3747dd
rename GetPresignedUploadUrl to GetPresignedUploadUrlResponse
v-tarasevich-blitz-brain ec49fd4
remove unnecessary code
v-tarasevich-blitz-brain f25bf45
clean up dependecies
v-tarasevich-blitz-brain 229e176
use StsAssumeRoleCredentialsProvider with automatic refresh
v-tarasevich-blitz-brain 3628917
move hardcoded values from resolver to config
v-tarasevich-blitz-brain f86a791
remove quickstartDebugAws
v-tarasevich-blitz-brain 6f48d92
remove validation of file extensions
v-tarasevich-blitz-brain 6af0296
add new options to NON_SENSITIVE_PROPERTIES to fix tests
v-tarasevich-blitz-brain File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Some comments aren't visible on the classic Files Changed page.
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
114 changes: 114 additions & 0 deletions
114
...main/java/com/linkedin/datahub/graphql/resolvers/files/GetPresignedUploadUrlResolver.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,114 @@ | ||
| package com.linkedin.datahub.graphql.resolvers.files; | ||
|
|
||
| import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; | ||
|
|
||
| import com.linkedin.common.urn.UrnUtils; | ||
| import com.linkedin.datahub.graphql.QueryContext; | ||
| import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils; | ||
| import com.linkedin.datahub.graphql.exception.AuthorizationException; | ||
| import com.linkedin.datahub.graphql.generated.GetPresignedUploadUrlInput; | ||
| import com.linkedin.datahub.graphql.generated.GetPresignedUploadUrlResponse; | ||
| import com.linkedin.datahub.graphql.generated.UploadDownloadScenario; | ||
| import com.linkedin.datahub.graphql.resolvers.mutate.DescriptionUtils; | ||
| import com.linkedin.datahub.graphql.util.S3Util; | ||
| import com.linkedin.metadata.config.S3Configuration; | ||
| import graphql.schema.DataFetcher; | ||
| import graphql.schema.DataFetchingEnvironment; | ||
| import java.util.UUID; | ||
| import java.util.concurrent.CompletableFuture; | ||
| import lombok.extern.slf4j.Slf4j; | ||
| import org.springframework.stereotype.Component; | ||
|
|
||
| @Slf4j | ||
| @Component | ||
| public class GetPresignedUploadUrlResolver | ||
| implements DataFetcher<CompletableFuture<GetPresignedUploadUrlResponse>> { | ||
|
|
||
| private final S3Util s3Util; | ||
| private final S3Configuration s3Configuration; | ||
|
|
||
| public GetPresignedUploadUrlResolver(S3Util s3Util, S3Configuration s3Configuration) { | ||
| this.s3Util = s3Util; | ||
| this.s3Configuration = s3Configuration; | ||
| } | ||
|
|
||
| @Override | ||
| public CompletableFuture<GetPresignedUploadUrlResponse> get(DataFetchingEnvironment environment) | ||
| throws Exception { | ||
| if (s3Util == null) { | ||
| throw new IllegalArgumentException("S3Util isn't provided"); | ||
| } | ||
|
|
||
| String bucketName = s3Configuration.getBucketName(); | ||
|
|
||
| if (bucketName == null || bucketName.isEmpty()) { | ||
| throw new IllegalArgumentException("Bucket name isn't provided"); | ||
| } | ||
|
|
||
| final GetPresignedUploadUrlInput input = | ||
| bindArgument(environment.getArgument("input"), GetPresignedUploadUrlInput.class); | ||
|
|
||
| final QueryContext context = environment.getContext(); | ||
|
|
||
| validateInput(context, input); | ||
|
|
||
| String newFileId = generateNewFileId(input); | ||
| String s3Key = getS3Key(input, newFileId, bucketName); | ||
| String contentType = input.getContentType(); | ||
|
|
||
| return GraphQLConcurrencyUtils.supplyAsync( | ||
| () -> { | ||
| String presignedUploadUrl = | ||
| s3Util.generatePresignedUploadUrl( | ||
| bucketName, | ||
| s3Key, | ||
| s3Configuration.getPresignedUploadUrlExpirationSeconds(), | ||
| contentType); | ||
|
|
||
| GetPresignedUploadUrlResponse result = new GetPresignedUploadUrlResponse(); | ||
| result.setUrl(presignedUploadUrl); | ||
| result.setFileId(newFileId); | ||
| return result; | ||
| }, | ||
| this.getClass().getSimpleName(), | ||
| "get"); | ||
| } | ||
|
|
||
| private void validateInput(final QueryContext context, final GetPresignedUploadUrlInput input) { | ||
| UploadDownloadScenario scenario = input.getScenario(); | ||
|
|
||
| if (scenario == UploadDownloadScenario.ASSET_DOCUMENTATION) { | ||
| validateInputForAssetDocumentationScenario(context, input); | ||
| } | ||
| } | ||
|
|
||
| private void validateInputForAssetDocumentationScenario( | ||
| final QueryContext context, final GetPresignedUploadUrlInput input) { | ||
| String assetUrn = input.getAssetUrn(); | ||
|
|
||
| if (assetUrn == null) { | ||
| throw new IllegalArgumentException("assetUrn is required for ASSET_DOCUMENTATION scenario"); | ||
| } | ||
|
|
||
| if (!DescriptionUtils.isAuthorizedToUpdateDescription(context, UrnUtils.getUrn(assetUrn))) { | ||
| throw new AuthorizationException("Unauthorized to edit documentation for asset: " + assetUrn); | ||
| } | ||
| } | ||
|
|
||
| private String generateNewFileId(final GetPresignedUploadUrlInput input) { | ||
| return String.format("%s-%s", UUID.randomUUID().toString(), input.getFileName()); | ||
| } | ||
|
|
||
| private String getS3Key( | ||
| final GetPresignedUploadUrlInput input, final String fileId, final String bucketName) { | ||
| UploadDownloadScenario scenario = input.getScenario(); | ||
|
|
||
| if (scenario == UploadDownloadScenario.ASSET_DOCUMENTATION) { | ||
| return String.format( | ||
| "%s/%s/%s", | ||
| s3Configuration.getBucketName(), s3Configuration.getAssetPathPrefix(), fileId); | ||
| } else { | ||
| throw new IllegalArgumentException("Unsupported upload scenario: " + scenario); | ||
| } | ||
| } | ||
| } |
172 changes: 172 additions & 0 deletions
172
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/util/S3Util.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,172 @@ | ||
| package com.linkedin.datahub.graphql.util; | ||
|
|
||
| import com.linkedin.entity.client.EntityClient; | ||
| import java.time.Duration; | ||
| import javax.annotation.Nonnull; | ||
| import javax.annotation.Nullable; | ||
| import lombok.extern.slf4j.Slf4j; | ||
| import software.amazon.awssdk.services.s3.S3Client; | ||
| import software.amazon.awssdk.services.s3.model.*; | ||
| import software.amazon.awssdk.services.s3.presigner.S3Presigner; | ||
| import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest; | ||
| import software.amazon.awssdk.services.s3.presigner.model.PresignedGetObjectRequest; | ||
| import software.amazon.awssdk.services.s3.presigner.model.PresignedPutObjectRequest; | ||
| import software.amazon.awssdk.services.s3.presigner.model.PutObjectPresignRequest; | ||
| import software.amazon.awssdk.services.sts.StsClient; | ||
| import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; | ||
|
|
||
| @Slf4j | ||
| public class S3Util { | ||
|
|
||
| private final S3Client s3Client; | ||
| private final EntityClient entityClient; | ||
|
|
||
| // Optional S3Presigner for testing purposes | ||
| @Nullable private final S3Presigner s3Presigner; | ||
|
|
||
| public S3Util(@Nonnull S3Client s3Client, @Nonnull EntityClient entityClient) { | ||
| this(s3Client, entityClient, null); | ||
| } | ||
|
|
||
| public S3Util( | ||
| @Nonnull S3Client s3Client, | ||
| @Nonnull EntityClient entityClient, | ||
| @Nullable S3Presigner s3Presigner) { | ||
| this.s3Client = s3Client; | ||
| this.entityClient = entityClient; | ||
| this.s3Presigner = s3Presigner; | ||
| } | ||
|
|
||
| public S3Util( | ||
| @Nonnull EntityClient entityClient, @Nonnull StsClient stsClient, @Nonnull String roleArn) { | ||
| this(entityClient, stsClient, roleArn, null); | ||
| } | ||
|
|
||
| public S3Util( | ||
| @Nonnull EntityClient entityClient, | ||
| @Nonnull StsClient stsClient, | ||
| @Nonnull String roleArn, | ||
| @Nullable S3Presigner s3Presigner) { | ||
| this.entityClient = entityClient; | ||
| this.s3Presigner = s3Presigner; | ||
| this.s3Client = createS3Client(stsClient, roleArn); | ||
| } | ||
|
|
||
| /** Creates S3Client with StsAssumeRoleCredentialsProvider for automatic credential refresh. */ | ||
| private static S3Client createS3Client(@Nonnull StsClient stsClient, @Nonnull String roleArn) { | ||
| try { | ||
| log.info("Creating S3Client for role: {}", roleArn); | ||
|
|
||
| StsAssumeRoleCredentialsProvider credentialsProvider = | ||
| StsAssumeRoleCredentialsProvider.builder() | ||
| .stsClient(stsClient) | ||
| .refreshRequest(r -> r.roleArn(roleArn).roleSessionName("s3-session")) | ||
| .asyncCredentialUpdateEnabled(true) // Enable background credential refresh | ||
| .build(); | ||
|
|
||
| var clientBuilder = S3Client.builder().credentialsProvider(credentialsProvider); | ||
|
|
||
| // Configure endpoint URL if provided (for LocalStack or custom S3 endpoints) | ||
| String endpointUrl = System.getenv("AWS_ENDPOINT_URL"); | ||
| if (endpointUrl != null && !endpointUrl.isEmpty()) { | ||
| clientBuilder.endpointOverride(java.net.URI.create(endpointUrl)); | ||
| // Force path-style access for LocalStack compatibility | ||
| clientBuilder.forcePathStyle(true); | ||
| } | ||
|
|
||
| S3Client client = clientBuilder.build(); | ||
| log.info("Successfully created S3Client for role: {}", roleArn); | ||
| return client; | ||
|
|
||
| } catch (Exception e) { | ||
| log.error("Failed to create S3 client: roleArn={}", roleArn, e); | ||
| throw new RuntimeException("Failed to create S3 clien: " + e.getMessage(), e); | ||
| } | ||
| } | ||
|
|
||
| private S3Presigner getPresigner() { | ||
| if (this.s3Presigner != null) { | ||
| return this.s3Presigner; | ||
| } | ||
|
|
||
| return S3Presigner.builder() | ||
| .credentialsProvider(s3Client.serviceClientConfiguration().credentialsProvider()) | ||
| .region(s3Client.serviceClientConfiguration().region()) | ||
| .build(); | ||
| } | ||
|
|
||
| /** | ||
| * Generate a pre-signed URL for downloading an S3 object | ||
| * | ||
| * @param bucket The S3 bucket name | ||
| * @param key The S3 object key | ||
| * @param expirationSeconds The expiration time in seconds | ||
| * @return The pre-signed URL | ||
| */ | ||
| public String generatePresignedDownloadUrl( | ||
v-tarasevich-blitz-brain marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| @Nonnull String bucket, @Nonnull String key, int expirationSeconds) { | ||
| try { | ||
| // Create a pre-signer using the same configuration as the S3 client | ||
| try (S3Presigner presigner = getPresigner()) { | ||
|
|
||
| // Create the GetObjectRequest | ||
| GetObjectRequest getObjectRequest = | ||
| GetObjectRequest.builder().bucket(bucket).key(key).build(); | ||
|
|
||
| // Create the presign request | ||
| GetObjectPresignRequest presignRequest = | ||
| GetObjectPresignRequest.builder() | ||
| .signatureDuration(Duration.ofSeconds(expirationSeconds)) | ||
| .getObjectRequest(getObjectRequest) | ||
| .build(); | ||
|
|
||
| // Generate the presigned URL | ||
| PresignedGetObjectRequest presignedRequest = presigner.presignGetObject(presignRequest); | ||
| return presignedRequest.url().toString(); | ||
| } | ||
| } catch (Exception e) { | ||
| log.error("Failed to generate presigned URL for bucket: {}, key: {}", bucket, key, e); | ||
| throw new RuntimeException("Failed to generate presigned URL: " + e.getMessage(), e); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Generate a pre-signed URL for uploading an S3 object | ||
| * | ||
| * @param bucket The S3 bucket name | ||
| * @param key The S3 object key | ||
| * @param expirationSeconds The expiration time in seconds | ||
| * @param contentType The content type of the object to be uploaded (e.g., "image/jpeg", | ||
| * "application/pdf") | ||
| * @return The pre-signed URL | ||
| */ | ||
| public String generatePresignedUploadUrl( | ||
| @Nonnull String bucket, | ||
| @Nonnull String key, | ||
| int expirationSeconds, | ||
| @Nullable String contentType) { | ||
| try { | ||
| // Create a pre-signer using the same configuration as the S3 client | ||
| try (S3Presigner presigner = getPresigner()) { | ||
|
|
||
| // Create the PutObjectRequest | ||
| PutObjectRequest putObjectRequest = | ||
| PutObjectRequest.builder().bucket(bucket).contentType(contentType).key(key).build(); | ||
|
|
||
| // Create the presign request | ||
| PutObjectPresignRequest presignRequest = | ||
| PutObjectPresignRequest.builder() | ||
| .signatureDuration(Duration.ofSeconds(expirationSeconds)) | ||
| .putObjectRequest(putObjectRequest) | ||
| .build(); | ||
|
|
||
| // Generate the presigned URL | ||
| PresignedPutObjectRequest presignedRequest = presigner.presignPutObject(presignRequest); | ||
| return presignedRequest.url().toString(); | ||
| } | ||
| } catch (Exception e) { | ||
| log.error("Failed to generate presigned upload URL for bucket: {}, key: {}", bucket, key, e); | ||
| throw new RuntimeException("Failed to generate presigned upload URL: " + e.getMessage(), e); | ||
| } | ||
| } | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
let's double check @anshbansal has reviewed this file since he was the original implementer