From 12bf1ee5f54280651e514f93586051050b36edf5 Mon Sep 17 00:00:00 2001
From: Fabio Buso <fabio@logicalclocks.com>
Date: Fri, 22 Jan 2021 12:28:37 +0100
Subject: [PATCH] [HOPSWORKS-2197] Add ADLS connector support (#221)

---
 .../logicalclocks/hsfs/StorageConnector.java  |  44 +++++-
 .../hsfs/StorageConnectorType.java            |   3 +-
 .../logicalclocks/hsfs/TrainingDataset.java   |  11 +-
 .../hsfs/engine/SparkEngine.java              | 110 ++++++++------
 .../hsfs/engine/TrainingDatasetEngine.java    |   4 -
 .../com/logicalclocks/hsfs/engine/Utils.java  |  12 ++
 .../logicalclocks/hsfs/metadata/Option.java   |  34 +++++
 python/hsfs/engine/spark.py                   |  22 ++-
 python/hsfs/storage_connector.py              | 139 +++++++++++++++---
 python/hsfs/training_dataset.py               |  11 +-
 10 files changed, 292 insertions(+), 98 deletions(-)
 create mode 100644 java/src/main/java/com/logicalclocks/hsfs/metadata/Option.java

diff --git a/java/src/main/java/com/logicalclocks/hsfs/StorageConnector.java b/java/src/main/java/com/logicalclocks/hsfs/StorageConnector.java
index 467811035c..232ee2b139 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/StorageConnector.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/StorageConnector.java
@@ -17,6 +17,7 @@
 package com.logicalclocks.hsfs;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.logicalclocks.hsfs.metadata.Option;
 import com.logicalclocks.hsfs.util.Constants;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
@@ -28,6 +29,7 @@
 import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 
@@ -120,18 +122,48 @@ public class StorageConnector {
   @Setter
   private String arguments;
 
+  @Getter
+  @Setter
+  private Integer generation;
+
+  @Getter
+  @Setter
+  private String directoryId;
+
+  @Getter
+  @Setter
+  private String applicationId;
+
+  @Getter
+  @Setter
+  private String serviceCredentials;
+
+  @Getter
+  @Setter
+  private String accountName;
+
+  @Getter
+  @Setter
+  private String containerName;
+
+  @Getter
+  @Setter
+  private List<Option> sparkOptions;
+
   @Getter
   @Setter
   private StorageConnectorType storageConnectorType;
 
   @JsonIgnore
-  public Map<String, String> getSparkOptions() throws FeatureStoreException {
-    if (StorageConnectorType.JDBC.equals(storageConnectorType)) {
-      return getJdbcOptions();
-    } else if (StorageConnectorType.REDSHIFT.equals(storageConnectorType)) {
-      return getRedshiftOptions();
+  public Map<String, String> getSparkOptionsInt() throws FeatureStoreException {
+    switch (storageConnectorType) {
+      case JDBC:
+        return getJdbcOptions();
+      case REDSHIFT:
+        return getRedshiftOptions();
+      default:
+        throw new FeatureStoreException("Spark options are not supported for connector " + storageConnectorType);
     }
-    throw new FeatureStoreException("Spark options are not supported for connector " + storageConnectorType);
   }
 
   @JsonIgnore
diff --git a/java/src/main/java/com/logicalclocks/hsfs/StorageConnectorType.java b/java/src/main/java/com/logicalclocks/hsfs/StorageConnectorType.java
index d623ab5e16..0b39c0d240 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/StorageConnectorType.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/StorageConnectorType.java
@@ -20,5 +20,6 @@ public enum StorageConnectorType {
   HOPSFS,
   S3,
   JDBC,
-  REDSHIFT
+  REDSHIFT,
+  ADLS
 }
diff --git a/java/src/main/java/com/logicalclocks/hsfs/TrainingDataset.java b/java/src/main/java/com/logicalclocks/hsfs/TrainingDataset.java
index 8039f0b26f..c4e3dad9cb 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/TrainingDataset.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/TrainingDataset.java
@@ -21,6 +21,7 @@
 import com.logicalclocks.hsfs.engine.StatisticsEngine;
 import com.logicalclocks.hsfs.engine.TrainingDatasetEngine;
 import com.logicalclocks.hsfs.constructor.Query;
+import com.logicalclocks.hsfs.engine.Utils;
 import com.logicalclocks.hsfs.metadata.Statistics;
 import lombok.Builder;
 import lombok.Getter;
@@ -73,7 +74,6 @@ public class TrainingDataset {
 
   @Getter
   @Setter
-  @JsonIgnore
   private StorageConnector storageConnector;
 
   @Getter
@@ -103,6 +103,7 @@ public class TrainingDataset {
 
   private TrainingDatasetEngine trainingDatasetEngine = new TrainingDatasetEngine();
   private StatisticsEngine statisticsEngine = new StatisticsEngine(EntityEndpointType.TRAINING_DATASET);
+  private Utils utils = new Utils();
 
   @Builder
   public TrainingDataset(@NonNull String name, Integer version, String description, DataFormat dataFormat,
@@ -115,13 +116,7 @@ public TrainingDataset(@NonNull String name, Integer version, String description
     this.location = location;
     this.storageConnector = storageConnector;
 
-    if (storageConnector != null) {
-      if (storageConnector.getStorageConnectorType() == StorageConnectorType.S3) {
-        // Default it's already HOPSFS_TRAINING_DATASET
-        this.trainingDatasetType = TrainingDatasetType.EXTERNAL_TRAINING_DATASET;
-      }
-    }
-
+    this.trainingDatasetType = utils.getTrainingDatasetType(storageConnector);
     this.splits = splits;
     this.seed = seed;
     this.featureStore = featureStore;
diff --git a/java/src/main/java/com/logicalclocks/hsfs/engine/SparkEngine.java b/java/src/main/java/com/logicalclocks/hsfs/engine/SparkEngine.java
index 54748f869e..b3b5bb7748 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/engine/SparkEngine.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/engine/SparkEngine.java
@@ -26,10 +26,10 @@
 import com.logicalclocks.hsfs.OnDemandFeatureGroup;
 import com.logicalclocks.hsfs.Split;
 import com.logicalclocks.hsfs.StorageConnector;
-import com.logicalclocks.hsfs.StorageConnectorType;
 import com.logicalclocks.hsfs.TimeTravelFormat;
 import com.logicalclocks.hsfs.TrainingDataset;
 import com.logicalclocks.hsfs.metadata.OnDemandOptions;
+import com.logicalclocks.hsfs.metadata.Option;
 import com.logicalclocks.hsfs.util.Constants;
 import lombok.Getter;
 import org.apache.hadoop.fs.Path;
@@ -80,7 +80,7 @@ public Dataset<Row> sql(String query) {
   }
 
   public Dataset<Row> jdbc(String query, StorageConnector storageConnector) throws FeatureStoreException {
-    Map<String, String> readOptions = storageConnector.getSparkOptions();
+    Map<String, String> readOptions = storageConnector.getSparkOptionsInt();
     if (!Strings.isNullOrEmpty(query)) {
       readOptions.put("query", query);
     }
@@ -124,42 +124,6 @@ public void registerHudiTemporaryTable(FeatureGroup featureGroup, String alias,
         leftFeaturegroupStartTimestamp, leftFeaturegroupEndTimestamp, readOptions);
   }
 
-  public void configureConnector(StorageConnector storageConnector) {
-    if (storageConnector.getStorageConnectorType() == StorageConnectorType.S3) {
-      configureS3Connector(storageConnector);
-    }
-  }
-
-  public static String sparkPath(String path) {
-    if (path.startsWith(Constants.S3_SCHEME)) {
-      return path.replaceFirst(Constants.S3_SCHEME, Constants.S3_SPARK_SCHEME);
-    }
-    return path;
-  }
-
-  private void configureS3Connector(StorageConnector storageConnector) {
-    if (!Strings.isNullOrEmpty(storageConnector.getAccessKey())
-        && Strings.isNullOrEmpty(storageConnector.getSessionToken())) {
-      sparkSession.conf().set(Constants.S3_ACCESS_KEY_ENV, storageConnector.getAccessKey());
-      sparkSession.conf().set(Constants.S3_SECRET_KEY_ENV, storageConnector.getSecretKey());
-    }
-    if (!Strings.isNullOrEmpty(storageConnector.getSessionToken())) {
-      sparkSession.conf().set(Constants.S3_CREDENTIAL_PROVIDER_ENV, Constants.S3_TEMPORARY_CREDENTIAL_PROVIDER);
-      sparkSession.conf().set(Constants.S3_ACCESS_KEY_ENV, storageConnector.getAccessKey());
-      sparkSession.conf().set(Constants.S3_SECRET_KEY_ENV, storageConnector.getSecretKey());
-      sparkSession.conf().set(Constants.S3_SESSION_KEY_ENV, storageConnector.getSessionToken());
-    }
-    if (!Strings.isNullOrEmpty(storageConnector.getServerEncryptionAlgorithm())) {
-      sparkSession.conf().set(
-          "fs.s3a.server-side-encryption-algorithm",
-          storageConnector.getServerEncryptionAlgorithm()
-      );
-    }
-    if (!Strings.isNullOrEmpty(storageConnector.getServerEncryptionKey())) {
-      sparkSession.conf().set("fs.s3a.server-side-encryption.key", storageConnector.getServerEncryptionKey());
-    }
-  }
-
   /**
    * Setup Spark to write the data on the File System.
    *
@@ -171,9 +135,8 @@ private void configureS3Connector(StorageConnector storageConnector) {
   public void write(TrainingDataset trainingDataset, Dataset<Row> dataset,
                     Map<String, String> writeOptions, SaveMode saveMode) {
 
-    if (trainingDataset.getStorageConnector() != null) {
-      SparkEngine.getInstance().configureConnector(trainingDataset.getStorageConnector());
-    }
+    setupConnectorHadoopConf(trainingDataset.getStorageConnector());
+
     if (trainingDataset.getSplits() == null) {
       // Write a single dataset
 
@@ -296,10 +259,7 @@ private void writeSingle(Dataset<Row> dataset, DataFormat dataFormat,
   // OnDemand Feature Group in TFRecords format. However Spark does not use an enum but a string.
   public Dataset<Row> read(StorageConnector storageConnector, String dataFormat,
                            Map<String, String> readOptions, String path) {
-
-    if (storageConnector.getStorageConnectorType() == StorageConnectorType.S3) {
-      configureS3Connector(storageConnector);
-    }
+    setupConnectorHadoopConf(storageConnector);
 
     return SparkEngine.getInstance().getSparkSession()
         .read()
@@ -322,7 +282,7 @@ public Dataset<Row> read(StorageConnector storageConnector, String dataFormat,
   public Map<String, String> getOnlineOptions(Map<String, String> providedWriteOptions,
                                               FeatureGroup featureGroup,
                                               StorageConnector storageConnector) throws FeatureStoreException {
-    Map<String, String> writeOptions = storageConnector.getSparkOptions();
+    Map<String, String> writeOptions = storageConnector.getSparkOptionsInt();
     writeOptions.put(Constants.JDBC_TABLE, utils.getFgName(featureGroup));
 
     // add user provided configuration
@@ -400,4 +360,62 @@ public String profile(Dataset<Row> df, boolean correlation, boolean histogram) {
   public String profile(Dataset<Row> df) {
     return profile(df, null, true, true);
   }
+
+  public void setupConnectorHadoopConf(StorageConnector storageConnector) {
+    if (storageConnector == null) {
+      return;
+    }
+
+    switch (storageConnector.getStorageConnectorType()) {
+      case S3:
+        setupS3ConnectorHadoopConf(storageConnector);
+        break;
+      case ADLS:
+        setupAdlsConnectorHadoopConf(storageConnector);
+        break;
+      default:
+        // No-OP
+        break;
+    }
+  }
+
+  public static String sparkPath(String path) {
+    if (path.startsWith(Constants.S3_SCHEME)) {
+      return path.replaceFirst(Constants.S3_SCHEME, Constants.S3_SPARK_SCHEME);
+    }
+    return path;
+  }
+
+  private void setupS3ConnectorHadoopConf(StorageConnector storageConnector) {
+    if (!Strings.isNullOrEmpty(storageConnector.getAccessKey())) {
+      sparkSession.sparkContext().hadoopConfiguration()
+          .set(Constants.S3_ACCESS_KEY_ENV, storageConnector.getAccessKey());
+    }
+    if (!Strings.isNullOrEmpty(storageConnector.getSecretKey())) {
+      sparkSession.sparkContext().hadoopConfiguration()
+          .set(Constants.S3_SECRET_KEY_ENV, storageConnector.getSecretKey());
+    }
+    if (!Strings.isNullOrEmpty(storageConnector.getServerEncryptionAlgorithm())) {
+      sparkSession.sparkContext().hadoopConfiguration().set(
+          "fs.s3a.server-side-encryption-algorithm",
+          storageConnector.getServerEncryptionAlgorithm()
+      );
+    }
+    if (!Strings.isNullOrEmpty(storageConnector.getServerEncryptionKey())) {
+      sparkSession.sparkContext().hadoopConfiguration()
+          .set("fs.s3a.server-side-encryption.key", storageConnector.getServerEncryptionKey());
+    }
+    if (!Strings.isNullOrEmpty(storageConnector.getSessionToken())) {
+      sparkSession.sparkContext().hadoopConfiguration()
+          .set(Constants.S3_CREDENTIAL_PROVIDER_ENV, Constants.S3_TEMPORARY_CREDENTIAL_PROVIDER);
+      sparkSession.sparkContext().hadoopConfiguration()
+          .set(Constants.S3_SESSION_KEY_ENV, storageConnector.getSessionToken());
+    }
+  }
+
+  private void setupAdlsConnectorHadoopConf(StorageConnector storageConnector) {
+    for (Option confOption : storageConnector.getSparkOptions()) {
+      sparkSession.sparkContext().hadoopConfiguration().set(confOption.getName(), confOption.getValue());
+    }
+  }
 }
diff --git a/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java b/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java
index 9a3788353d..f579e23175 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java
@@ -114,10 +114,6 @@ public void insert(TrainingDataset trainingDataset, Dataset<Row> dataset,
   }
 
   public Dataset<Row> read(TrainingDataset trainingDataset, String split, Map<String, String> providedOptions) {
-    if (trainingDataset.getStorageConnector() != null) {
-      SparkEngine.getInstance().configureConnector(trainingDataset.getStorageConnector());
-    }
-
     String path = "";
     if (com.google.common.base.Strings.isNullOrEmpty(split)) {
       // ** glob means "all sub directories"
diff --git a/java/src/main/java/com/logicalclocks/hsfs/engine/Utils.java b/java/src/main/java/com/logicalclocks/hsfs/engine/Utils.java
index f7d0b5b63d..7ca292d0c3 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/engine/Utils.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/engine/Utils.java
@@ -19,8 +19,10 @@
 import com.logicalclocks.hsfs.Feature;
 import com.logicalclocks.hsfs.FeatureStoreException;
 import com.logicalclocks.hsfs.FeatureGroup;
+import com.logicalclocks.hsfs.StorageConnectorType;
 import com.logicalclocks.hsfs.TrainingDatasetFeature;
 import com.logicalclocks.hsfs.StorageConnector;
+import com.logicalclocks.hsfs.TrainingDatasetType;
 import com.logicalclocks.hsfs.metadata.StorageConnectorApi;
 import org.apache.commons.io.FileUtils;
 import org.apache.spark.sql.Dataset;
@@ -84,6 +86,16 @@ public void trainingDatasetSchemaMatch(Dataset<Row> dataset, List<TrainingDatase
     }
   }
 
+  public TrainingDatasetType getTrainingDatasetType(StorageConnector storageConnector) {
+    if (storageConnector == null) {
+      return TrainingDatasetType.HOPSFS_TRAINING_DATASET;
+    } else if (storageConnector.getStorageConnectorType() == StorageConnectorType.HOPSFS) {
+      return TrainingDatasetType.HOPSFS_TRAINING_DATASET;
+    } else {
+      return TrainingDatasetType.EXTERNAL_TRAINING_DATASET;
+    }
+  }
+
   // TODO(Fabio): this should be moved in the backend
   public String getTableName(FeatureGroup offlineFeatureGroup) {
     return offlineFeatureGroup.getFeatureStore().getName() + "."
diff --git a/java/src/main/java/com/logicalclocks/hsfs/metadata/Option.java b/java/src/main/java/com/logicalclocks/hsfs/metadata/Option.java
new file mode 100644
index 0000000000..ef5b72a2f9
--- /dev/null
+++ b/java/src/main/java/com/logicalclocks/hsfs/metadata/Option.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2021 Logical Clocks AB
+ *
+ * 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.logicalclocks.hsfs.metadata;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.Setter;
+
+@AllArgsConstructor
+@NoArgsConstructor
+public class Option {
+  @Getter
+  @Setter
+  private String name;
+
+  @Getter
+  @Setter
+  private String value;
+}
diff --git a/python/hsfs/engine/spark.py b/python/hsfs/engine/spark.py
index 936c56e131..82fc886945 100644
--- a/python/hsfs/engine/spark.py
+++ b/python/hsfs/engine/spark.py
@@ -305,8 +305,7 @@ def _write_training_dataset_single(
         if data_format.lower() == "tsv":
             data_format = "csv"
 
-        if storage_connector.connector_type == StorageConnector.S3:
-            path = self._setup_s3(storage_connector, path)
+        path = self._setup_storage_connector(storage_connector, path)
 
         feature_dataframe.write.format(data_format).options(**write_options).mode(
             save_mode
@@ -317,8 +316,7 @@ def read(self, storage_connector, data_format, read_options, path):
         if data_format.lower() == "tsv":
             data_format = "csv"
 
-        if storage_connector.connector_type == StorageConnector.S3:
-            path = self._setup_s3(storage_connector, path)
+        path = self._setup_storage_connector(storage_connector, path)
 
         return (
             self._spark_session.read.format(data_format)
@@ -419,7 +417,15 @@ def training_dataset_schema_match(self, dataframe, schema):
 
             i += 1
 
-    def _setup_s3(self, storage_connector, path):
+    def _setup_storage_connector(self, storage_connector, path=None):
+        if storage_connector.connector_type == StorageConnector.S3:
+            return self._setup_s3_hadoop_conf(storage_connector, path)
+        elif storage_connector.connector_type == StorageConnector.ADLS:
+            return self._setup_adls_hadoop_conf(storage_connector, path)
+        else:
+            return path
+
+    def _setup_s3_hadoop_conf(self, storage_connector, path):
         if storage_connector.access_key:
             self._spark_context._jsc.hadoopConfiguration().set(
                 "fs.s3a.access.key", storage_connector.access_key
@@ -449,6 +455,12 @@ def _setup_s3(self, storage_connector, path):
             )
         return path.replace("s3", "s3a", 1)
 
+    def _setup_adls_hadoop_conf(self, storage_connector, path):
+        for k, v in storage_connector.spark_options().items():
+            self._spark_context._jsc.hadoopConfiguration().set(k, v)
+
+        return path
+
     def _setup_pydoop(self):
         # Import Pydoop only here, so it doesn't trigger if the execution environment
         # does not support Pydoop. E.g. Sagemaker
diff --git a/python/hsfs/storage_connector.py b/python/hsfs/storage_connector.py
index 961761bef4..c4a47ffafd 100644
--- a/python/hsfs/storage_connector.py
+++ b/python/hsfs/storage_connector.py
@@ -22,18 +22,20 @@ class StorageConnector:
     S3 = "S3"
     JDBC = "JDBC"
     REDSHIFT = "REDSHIFT"
+    ADLS = "ADLS"
     HOPSFS_DTO = "featurestoreHopsfsConnectorDTO"
     JDBC_DTO = "featurestoreJdbcConnectorDTO"
     S3_DTO = "featurestoreS3ConnectorDTO"
     REDSHIFT_DTO = "featurestoreRedshiftConnectorDTO"
+    ADLS_DTO = "featurestoreADLSConnectorDTO"
 
     def __init__(
         self,
         id,
         name,
-        description,
         featurestore_id,
         storage_connector_type,
+        description=None,
         # members specific to type of connector
         hopsfs_path=None,
         dataset_name=None,
@@ -57,19 +59,32 @@ def __init__(
         connection_string=None,
         arguments=None,
         expiration=None,
+        generation=None,
+        directory_id=None,
+        application_id=None,
+        service_credential=None,
+        account_name=None,
+        container_name=None,
+        spark_options=None,
     ):
         self._id = id
         self._name = name
         self._description = description
         self._feature_store_id = featurestore_id
         self._storage_connector_type = storage_connector_type
+
+        # HopsFS
         self._hopsfs_path = hopsfs_path
         self._dataset_name = dataset_name
+
+        # S3
         self._access_key = access_key
         self._secret_key = secret_key
         self._server_encryption_algorithm = server_encryption_algorithm
         self._server_encryption_key = server_encryption_key
         self._bucket = bucket
+
+        # Redshift
         self._cluster_identifier = cluster_identifier
         self._database_driver = database_driver
         self._database_endpoint = database_endpoint
@@ -86,6 +101,16 @@ def __init__(
         self._arguments = arguments
         self._expiration = expiration
 
+        # ADL
+        self._generation = generation
+        self._directory_id = directory_id
+        self._application_id = application_id
+        self._account_name = account_name
+        self._service_credential = service_credential
+        self._container_name = container_name
+
+        self._spark_options = {opt["name"]: opt["value"] for opt in spark_options}
+
     @classmethod
     def from_response_json(cls, json_dict):
         json_decamelized = humps.decamelize(json_dict)
@@ -110,7 +135,7 @@ def connector_type(self):
     @property
     def access_key(self):
         """Access key."""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return self._access_key
         else:
             raise Exception(
@@ -121,7 +146,7 @@ def access_key(self):
     @property
     def secret_key(self):
         """Secret key."""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return self._secret_key
         else:
             raise Exception(
@@ -132,7 +157,7 @@ def secret_key(self):
     @property
     def server_encryption_algorithm(self):
         """Encryption algorithm if server-side S3 bucket encryption is enabled."""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return self._server_encryption_algorithm
         else:
             raise Exception(
@@ -143,7 +168,7 @@ def server_encryption_algorithm(self):
     @property
     def server_encryption_key(self):
         """Encryption key if server-side S3 bucket encryption is enabled."""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return self._server_encryption_key
         else:
             raise Exception(
@@ -154,7 +179,7 @@ def server_encryption_key(self):
     @property
     def cluster_identifier(self):
         """Cluster identifier for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._cluster_identifier
         else:
             raise Exception(
@@ -165,7 +190,7 @@ def cluster_identifier(self):
     @property
     def database_driver(self):
         """Database endpoint for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_driver
         else:
             raise Exception(
@@ -176,7 +201,7 @@ def database_driver(self):
     @property
     def database_endpoint(self):
         """Database endpoint for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_endpoint
         else:
             raise Exception(
@@ -187,7 +212,7 @@ def database_endpoint(self):
     @property
     def database_name(self):
         """Database name for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_name
         else:
             raise Exception(
@@ -198,7 +223,7 @@ def database_name(self):
     @property
     def database_port(self):
         """Database port for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_port
         else:
             raise Exception(
@@ -209,7 +234,7 @@ def database_port(self):
     @property
     def table_name(self):
         """Table name for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._table_name
         else:
             raise Exception(
@@ -220,7 +245,7 @@ def table_name(self):
     @property
     def database_user_name(self):
         """Database username for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_user_name
         else:
             raise Exception(
@@ -231,7 +256,7 @@ def database_user_name(self):
     @property
     def auto_create(self):
         """Database username for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._auto_create
         else:
             raise Exception(
@@ -242,7 +267,7 @@ def auto_create(self):
     @property
     def database_group(self):
         """Database username for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_group
         else:
             raise Exception(
@@ -253,7 +278,7 @@ def database_group(self):
     @property
     def database_password(self):
         """Database password for redshift cluster."""
-        if self._storage_connector_type.upper() == "REDSHIFT":
+        if self._storage_connector_type.upper() == self.REDSHIFT:
             return self._database_password
         else:
             raise Exception(
@@ -264,7 +289,7 @@ def database_password(self):
     @property
     def session_token(self):
         """Session token."""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return self._session_token
         else:
             raise Exception(
@@ -280,7 +305,7 @@ def iam_role(self):
     @property
     def expiration(self):
         """Cluster temporary credential expiration time."""
-        if self._storage_connector_type.upper() in ["S3", "REDSHIFT"]:
+        if self._storage_connector_type.upper() in [self.S3, self.REDSHIFT]:
             return self._expiration
         else:
             raise Exception(
@@ -291,7 +316,7 @@ def expiration(self):
     @property
     def bucket(self):
         """Return the bucket for S3 connectors."""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return self._bucket
         else:
             raise Exception(
@@ -301,7 +326,7 @@ def bucket(self):
     @property
     def connection_string(self):
         """JDBC connection string."""
-        if self._storage_connector_type.upper() == "JDBC":
+        if self._storage_connector_type.upper() == self.JDBC:
             return self._connection_string
         else:
             raise Exception(
@@ -312,7 +337,7 @@ def connection_string(self):
     @property
     def arguments(self):
         """Additional JDBC arguments."""
-        if self._storage_connector_type.upper() == "JDBC":
+        if self._storage_connector_type.upper() == self.JDBC:
             return self._arguments
         else:
             raise Exception(
@@ -323,25 +348,91 @@ def arguments(self):
     @property
     def path(self):
         """If the connector refers to a path (e.g. S3) - return the path of the connector"""
-        if self._storage_connector_type.upper() == "S3":
+        if self._storage_connector_type.upper() == self.S3:
             return "s3://" + self._bucket
         else:
             raise Exception(
                 "Path is not supported for connector " + self._storage_connector_type
             )
 
+    @property
+    def generation(self):
+        """Generation of the ADLS storage connector"""
+        if self._storage_connector_type.upper() == self.ADLS:
+            self._generation
+        else:
+            raise Exception(
+                "Generation is not supported for connector "
+                + self._storage_connector_type
+            )
+
+    @property
+    def directory_id(self):
+        """Directory ID of the ADLS storage connector"""
+        if self._storage_connector_type.upper() == self.ADLS:
+            self._directory_id
+        else:
+            raise Exception(
+                "Directory ID is not supported for connector "
+                + self._storage_connector_type
+            )
+
+    @property
+    def application_id(self):
+        """Application ID of the ADLS storage connector"""
+        if self._storage_connector_type.upper() == self.ADLS:
+            self._application_id
+        else:
+            raise Exception(
+                "Application ID is not supported for connector "
+                + self._storage_connector_type
+            )
+
+    @property
+    def account_name(self):
+        """Account name of the ADLS storage connector"""
+        if self._storage_connector_type.upper() == self.ADLS:
+            self._account_name
+        else:
+            raise Exception(
+                "Account name is not supported for connector "
+                + self._storage_connector_type
+            )
+
+    @property
+    def container_name(self):
+        """Container name of the ADLS storage connector"""
+        if self._storage_connector_type.upper() == self.ADLS:
+            self._container_name
+        else:
+            raise Exception(
+                "Container name is not supported for connector "
+                + self._storage_connector_type
+            )
+
+    @property
+    def service_credential(self):
+        """Service credential of the ADLS storage connector"""
+        if self._storage_connector_type.upper() == self.ADLS:
+            self._service_credential
+        else:
+            raise Exception(
+                "Service Credential is not supported for connector "
+                + self._storage_connector_type
+            )
+
     def spark_options(self):
         """Return prepared options to be passed to Spark, based on the additional
         arguments.
         """
-        if self._storage_connector_type.upper() == "JDBC":
+        if self._storage_connector_type.upper() == self.JDBC:
             args = [arg.split("=") for arg in self._arguments.split(",")]
 
             options = {a[0]: a[1] for a in args}
             options["url"] = self._connection_string
 
             return options
-        elif self._storage_connector_type.upper() == "REDSHIFT":
+        elif self._storage_connector_type.upper() == self.REDSHIFT:
             connstr = (
                 "jdbc:redshift://"
                 + self._cluster_identifier
@@ -363,6 +454,8 @@ def spark_options(self):
             if self._table_name is not None:
                 props["dbtable"] = self._table_name
             return props
+        elif self._storage_connector_type.upper() == self.ADLS:
+            return self._spark_options
         else:
             raise Exception(
                 "Spark options are not supported for connector "
diff --git a/python/hsfs/training_dataset.py b/python/hsfs/training_dataset.py
index b4ace9a9d6..ac66b27839 100644
--- a/python/hsfs/training_dataset.py
+++ b/python/hsfs/training_dataset.py
@@ -330,12 +330,13 @@ def update_from_response_json(self, json_dict):
         return self
 
     def _infer_training_dataset_type(self, connector_type):
-        if connector_type == StorageConnector.HOPSFS:
+        if connector_type == StorageConnector.HOPSFS or connector_type is None:
             return self.HOPSFS
-        elif connector_type == StorageConnector.S3:
+        elif (
+            connector_type == StorageConnector.S3
+            or connector_type == StorageConnector.ADLS
+        ):
             return self.EXTERNAL
-        elif connector_type is None:
-            return self.HOPSFS
         else:
             raise TypeError(
                 "Storage connectors of type {} are currently not supported for training datasets.".format(
@@ -352,7 +353,7 @@ def to_dict(self):
             "version": self._version,
             "description": self._description,
             "dataFormat": self._data_format,
-            "storageConnectorId": self._storage_connector.id,
+            "storageConnector": self._storage_connector,
             "location": self._location,
             "trainingDatasetType": self._training_dataset_type,
             "features": self._features,