feat(uploadFiles): disable feature flag when bucket name was not provided

build.gradle

@@ -128,6 +128,8 @@ project.ext.externalDependency = [
     'avroCompiler': 'org.apache.avro:avro-compiler:1.11.4',
     'awsGlueSchemaRegistrySerde': 'software.amazon.glue:schema-registry-serde:1.1.23',
     'awsMskIamAuth': 'software.amazon.msk:aws-msk-iam-auth:2.3.2',
+    'awsJavaSdkCore': 'software.amazon.awssdk:sdk-core:2.17.261',
+    'awsJavaSdkSts': 'software.amazon.awssdk:sts:2.17.261',
     'awsS3': "software.amazon.awssdk:s3:$awsSdk2Version",
     'awsSecretsManagerJdbc': 'com.amazonaws.secretsmanager:aws-secretsmanager-jdbc:1.0.15',
     'awsPostgresIamAuth': 'software.amazon.jdbc:aws-advanced-jdbc-wrapper:2.5.4',

datahub-graphql-core/build.gradle

@@ -22,6 +22,11 @@ dependencies {
     implementation externalDependency.guava
     implementation externalDependency.opentelemetryAnnotations
 
+    implementation platform('software.amazon.awssdk:bom:2.23.6')
+    implementation 'software.amazon.awssdk:regions'
+    implementation 'software.amazon.awssdk:sts'
+    implementation 'software.amazon.awssdk:s3'
+
     implementation externalDependency.slf4jApi
     implementation externalDependency.springContext
     compileOnly externalDependency.lombok

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java

@@ -34,6 +34,7 @@ private Constants() {}
   public static final String TIMESERIES_SCHEMA_FILE = "timeseries.graphql";
   public static final String LOGICAL_SCHEMA_FILE = "logical.graphql";
   public static final String SETTINGS_SCHEMA_FILE = "settings.graphql";
+  public static final String FILES_SCHEMA_FILE = "files.graphql";
 
   public static final String QUERY_SCHEMA_FILE = "query.graphql";
   public static final String TEMPLATE_SCHEMA_FILE = "template.graphql";

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -82,6 +82,7 @@
 import com.linkedin.datahub.graphql.resolvers.entity.EntityPrivilegesResolver;
 import com.linkedin.datahub.graphql.resolvers.entity.versioning.LinkAssetVersionResolver;
 import com.linkedin.datahub.graphql.resolvers.entity.versioning.UnlinkAssetVersionResolver;
+import com.linkedin.datahub.graphql.resolvers.files.GetPresignedUploadUrlResolver;
 import com.linkedin.datahub.graphql.resolvers.form.BatchAssignFormResolver;
 import com.linkedin.datahub.graphql.resolvers.form.BatchRemoveFormResolver;
 import com.linkedin.datahub.graphql.resolvers.form.CreateDynamicFormAssignmentResolver;
@@ -315,6 +316,7 @@
 import com.linkedin.datahub.graphql.types.test.TestType;
 import com.linkedin.datahub.graphql.types.versioning.VersionSetType;
 import com.linkedin.datahub.graphql.types.view.DataHubViewType;
+import com.linkedin.datahub.graphql.util.S3Util;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.entity.client.SystemEntityClient;
 import com.linkedin.metadata.client.UsageStatsJavaClient;
@@ -489,6 +491,8 @@ public class GmsGraphQLEngine {
   private final GraphQLConfiguration graphQLConfiguration;
   private final MetricUtils metricUtils;
 
+  private final S3Util s3Util;
+
   private final BusinessAttributeType businessAttributeType;
 
   /** A list of GraphQL Plugins that extend the core engine */
@@ -621,6 +625,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
     this.dataHubPageModuleType = new PageModuleType(entityClient);
     this.graphQLConfiguration = args.graphQLConfiguration;
     this.metricUtils = args.metricUtils;
+    this.s3Util = args.s3Util;
 
     this.businessAttributeType = new BusinessAttributeType(entityClient);
     // Init Lists
@@ -845,7 +850,8 @@ public GraphQLEngine.Builder builder() {
         .addSchema(fileBasedSchema(QUERY_SCHEMA_FILE))
         .addSchema(fileBasedSchema(TEMPLATE_SCHEMA_FILE))
         .addSchema(fileBasedSchema(MODULE_SCHEMA_FILE))
-        .addSchema(fileBasedSchema(SETTINGS_SCHEMA_FILE));
+        .addSchema(fileBasedSchema(SETTINGS_SCHEMA_FILE))
+        .addSchema(fileBasedSchema(FILES_SCHEMA_FILE));
 
     for (GmsGraphQLPlugin plugin : this.graphQLPlugins) {
       List<String> pluginSchemaFiles = plugin.getSchemaFiles();
@@ -1108,7 +1114,11 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
                     new DocPropagationSettingsResolver(this.settingsService))
                 .dataFetcher(
                     "globalHomePageSettings",
-                    new GlobalHomePageSettingsResolver(this.settingsService)));
+                    new GlobalHomePageSettingsResolver(this.settingsService))
+                .dataFetcher(
+                    "getPresignedUploadUrl",
+                    new GetPresignedUploadUrlResolver(
+                        this.s3Util, this.datahubConfiguration.getS3().getBucketName())));
   }
 
   private DataFetcher getEntitiesResolver() {
@@ -3662,4 +3672,18 @@ private void configureAssetSettingsResolver(final RuntimeWiring.Builder builder)
                       return null;
                     })));
   }
