From 70ae2b3803b7f6dc5d37786a87873198b3774fa2 Mon Sep 17 00:00:00 2001 From: Doyoung Kim Date: Mon, 27 May 2024 01:09:06 +0000 Subject: [PATCH] move sas token generating functions from data_utils to adaptors.azure --- sky/adaptors/azure.py | 67 ++++++++++++++++++++++++++++++++++++++++++ sky/cloud_stores.py | 5 ++-- sky/data/data_utils.py | 63 --------------------------------------- 3 files changed, 70 insertions(+), 65 deletions(-) diff --git a/sky/adaptors/azure.py b/sky/adaptors/azure.py index 6da151fe188..faf1c1b30eb 100644 --- a/sky/adaptors/azure.py +++ b/sky/adaptors/azure.py @@ -1,6 +1,7 @@ """Azure cli adaptor""" # pylint: disable=import-outside-toplevel +import datetime import functools import logging import threading @@ -114,6 +115,72 @@ def get_client(name: str, raise ValueError(f'Client not supported: "{name}"') +@functools.lru_cache() +@common.load_lazy_modules(modules=_LAZY_MODULES) +def get_az_container_sas_token( + storage_account_name: str, + storage_account_key: str, + container_name: str, +) -> str: + """Returns SAS token used to access container. + + Args: + storage_account_name: str; Name of the storage account + storage_account_key: str; access key for the given storage + account + container_name: str; name of the mounting container + + Returns: + SAS token prepended with the delimiter character, "?" + """ + from azure.storage.blob import generate_container_sas, ContainerSasPermissions + sas_token = generate_container_sas( + account_name=storage_account_name, + container_name=container_name, + account_key=storage_account_key, + permission=ContainerSasPermissions( + read=True, write=True, list=True, create=True), + expiry=datetime.datetime.now( + datetime.timezone.utc) + datetime.timedelta(hours=1) + ) + # "?" is a delimiter character used when SAS token is attached to the + # container endpoint. + # Reference: https://learn.microsoft.com/en-us/azure/ai-services/translator/document-translation/how-to-guides/create-sas-tokens?tabs=Containers # pylint: disable=line-too-long + return f'?{sas_token}' + + +@functools.lru_cache() +@common.load_lazy_modules(modules=_LAZY_MODULES) +def get_az_blob_sas_token(storage_account_name: str, storage_account_key: str, + container_name: str, blob_name: str) -> str: + """Returns SAS token used to access a blob. + + Args: + storage_account_name: str; Name of the storage account + storage_account_key: str; access key for the given storage + account + container_name: str; name of the mounting container + blob_name: str; path to the blob(file) + + Returns: + SAS token prepended with the delimiter character, "?" + """ + from azure.storage.blob import generate_blob_sas, BlobSasPermissions + sas_token = generate_blob_sas( + account_name=storage_account_name, + container_name=container_name, + blob_name=blob_name, + account_key=storage_account_key, + permission=BlobSasPermissions( + read=True, write=True, list=True, create=True), + expiry=datetime.datetime.now( + datetime.timezone.utc) + datetime.timedelta(hours=1) + ) + # "?" is a delimiter character used when SAS token is attached to the + # blob endpoint. + return f'?{sas_token}' + + @common.load_lazy_modules(modules=_LAZY_MODULES) def create_security_rule(**kwargs): from azure.mgmt.network.models import SecurityRule diff --git a/sky/cloud_stores.py b/sky/cloud_stores.py index 7893572a7a7..ce87f2e5ed8 100644 --- a/sky/cloud_stores.py +++ b/sky/cloud_stores.py @@ -12,6 +12,7 @@ import urllib.parse from sky.adaptors import aws +from sky.adaptors import azure from sky.adaptors import cloudflare from sky.adaptors import ibm from sky.clouds import gcp @@ -224,7 +225,7 @@ def make_sync_dir_command(self, source: str, destination: str) -> str: # public containers does not require SAS token for access sas_token = '' else: - sas_token = data_utils.get_az_container_sas_token( + sas_token = azure.get_az_container_sas_token( storage_account_name, storage_account_key, container_name) source = (f'https://{storage_account_name}.blob.core.windows.net/' f'{container_name}/{sas_token}') @@ -248,7 +249,7 @@ def make_sync_file_command(self, source: str, destination: str) -> str: # public containers does not require SAS token for access sas_token = '' else: - sas_token = data_utils.get_az_blob_sas_token( + sas_token = azure.get_az_blob_sas_token( storage_account_name, storage_account_key, container_name, blob_path) source = (f'https://{storage_account_name}.blob.core.windows.net/' diff --git a/sky/data/data_utils.py b/sky/data/data_utils.py index 0f3c0c31a5c..c8bc193d1c0 100644 --- a/sky/data/data_utils.py +++ b/sky/data/data_utils.py @@ -1,7 +1,6 @@ """Miscellaneous Utils for Sky Data """ import concurrent.futures -import datetime from enum import Enum from multiprocessing import pool import os @@ -253,68 +252,6 @@ def get_az_storage_account_key( return storage_account_key -def get_az_container_sas_token( - storage_account_name: str, - storage_account_key: str, - container_name: str, -) -> str: - """Returns SAS token used to access container. - - Args: - storage_account_name: str; Name of the storage account - storage_account_key: str; access key for the given storage - account - container_name: str; name of the mounting container - - Returns: - SAS token prepended with the delimiter character, "?" - """ - from azure.storage.blob import generate_container_sas, ContainerSasPermissions - sas_token = generate_container_sas( - account_name=storage_account_name, - container_name=container_name, - account_key=storage_account_key, - permission=ContainerSasPermissions( - read=True, write=True, list=True, create=True), - expiry=datetime.datetime.now( - datetime.timezone.utc) + datetime.timedelta(hours=1) - ) - # "?" is a delimiter character used when SAS token is attached to the - # container endpoint. - # Reference: https://learn.microsoft.com/en-us/azure/ai-services/translator/document-translation/how-to-guides/create-sas-tokens?tabs=Containers # pylint: disable=line-too-long - return f'?{sas_token}' - - -def get_az_blob_sas_token(storage_account_name: str, storage_account_key: str, - container_name: str, blob_name: str) -> str: - """Returns SAS token used to access a blob. - - Args: - storage_account_name: str; Name of the storage account - storage_account_key: str; access key for the given storage - account - container_name: str; name of the mounting container - blob_name: str; path to the blob(file) - - Returns: - SAS token prepended with the delimiter character, "?" - """ - from azure.storage.blob import generate_blob_sas, BlobSasPermissions - sas_token = generate_blob_sas( - account_name=storage_account_name, - container_name=container_name, - blob_name=blob_name, - account_key=storage_account_key, - permission=BlobSasPermissions( - read=True, write=True, list=True, create=True), - expiry=datetime.datetime.now( - datetime.timezone.utc) + datetime.timedelta(hours=1) - ) - # "?" is a delimiter character used when SAS token is attached to the - # blob endpoint. - return f'?{sas_token}' - - def create_r2_client(region: str = 'auto') -> Client: """Helper method that connects to Boto3 client for R2 Bucket