Add support for x-apigw-api-id request header in Glue

docs/src/main/sphinx/connector/hive.rst

@@ -654,6 +654,9 @@ Property Name                                        Description
                                                      required when running in a GovCloud region.
                                                      Example: ``us-gov-east-1``
 
+``hive.metastore.glue.proxy-api-id``                 The ID of the Glue Proxy API, when accessing Glue via an VPC
+                                                     endpoint in API Gateway.
+
 ``hive.metastore.glue.sts.endpoint``                 STS endpoint URL to use when authenticating to Glue (optional).
                                                      Example: ``https://sts.us-gov-east-1.amazonaws.com``
 

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java

@@ -18,18 +18,22 @@
 import io.airlift.configuration.ConfigSecuritySensitive;
 import io.airlift.configuration.DefunctConfig;
 
+import javax.annotation.PostConstruct;
 import javax.validation.constraints.Max;
 import javax.validation.constraints.Min;
 
 import java.util.Optional;
 
+import static com.google.common.base.Preconditions.checkState;
+
 @DefunctConfig("hive.metastore.glue.use-instance-credentials")
 public class GlueHiveMetastoreConfig
 {
     private Optional<String> glueRegion = Optional.empty();
     private Optional<String> glueEndpointUrl = Optional.empty();
     private Optional<String> glueStsRegion = Optional.empty();
     private Optional<String> glueStsEndpointUrl = Optional.empty();
+    private Optional<String> glueProxyApiId = Optional.empty();
     private boolean pinGlueClientToCurrentRegion;
     private int maxGlueErrorRetries = 10;
     private int maxGlueConnections = 30;
@@ -98,6 +102,19 @@ public GlueHiveMetastoreConfig setGlueStsEndpointUrl(String glueStsEndpointUrl)
         return this;
     }
 
+    public Optional<String> getGlueProxyApiId()
+    {
+        return glueProxyApiId;
+    }
+
+    @Config("hive.metastore.glue.proxy-api-id")
+    @ConfigDescription("ID of Glue Proxy API")
+    public GlueHiveMetastoreConfig setGlueProxyApiId(String glueProxyApiId)
+    {
+        this.glueProxyApiId = Optional.ofNullable(glueProxyApiId);
+        return this;
+    }
+
     public boolean getPinGlueClientToCurrentRegion()
     {
         return pinGlueClientToCurrentRegion;
@@ -300,4 +317,13 @@ public GlueHiveMetastoreConfig setWriteStatisticsThreads(int writeStatisticsThre
         this.writeStatisticsThreads = writeStatisticsThreads;
         return this;
     }
+
+    @PostConstruct
+    public void validate()
+    {
+        if (getGlueProxyApiId().isPresent()) {
+            checkState(getGlueRegion().isPresent() && getGlueEndpointUrl().isPresent(),
+                    "Both Glue region and Glue endpoint URL must be provided when Glue proxy API ID is present");
+        }
+    }
 }

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java

@@ -48,7 +48,10 @@ public class GlueMetastoreModule
     @Override
     protected void setup(Binder binder)
     {
-        configBinder(binder).bindConfig(GlueHiveMetastoreConfig.class);
+        GlueHiveMetastoreConfig glueConfig = buildConfigObject(GlueHiveMetastoreConfig.class);
+        glueConfig.getGlueProxyApiId().ifPresent(glueProxyApiId -> binder
+                .bind(Key.get(RequestHandler2.class, ForGlueHiveMetastore.class))
+                .toInstance(new ProxyApiRequestHandler(glueProxyApiId)));
         configBinder(binder).bindConfig(HiveConfig.class);
         binder.bind(AWSCredentialsProvider.class).toProvider(GlueCredentialsProvider.class).in(Scopes.SINGLETON);
         newOptionalBinder(binder, Key.get(RequestHandler2.class, ForGlueHiveMetastore.class));

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ProxyApiRequestHandler.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed 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 io.trino.plugin.hive.metastore.glue;
+
+import com.amazonaws.AmazonWebServiceRequest;
+import com.amazonaws.Request;
+import com.amazonaws.handlers.RequestHandler2;
+
+import static java.util.Objects.requireNonNull;
+
+public class ProxyApiRequestHandler
+        extends RequestHandler2
+{
+    private final String proxyApiId;
+
+    public ProxyApiRequestHandler(String proxyApiId)
+    {
+        this.proxyApiId = requireNonNull(proxyApiId, "proxyApiId is null");
+    }
+
+    @Override
+    public AmazonWebServiceRequest beforeExecution(AmazonWebServiceRequest request)
+    {
+        request.putCustomRequestHeader("x-apigw-api-id", proxyApiId);
+        return request;
+    }
+
+    @Override
+    public void beforeRequest(Request<?> request)
+    {
+        // AWS Glue SDK will append "X-Amz-Target" header to requests (with "AWSGlue" prefix).
+        // This misleads API Gateway (Glue proxy) that it's not the target of the REST call. Therefore, we
+        // need to pass "X-Amz-Target" value in a special HTTP header that is translated back to "X-Amz-Target"
+        // when API Gateway makes request to AWSGlue.
+        request.getHeaders().put("X-Trino-Amz-Target-Proxy", request.getHeaders().remove("X-Amz-Target"));
+    }
+}

plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java

@@ -32,6 +32,7 @@ public void testDefaults()
                 .setGlueEndpointUrl(null)
                 .setGlueStsRegion(null)
                 .setGlueStsEndpointUrl(null)
+                .setGlueProxyApiId(null)
                 .setPinGlueClientToCurrentRegion(false)
                 .setMaxGlueConnections(30)
                 .setMaxGlueErrorRetries(10)
@@ -57,6 +58,7 @@ public void testExplicitPropertyMapping()
                 .put("hive.metastore.glue.endpoint-url", "http://foo.bar")
                 .put("hive.metastore.glue.sts.region", "us-east-3")
                 .put("hive.metastore.glue.sts.endpoint", "http://sts.foo.bar")
+                .put("hive.metastore.glue.proxy-api-id", "abc123")
                 .put("hive.metastore.glue.pin-client-to-current-region", "true")
                 .put("hive.metastore.glue.max-connections", "10")
                 .put("hive.metastore.glue.max-error-retries", "20")
@@ -79,6 +81,7 @@ public void testExplicitPropertyMapping()
                 .setGlueEndpointUrl("http://foo.bar")
                 .setGlueStsRegion("us-east-3")
                 .setGlueStsEndpointUrl("http://sts.foo.bar")
+                .setGlueProxyApiId("abc123")
                 .setPinGlueClientToCurrentRegion(true)
                 .setMaxGlueConnections(10)
                 .setMaxGlueErrorRetries(20)