change azure https to wasbs and add azure creds to spark (feast-dev#1258)

jklegar · web-flow · commit 99ff094d5aa5 · 2021-01-07T19:16:24.000-05:00
* change azure https to wasbs and add azure creds to spark

Signed-off-by: Jacob Klegar &lt;jacob@tecton.ai&gt;

* lint

Signed-off-by: Jacob Klegar &lt;jacob@tecton.ai&gt;
diff --git a/.prow/config.yaml b/.prow/config.yaml
@@ -296,7 +296,7 @@ presubmits:
           - name: DOCKER_REPOSITORY
             value: gcr.io/kf-feast
           - name: STAGING_PATH
-            value: https://feastcicd.blob.core.windows.net/staging/cicd-staging
+            value: wasbs://staging@feastcicd.blob.core.windows.net/cicd-staging
           - name: AZ_SERVICE_PRINCIPAL_ID
             valueFrom:
               secretKeyRef:
diff --git a/sdk/python/feast/loaders/file.py b/sdk/python/feast/loaders/file.py
@@ -53,7 +53,7 @@ def export_source_to_staging_location(
             Examples:
                 * gs://bucket/path/
                 * s3://bucket/path/
-                * https://account_name.blob.core.windows.net/bucket/path/
+                * wasbs://bucket@account_name.blob.core.windows.net/path/
                 * file:///data/subfolder/
 
     Returns:
diff --git a/sdk/python/feast/pyspark/launchers/k8s/k8s.py b/sdk/python/feast/pyspark/launchers/k8s/k8s.py
@@ -184,6 +184,20 @@ def _get_staging_client(self):
         uri = urlparse(self._staging_location)
         return get_staging_client(uri.scheme, self._config)
 
+    def _get_azure_credentials(self):
+        uri = urlparse(self._staging_location)
+        if uri.scheme != "wasbs":
+            return {}
+        account_name = self._config.get(opt.AZURE_BLOB_ACCOUNT_NAME)
+        account_key = self._config.get(opt.AZURE_BLOB_ACCOUNT_ACCESS_KEY)
+        if account_name is None or account_key is None:
+            raise Exception(
+                f"Using Azure blob storage requires {opt.AZURE_BLOB_ACCOUNT_NAME} and {opt.AZURE_BLOB_ACCOUNT_ACCESS_KEY} to be set in config"
+            )
+        return {
+            f"spark.hadoop.fs.azure.account.key.{account_name}.blob.core.windows.net": f"{account_key}"
+        }
+
     def historical_feature_retrieval(
         self, job_params: RetrievalJobParameters
     ) -> RetrievalJob:
@@ -221,6 +235,7 @@ def historical_feature_retrieval(
             packages=[],
             jars=[],
             extra_metadata={METADATA_OUTPUT_URI: job_params.get_destination_path()},
+            azure_credentials=self._get_azure_credentials(),
             arguments=job_params.get_arguments(),
             namespace=self._namespace,
         )
@@ -275,6 +290,7 @@ def offline_to_online_ingestion(
             packages=[BQ_SPARK_PACKAGE],
             jars=[],
             extra_metadata={},
+            azure_credentials=self._get_azure_credentials(),
             arguments=ingestion_job_params.get_arguments(),
             namespace=self._namespace,
         )
@@ -317,6 +333,7 @@ def start_stream_to_online_ingestion(
             packages=[BQ_SPARK_PACKAGE],
             jars=extra_jar_paths,
             extra_metadata={METADATA_JOBHASH: job_hash},
+            azure_credentials=self._get_azure_credentials(),
             arguments=ingestion_job_params.get_arguments(),
             namespace=self._namespace,
         )
diff --git a/sdk/python/feast/pyspark/launchers/k8s/k8s_utils.py b/sdk/python/feast/pyspark/launchers/k8s/k8s_utils.py
@@ -111,6 +111,7 @@ def _prepare_job_resource(
     packages: List[str],
     jars: List[str],
     extra_metadata: Dict[str, str],
+    azure_credentials: Dict[str, str],
     arguments: List[str],
     namespace: str,
 ) -> Dict[str, Any]:
@@ -130,6 +131,7 @@ def _prepare_job_resource(
     _add_keys(job, ("spec",), dict(arguments=arguments))
 
     _add_keys(job, ("spec", "sparkConf"), extra_metadata)
+    _add_keys(job, ("spec", "sparkConf"), azure_credentials)
 
     _append_items(job, ("spec", "deps", "packages"), packages)
     _append_items(job, ("spec", "deps", "jars"), jars)
diff --git a/sdk/python/feast/staging/storage_client.py b/sdk/python/feast/staging/storage_client.py
@@ -30,7 +30,7 @@
 GS = "gs"
 S3 = "s3"
 S3A = "s3a"
-AZURE_SCHEME = "https"
+AZURE_SCHEME = "wasbs"
 LOCAL_FILE = "file"
 
 
@@ -326,17 +326,18 @@ def __init__(self, account_name: str, account_access_key: str):
                 "Install package azure-storage-blob for azure blob staging support"
                 "run ```pip install azure-storage-blob```"
             )
-        self.account_url = f"https://{account_name}.blob.core.windows.net"
+        self.account_name = account_name
+        account_url = f"https://{account_name}.blob.core.windows.net"
         self.blob_service_client = BlobServiceClient(
-            account_url=self.account_url, credential=account_access_key
+            account_url=account_url, credential=account_access_key
         )
 
     def download_file(self, uri: ParseResult) -> IO[bytes]:
         """
         Downloads a file from Azure blob storage and returns a TemporaryFile object
 
         Args:
-            uri (urllib.parse.ParseResult): Parsed uri of the file ex: urlparse("https://account_name.blob.core.windows.net/bucket/file.avro")
+            uri (urllib.parse.ParseResult): Parsed uri of the file ex: urlparse("wasbs://bucket@account_name.blob.core.windows.net/file.avro")
 
         Returns:
              TemporaryFile object
@@ -366,17 +367,19 @@ def list_files(self, uri: ParseResult) -> List[str]:
             )
             # File path should not be in path (file path must be longer than path)
             return [
-                f"{self.account_url}/{bucket}/{file}"
+                f"wasbs://{bucket}@{self.account_name}.blob.core.windows.net/{file}"
                 for file in [x.name for x in blob_list]
                 if re.match(regex, file) and file not in path
             ]
         else:
-            return [f"{self.account_url}/{bucket}/{path}"]
+            return [
+                f"wasbs://{bucket}@{self.account_name}.blob.core.windows.net/{path}"
+            ]
 
     def _uri_to_bucket_key(self, uri: ParseResult) -> Tuple[str, str]:
-        assert uri.hostname == urlparse(self.account_url).hostname
-        bucket = uri.path.lstrip("/").split("/")[0]
-        key = uri.path.lstrip("/").split("/", 1)[1]
+        assert uri.hostname == f"{self.account_name}.blob.core.windows.net"
+        bucket = uri.username
+        key = uri.path.lstrip("/")
         return bucket, key
 
     def upload_fileobj(
@@ -485,7 +488,7 @@ def _local_fs_client(config: Config = None):
     GS: _gcs_client,
     S3: _s3_client,
     S3A: _s3a_client,
-    AZURE_SCHEME: _azure_blob_client,  # note we currently interpret all uris beginning https:// as Azure blob uris
+    AZURE_SCHEME: _azure_blob_client,
     LOCAL_FILE: _local_fs_client,
 }
 
@@ -505,5 +508,5 @@ def get_staging_client(scheme, config: Config = None) -> AbstractStagingClient:
         return storage_clients[scheme](config)
     except ValueError:
         raise Exception(
-            f"Could not identify file scheme {scheme}. Only gs://, file://, s3:// and https:// (for Azure) are supported"
+            f"Could not identify file scheme {scheme}. Only gs://, file://, s3:// and wasbs:// (for Azure) are supported"
         )
