move sas token generating functions from data_utils to adaptors.azure

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

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/'

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