+
+  private void configureFilesResolver(final RuntimeWiring.Builder builder) {
+    builder.type(
+        "GetPresignedUploadUrlResponse",
+        typeWiring ->
+            typeWiring
+                .dataFetcher(
+                    "latestVersion",
+                    new EntityTypeResolver(
+                        entityTypes, (env) -> ((VersionSet) env.getSource()).getLatestVersion()))
+                .dataFetcher(
+                    "versionsSearch",
+                    new VersionsSearchResolver(this.entityClient, this.viewService)));
+  }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java

@@ -10,6 +10,7 @@
 import com.datahub.authorization.role.RoleService;
 import com.linkedin.datahub.graphql.analytics.service.AnalyticsService;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
+import com.linkedin.datahub.graphql.util.S3Util;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.entity.client.SystemEntityClient;
 import com.linkedin.metadata.client.UsageStatsJavaClient;
@@ -97,6 +98,6 @@ public class GmsGraphQLEngineArgs {
   PageModuleService pageModuleService;
   boolean systemTelemetryEnabled;
   MetricUtils metricUtils;
-
+  S3Util s3Util;
   // any fork specific args should go below this line
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java

@@ -266,7 +266,7 @@ public CompletableFuture<AppConfig> get(final DataFetchingEnvironment environmen
             .setLogicalModelsEnabled(_featureFlags.isLogicalModelsEnabled())
             .setShowHomepageUserRole(_featureFlags.isShowHomepageUserRole())
             .setAssetSummaryPageV1(_featureFlags.isAssetSummaryPageV1())
-            .setDocumentationFileUploadV1(_featureFlags.isDocumentationFileUploadV1())
+            .setDocumentationFileUploadV1(isDocumentationFileUploadV1Enabled())
             .build();
 
     appConfig.setFeatureFlags(featureFlagsConfig);
@@ -368,4 +368,18 @@ private EntityType mapResourceTypeToEntityType(final String resourceType) {
       return null;
     }
   }
+
+  private boolean isDocumentationFileUploadV1Enabled() {
+    boolean isEnabledInConfig = _featureFlags.isDocumentationFileUploadV1();
+    if (!isEnabledInConfig) return false;
+
+    // Check if S3 bucket name is configured
+    String bucketName = _datahubConfiguration.getS3().getBucketName();
+    if (bucketName == null || bucketName.isEmpty()) {
+      log.debug("DocumentationFileUploadV1 disabled: DATAHUB_BUCKET_NAME not configured");
+      return false;
+    }
+
+    return true;
+  }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/files/GetPresignedUploadUrlResolver.java

@@ -0,0 +1,127 @@
+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.GetPresignedUploadUrl;
+import com.linkedin.datahub.graphql.generated.GetPresignedUploadUrlInput;
+import com.linkedin.datahub.graphql.generated.UploadDownloadScenario;
+import com.linkedin.datahub.graphql.resolvers.mutate.DescriptionUtils;
+import com.linkedin.datahub.graphql.util.S3Util;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.stereotype.Component;
+
+@Slf4j
+@RequiredArgsConstructor
+@Component
+public class GetPresignedUploadUrlResolver
+    implements DataFetcher<CompletableFuture<GetPresignedUploadUrl>> {
+
+  private static final int EXPIRATION_SECONDS = 60 * 60; // 60 minutes
+  private static final Set<String> ALLOWED_FILE_EXTENSIONS =
+      new HashSet<>(
+          Arrays.asList(
+              "pdf", "jpeg", "jpg", "png", "pptx", "docx", "xls", "xml", "ppt", "gif", "xlsx",
+              "bmp", "doc", "rtf", "gz", "zip", "mp4", "mp3", "wmv", "tiff", "txt", "md", "csv"));
+
+  private final S3Util s3Util;
+  private final String bucketName;
+
+  @Override
+  public CompletableFuture<GetPresignedUploadUrl> get(DataFetchingEnvironment environment)
+      throws Exception {
+    if (s3Util == null) {
+      throw new IllegalArgumentException("S3Util isn't provided");
+    }
+
+    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, EXPIRATION_SECONDS, contentType);
+
+          GetPresignedUploadUrl result = new GetPresignedUploadUrl();
+          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();
+
+    validateFileName(input.getFileName());
+
+    if (scenario == UploadDownloadScenario.ASSET_DOCUMENTATION) {
+      validateInputForAssetDocumentationScenario(context, input);
+    }
+  }
+
+  private void validateFileName(final String fileName) {
+    String fileExtension = "";
+    int i = fileName.lastIndexOf('.');
+    if (i > 0) {
+      fileExtension = fileName.substring(i + 1);
+    }
+
+    if (!ALLOWED_FILE_EXTENSIONS.contains(fileExtension.toLowerCase())) {
+      throw new IllegalArgumentException(
+          String.format("Unsupported file extension: %s", fileExtension));
+    }
+  }
+
+  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/product-assets/%s", bucketName, fileId);
+    } else {
+      throw new IllegalArgumentException("Unsupported upload scenario: " + scenario);
+    }
+  }
+}