Add configuration variables for authenticating to Azure Blob

Author: Sanjay Sharma

presto-hive/src/main/java/com/facebook/presto/hive/HdfsConfigurationUpdater.java

@@ -75,9 +75,11 @@ public class HdfsConfigurationUpdater
     private final File s3StagingDirectory;
     private final boolean pinS3ClientToCurrentRegion;
     private final String s3UserAgentPrefix;
+    private final String wasbAccessKey;
+    private final String wasbStorageAccount;
 
     @Inject
-    public HdfsConfigurationUpdater(HiveClientConfig hiveClientConfig, HiveS3Config s3Config)
+    public HdfsConfigurationUpdater(HiveClientConfig hiveClientConfig, HiveS3Config s3Config, HiveWasbConfig wasbConfig)
     {
         requireNonNull(hiveClientConfig, "hiveClientConfig is null");
         checkArgument(hiveClientConfig.getDfsTimeout().toMillis() >= 1, "dfsTimeout must be at least 1 ms");
@@ -115,6 +117,9 @@ public HdfsConfigurationUpdater(HiveClientConfig hiveClientConfig, HiveS3Config
         this.s3StagingDirectory = s3Config.getS3StagingDirectory();
         this.pinS3ClientToCurrentRegion = s3Config.isPinS3ClientToCurrentRegion();
         this.s3UserAgentPrefix = s3Config.getS3UserAgentPrefix();
+
+        this.wasbAccessKey = wasbConfig.getWasbAccessKey();
+        this.wasbStorageAccount = wasbConfig.getWasbStorageAccount();
     }
 
     private static Configuration readConfiguration(List<String> resourcePaths)
@@ -208,6 +213,11 @@ public void updateConfiguration(PrestoHadoopConfiguration config)
         config.setLong(PrestoS3FileSystem.S3_MULTIPART_MIN_PART_SIZE, s3MultipartMinPartSize.toBytes());
         config.setBoolean(PrestoS3FileSystem.S3_PIN_CLIENT_TO_CURRENT_REGION, pinS3ClientToCurrentRegion);
         config.set(PrestoS3FileSystem.S3_USER_AGENT_PREFIX, s3UserAgentPrefix);
+
+        // set config for Azure Blob
+        if (wasbAccessKey != null && wasbStorageAccount != null) {
+            config.set(String.format("fs.azure.account.key.%s.blob.core.windows.net", wasbStorageAccount), wasbAccessKey);
+        }
     }
 
     public static void configureCompression(Configuration config, HiveCompressionCodec compressionCodec)

presto-hive/src/main/java/com/facebook/presto/hive/HiveClientModule.java

@@ -76,6 +76,7 @@ public void configure(Binder binder)
         binder.bind(DirectoryLister.class).to(HadoopDirectoryLister.class).in(Scopes.SINGLETON);
         configBinder(binder).bindConfig(HiveClientConfig.class);
         configBinder(binder).bindConfig(HiveS3Config.class);
+        configBinder(binder).bindConfig(HiveWasbConfig.class);
 
         binder.bind(HiveSessionProperties.class).in(Scopes.SINGLETON);
         binder.bind(HiveTableProperties.class).in(Scopes.SINGLETON);

presto-hive/src/main/java/com/facebook/presto/hive/HiveWasbConfig.java

@@ -0,0 +1,49 @@
+/*
+ * 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 com.facebook.presto.hive;
+
+import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigSecuritySensitive;
+
+public class HiveWasbConfig
+{
+    private String wasbStorageAccount;
+    private String wasbAccessKey;
+
+    public String getWasbStorageAccount()
+    {
+        return wasbStorageAccount;
+    }
+
+    @ConfigSecuritySensitive
+    @Config("hive.wasb.storage-account")
+    public HiveWasbConfig setWasbStorageAccount(String wasbStorageAccount)
+    {
+        this.wasbStorageAccount = wasbStorageAccount;
+        return this;
+    }
+
+    public String getWasbAccessKey()
+    {
+        return wasbAccessKey;
+    }
+
+    @ConfigSecuritySensitive
+    @Config("hive.wasb.access-key")
+    public HiveWasbConfig setWasbAccessKey(String wasbAccessKey)
+    {
+        this.wasbAccessKey = wasbAccessKey;
+        return this;
+    }
+}

presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java

@@ -513,7 +513,7 @@ protected final void setup(String databaseName, HiveClientConfig hiveClientConfi
         setupHive(connectorId.toString(), databaseName, hiveClientConfig.getTimeZone());
 
         metastoreClient = hiveMetastore;
-        HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveClientConfig, new HiveS3Config()));
+        HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveClientConfig, new HiveS3Config(), new HiveWasbConfig()));
         hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, hiveClientConfig, new NoHdfsAuthentication());
         locationService = new HiveLocationService(hdfsEnvironment);
         JsonCodec<PartitionUpdate> partitionUpdateCodec = JsonCodec.jsonCodec(PartitionUpdate.class);

presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClientS3.java

@@ -159,7 +159,7 @@ protected void setup(String host, int port, String databaseName, String awsAcces
         HiveConnectorId connectorId = new HiveConnectorId("hive-test");
         HiveCluster hiveCluster = new TestingHiveCluster(hiveClientConfig, host, port);
         ExecutorService executor = newCachedThreadPool(daemonThreadsNamed("hive-s3-%s"));
-        HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveClientConfig, s3Config));
+        HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveClientConfig, s3Config, new HiveWasbConfig()));
         HivePartitionManager hivePartitionManager = new HivePartitionManager(connectorId, TYPE_MANAGER, hiveClientConfig);
 
         hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, hiveClientConfig, new NoHdfsAuthentication());

presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java

@@ -91,7 +91,7 @@ public static DistributedQueryRunner createQueryRunner(Iterable<TpchTable<?>> ta
 
             HiveClientConfig hiveClientConfig = new HiveClientConfig();
             HiveS3Config s3Config = new HiveS3Config();
-            HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveClientConfig, s3Config));
+            HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveClientConfig, s3Config, new HiveWasbConfig()));
             HdfsEnvironment hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, hiveClientConfig, new NoHdfsAuthentication());
 
             FileHiveMetastore metastore = new FileHiveMetastore(hdfsEnvironment, baseDir.toURI().toString(), "test");

presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java

@@ -93,12 +93,12 @@ public static List<Type> getTypes(List<? extends ColumnHandle> columnHandles)
 
     public static HdfsEnvironment createTestHdfsEnvironment(HiveClientConfig config)
     {
-        return createTestHdfsEnvironment(config, new HiveS3Config());
+        return createTestHdfsEnvironment(config, new HiveS3Config(), new HiveWasbConfig());
     }
 
-    public static HdfsEnvironment createTestHdfsEnvironment(HiveClientConfig hiveConfig, HiveS3Config s3Config)
+    public static HdfsEnvironment createTestHdfsEnvironment(HiveClientConfig hiveConfig, HiveS3Config s3Config, HiveWasbConfig wasbConfig)
     {
-        HdfsConfiguration hdfsConfig = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveConfig, s3Config));
+        HdfsConfiguration hdfsConfig = new HiveHdfsConfiguration(new HdfsConfigurationUpdater(hiveConfig, s3Config, wasbConfig));
         return new HdfsEnvironment(hdfsConfig, hiveConfig, new NoHdfsAuthentication());
     }
 

presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientFileMetastore.java

@@ -27,7 +27,7 @@ protected ExtendedHiveMetastore createMetastore(File tempDir)
     {
         File baseDir = new File(tempDir, "metastore");
         HiveClientConfig hiveConfig = new HiveClientConfig();
-        HdfsConfigurationUpdater updator = new HdfsConfigurationUpdater(hiveConfig, new HiveS3Config());
+        HdfsConfigurationUpdater updator = new HdfsConfigurationUpdater(hiveConfig, new HiveS3Config(), new HiveWasbConfig());
         HdfsConfiguration hdfsConfiguration = new HiveHdfsConfiguration(updator);
         HdfsEnvironment hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, hiveConfig, new NoHdfsAuthentication());
         return new FileHiveMetastore(hdfsEnvironment, baseDir.toURI().toString(), "test");

presto-hive/src/test/java/com/facebook/presto/hive/TestHiveWasbConfig.java

@@ -0,0 +1,49 @@
+/*
+ * 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 com.facebook.presto.hive;
+
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+
+public class TestHiveWasbConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(HiveWasbConfig.class)
+                .setWasbAccessKey(null)
+                .setWasbStorageAccount(null));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+    {
+        Map<String, String> properties = new ImmutableMap.Builder<String, String>()
+                .put("hive.wasb.storage-account", "testwasbstorage")
+                .put("hive.wasb.access-key", "secret")
+                .build();
+
+        HiveWasbConfig expected = new HiveWasbConfig()
+                .setWasbStorageAccount("testwasbstorage")
+                .setWasbAccessKey("secret");
+
+        assertFullMapping(properties, expected);
+    }
+}