thds.adls 3.1.20250204215147__py3-none-any.whl → 3.1.20250205222211__py3-none-any.whl

thds/adls/__init__.py

@@ -2,6 +2,7 @@ from thds.core import meta
 
 from . import abfss, defaults, etag, fqn, named_roots, resource, source, uri  # noqa: F401
 from .cached_up_down import download_directory, download_to_cache, upload_through_cache  # noqa: F401
+from .copy import copy_file, copy_files, wait_for_copy  # noqa: F401
 from .errors import BlobNotFoundError  # noqa: F401
 from .fqn import *  # noqa: F401,F403
 from .global_client import get_global_client, get_global_fs_client  # noqa: F401

thds/adls/conf.py

@@ -31,3 +31,12 @@ UPLOAD_FILE_MAX_CONCURRENCY = config.item("upload_file_max_concurrency", 10, par
 UPLOAD_CHUNK_SIZE = config.item("upload_chunk_size", 2**20 * 100, parse=int)  # 100 MB
 
 CONNECTION_TIMEOUT = config.item("connection_timeout", 2000, parse=int)  # seconds
+
+# below are for SAS tokens
+_SAS_EXPIRY_DEFAULT = 900  # seconds
+# 15 minutes seems like a decent default for the time it would take to complete a SAS requiring operation
+USER_DELEGATION_KEY_EXPIRY = config.ConfigItem(
+    "users_delegation_key_expiry", _SAS_EXPIRY_DEFAULT, parse=int
+)  # seconds
+# above key needs to remain valid for signed SAS tokens to remain valid
+BLOB_SAS_EXPIRY = config.ConfigItem("blob_sas_expiry", _SAS_EXPIRY_DEFAULT, parse=int)  # seconds

thds/adls/copy.py

@@ -0,0 +1,183 @@
+"""Functions for copying blobs across remote locations."""
+import datetime
+import random
+import time
+import typing as ty
+
+from azure.storage.blob import BlobSasPermissions, BlobServiceClient, UserDelegationKey
+
+from thds.core import cache, log, parallel, thunks
+
+from .file_properties import exists, get_blob_properties, get_file_properties, is_directory
+from .fqn import AdlsFqn
+from .global_client import get_global_blob_container_client, get_global_blob_service_client
+from .sas_tokens import gen_blob_sas_token, get_user_delegation_key
+from .uri import UriIsh, parse_any
+
+logger = log.getLogger(__name__)
+
+_WAIT_TIMEOUT = 300  # seconds
+# purposely set shorter than SAS token expiry as this is more of a reasonable amount of time for a human to wait
+# rather than a reasonable amount of time for a file to finish copying.
+# the functions using this constant use it as a default argument, so it can be changed as needed.
+
+
+OverwriteMethod = ty.Literal["error", "warn", "skip", "silent"]
+
+CopyRequest = ty.Dict[str, ty.Union[str, datetime.datetime]]
+
+
+class CopyInfo(ty.NamedTuple):
+    src: AdlsFqn
+    dest: AdlsFqn
+    request: CopyRequest
+
+    @property
+    def copy_occurred(self) -> bool:
+        return bool(self.request)
+
+
+def _copy_file(
+    src: AdlsFqn,
+    dest: AdlsFqn,
+    account_key: ty.Union[str, UserDelegationKey],
+    overwrite_method: OverwriteMethod = "error",
+) -> CopyInfo:
+    if not exists(src):
+        raise ValueError(f"{src} does not exist!")
+
+    if is_directory(get_file_properties(src)):
+        raise ValueError(f"{src} is a directory!")
+
+    src_blob_client = get_global_blob_container_client(src.sa, src.container).get_blob_client(src.path)
+    dest_blob_client = get_global_blob_container_client(dest.sa, dest.container).get_blob_client(
+        dest.path
+    )
+
+    def md5s_exist_and_are_equal() -> bool:
+        if (
+            src_md5 := src_blob_client.get_blob_properties().content_settings.content_md5
+        ) and src_md5 == dest_blob_client.get_blob_properties().content_settings.content_md5:
+            return True
+        return False
+
+    if dest_blob_client.exists():
+        if md5s_exist_and_are_equal():
+            # no point in copying if the files are the same
+            logger.info(
+                "%s already exists with the same md5 as the file at %s, no copy will occur", dest, src
+            )
+            return CopyInfo(src, dest, dict())
+
+        logger.info(
+            "%s already exists but it and/or the file at %s have unknown md5s, or the md5s differ. "
+            "The overwrite method provided will determine the copy behavior.",
+            dest,
+            src,
+        )
+
+        if overwrite_method == "error":
+            raise ValueError(f"{dest} already exists!")
+        elif overwrite_method == "warn":
+            logger.warning("%s will be overwritten with the file from %s", dest, src)
+        elif overwrite_method == "skip":
+            logger.warning("%s already exists, skipping copy from %s", dest, src)
+            return CopyInfo(src, dest, dict())
+
+    sas_token = gen_blob_sas_token(
+        src, account_key=account_key, permissions=BlobSasPermissions(read=True)
+    )
+    # we may not always need a SAS token for the copy, but it would be hard to reason about when they are(n't)
+
+    logger.info("Copying %s to %s...", src, dest)
+    return CopyInfo(
+        src, dest, dest_blob_client.start_copy_from_url(f"{src_blob_client.url}?{sas_token}")
+    )
+
+
+def wait_for_copy(fqn: AdlsFqn, timeout: int = _WAIT_TIMEOUT) -> AdlsFqn:
+    """Set timeout to 0 or a negative int to wait indefinitely.
+
+    Keep in mind that, if the copy is authorized using a SAS token, that token expiry exists separate from the wait
+    timeout.
+    """
+    start_time = time.monotonic()
+
+    while (time.monotonic() - start_time < timeout) if timeout > 0 else True:
+        blob_copy_props = get_blob_properties(fqn).copy
+        # blob copy properties actually show copy progress so we could add a progress bar
+
+        if blob_copy_props.status == "success":
+            return fqn
+        elif blob_copy_props.status is None:
+            raise ValueError(f"{fqn} is not an a pending or completed copy.")
+        elif blob_copy_props.status != "pending":
+            raise ValueError(
+                f"The copy to {fqn} failed with the status: {blob_copy_props.status}. "
+                f"See blob copy properties: {blob_copy_props}"
+            )
+
+        time.sleep(random.uniform(1, 3))
+
+    raise TimeoutError(
+        f"Copying to {fqn} did not finish within {timeout} seconds. It may still be copying."
+    )
+
+
+def copy_file(
+    src: UriIsh,
+    dest: UriIsh,
+    overwrite_method: OverwriteMethod = "error",
+    wait: bool = True,
+    timeout: int = _WAIT_TIMEOUT,
+    get_account_key: ty.Callable[
+        [BlobServiceClient], ty.Union[str, UserDelegationKey]
+    ] = get_user_delegation_key,
+) -> CopyInfo:
+    src_fqn = parse_any(src)
+    dest_fqn = parse_any(dest)
+
+    func = thunks.thunking(_copy_file)(
+        src_fqn,
+        dest_fqn,
+        account_key=get_account_key(get_global_blob_service_client(src_fqn.sa)),
+        overwrite_method=overwrite_method,
+    )
+
+    if wait:
+        copy_info = func()
+        wait_for_copy(dest_fqn, timeout=timeout)
+        logger.info("Finished copying %s to %s", src_fqn, dest_fqn)
+        return copy_info
+
+    return func()
+
+
+def copy_files(
+    src_dest_pairs: ty.Collection[ty.Tuple[UriIsh, UriIsh]],
+    overwrite_method: OverwriteMethod = "error",
+    wait: bool = True,
+    timeout: int = _WAIT_TIMEOUT,
+    get_account_key: ty.Callable[
+        [BlobServiceClient], ty.Union[str, UserDelegationKey]
+    ] = get_user_delegation_key,
+) -> ty.List[CopyInfo]:
+    src_dest_fqn_pairs = [(parse_any(src), parse_any(dest)) for src, dest in src_dest_pairs]
+    get_account_key = cache.locking(get_account_key)
+
+    def copy_wrapper(src: AdlsFqn, dest: AdlsFqn) -> CopyInfo:
+        return copy_file(
+            src,
+            dest,
+            overwrite_method=overwrite_method,
+            wait=wait,
+            timeout=timeout,
+            get_account_key=get_account_key,
+        )
+
+    # would be cool to do this async, but using threads for quicker dev
+    return list(
+        parallel.yield_results(
+            [thunks.thunking(copy_wrapper)(src, dest) for src, dest in src_dest_fqn_pairs]
+        )
+    )

thds/adls/file_properties.py

@@ -1,7 +1,10 @@
+from azure.core.exceptions import AzureError, ResourceNotFoundError
+from azure.storage.blob import BlobProperties
 from azure.storage.filedatalake import FileProperties
 
+from .errors import translate_azure_error
 from .fqn import AdlsFqn
-from .global_client import get_global_fs_client
+from .global_client import get_global_blob_container_client, get_global_fs_client
 
 
 def is_directory(info: FileProperties) -> bool:
@@ -11,3 +14,28 @@ def is_directory(info: FileProperties) -> bool:
 
 def get_file_properties(fqn: AdlsFqn) -> FileProperties:
     return get_global_fs_client(fqn.sa, fqn.container).get_file_client(fqn.path).get_file_properties()
+
+
+def get_blob_properties(fqn: AdlsFqn) -> BlobProperties:
+    return (
+        get_global_blob_container_client(fqn.sa, fqn.container)
+        .get_blob_client(fqn.path)
+        .get_blob_properties()
+    )
+
+
+# At some point it may make sense to separate file and blob property modules,
+# but they also are very closely tied together. AFAIK all files are blobs, and given our usage of ADLS,
+# I don't know if we ever deal with things that are blobs but not files.
+
+
+def exists(fqn: AdlsFqn) -> bool:
+    try:
+        get_blob_properties(fqn)
+        # could generally use `get_file_properties` interchangeably,
+        # but blobs are a lower-level primitive than the ADLS file abstraction
+        return True
+    except ResourceNotFoundError:
+        return False
+    except AzureError as err:
+        translate_azure_error(get_global_fs_client(fqn.sa, fqn.container), fqn.path, err)

thds/adls/global_client.py

@@ -58,12 +58,9 @@ get_global_client = cache.locking(adls_fs_client)
 get_global_fs_client = get_global_client
 
 
-def adls_blob_container_client(
-    storage_account: str, container: str, connpool_size: int = DEFAULT_CONNECTION_POOL_SIZE()
-) -> ContainerClient:
-    """This seems to support an atomic write operation,
-    which is in theory going to be faster than two separate network requests.
-    """
+def adls_blob_service_client(
+    storage_account: str, connpool_size: int = DEFAULT_CONNECTION_POOL_SIZE()
+) -> BlobServiceClient:
     return BlobServiceClient(
         account_url=f"https://{storage_account}.blob.core.windows.net",
         credential=SharedCredential(),
@@ -72,7 +69,19 @@ def adls_blob_container_client(
         max_chunk_get_size=conf.MAX_CHUNK_GET_SIZE(),  # for downloads
         max_block_size=conf.MAX_BLOCK_SIZE(),  # for uploads
         max_single_put_size=conf.MAX_SINGLE_PUT_SIZE(),  # for_uploads
-    ).get_container_client(container)
+    )
+
+
+get_global_blob_service_client = cache.locking(adls_blob_service_client)
+
+
+def adls_blob_container_client(
+    storage_account: str, container: str, connpool_size: int = DEFAULT_CONNECTION_POOL_SIZE()
+) -> ContainerClient:
+    """This seems to support an atomic write operation,
+    which is in theory going to be faster than two separate network requests.
+    """
+    return get_global_blob_service_client(storage_account, connpool_size).get_container_client(container)
 
 
 get_global_blob_container_client = cache.locking(adls_blob_container_client)

thds/adls/meta.json

@@ -1,8 +1,8 @@
 {
-  "git_commit": "3cf733adc3c92ae7ff6ebd4bfbeea9d2de71f673",
+  "git_commit": "33bfc2c90d8350529a8067c6bf5f6115a35fd8de",
   "git_branch": "main",
   "git_is_clean": true,
-  "pyproject_version": "3.1.20250204215147",
+  "pyproject_version": "3.1.20250205222211",
   "thds_user": "runner",
   "misc": {}
 }

thds/adls/sas_tokens.py

@@ -0,0 +1,61 @@
+import datetime
+import typing as ty
+
+from azure.storage.blob import (
+    BlobSasPermissions,
+    BlobServiceClient,
+    UserDelegationKey,
+    generate_blob_sas,
+)
+
+from .conf import BLOB_SAS_EXPIRY, USER_DELEGATION_KEY_EXPIRY
+from .fqn import AdlsFqn
+
+
+def get_user_delegation_key(
+    blob_service_client: BlobServiceClient,
+    expiry: ty.Union[int, ty.Callable[[], int]] = USER_DELEGATION_KEY_EXPIRY,
+) -> UserDelegationKey:
+    """For using as an account key to generate SAS tokens.
+
+    This key must remain valid for the SAS tokens it signs to remain valid.
+
+    `expiry` is in seconds.
+    """
+    now = datetime.datetime.now(datetime.timezone.utc)
+
+    return blob_service_client.get_user_delegation_key(
+        key_start_time=now,
+        key_expiry_time=now + datetime.timedelta(seconds=expiry() if callable(expiry) else expiry),
+    )
+
+
+def gen_blob_sas_token(
+    fqn: AdlsFqn,
+    account_key: ty.Union[str, UserDelegationKey],
+    permissions: BlobSasPermissions,
+    expiry: ty.Union[int, ty.Callable[[], int]] = BLOB_SAS_EXPIRY,
+) -> str:
+    """Generates a SAS token for the blob present at the `fqn`.
+
+    If `account_key` is a `UserDelegationKey`, the expiry will be set to the `UserDelegationKey`'s expiry,
+    as when that key expires, the SAS tokens it signs will also expire.
+
+    `expiry` is in seconds.
+    """
+    if isinstance(account_key, UserDelegationKey):
+        expiry_datetime = account_key.signed_expiry
+    else:
+        expiry_datetime = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
+            seconds=expiry() if callable(expiry) else expiry
+        )
+
+    return generate_blob_sas(
+        fqn.sa,
+        fqn.container,
+        fqn.path,
+        account_key=account_key,  # type: ignore[arg-type]
+        # the Azure SDK has too restrictive a type here
+        permission=permissions,
+        expiry=expiry_datetime,
+    )

{thds.adls-3.1.20250204215147.dist-info → thds.adls-3.1.20250205222211.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: thds.adls
-Version: 3.1.20250204215147
+Version: 3.1.20250205222211
 Summary: ADLS tools
 Author: Trilliant Health
 Description-Content-Type: text/markdown

{thds.adls-3.1.20250204215147.dist-info → thds.adls-3.1.20250205222211.dist-info}/RECORD

@@ -1,24 +1,26 @@
-thds/adls/__init__.py,sha256=20oVbjjKkXWMsuhZo9I5NxQDUl5fL4r10hD-jIKPjeM,715
+thds/adls/__init__.py,sha256=Tc23Gy292GO7hC8qwxiAnVE3jGqLcEjHed2MKlSSpvI,784
 thds/adls/_progress.py,sha256=SVOiVRvfvOfUjp-y-UUgA6n0FhAnpKesxltaQ16kmOk,6394
 thds/adls/_upload.py,sha256=q6Sk0CRnNcAjUOUPiBj4CfO4tJD196SQY0lT25CTSE4,4364
 thds/adls/abfss.py,sha256=ZRJOLjDuXmS4bIbQAQpQxWWWeu74N9NKEKCNfXQek80,726
 thds/adls/cached_up_down.py,sha256=qTMrsHHrDq3YZ0gaLMFIPnGNraz0ctG9yoFd3HCVit4,1896
-thds/adls/conf.py,sha256=GmVO1L8nEvC0J0Tv7gUfC7II3yHZkejIddXEDPcCLM4,1522
+thds/adls/conf.py,sha256=q1SPrgb46NpobVzwt_Oyv71-BvsIbZLq9nRWS3LZjz0,1990
+thds/adls/copy.py,sha256=jUWbGvTpb4B3yRGS0nhGSbDzqRPzUqYgH0z1lFRJB3k,6365
 thds/adls/dbfs.py,sha256=pPAjbIZRKJsaXKQljDMUgqS_zy1yKeEZHGMueXbuv3g,2219
 thds/adls/defaults.py,sha256=VxyaGz1gCcz2rGTR9acNRybbxGOYhTldvD_-SdGLBh0,652
 thds/adls/download.py,sha256=Bw4caFGNq5QD9TMzSOVCNV5HVcQsW28WlNDhNItbb0U,16461
 thds/adls/download_lock.py,sha256=_JZj-kjCUfHk9FvrmEuYpJYknmbam5eReFhGNDgzdLQ,2520
 thds/adls/errors.py,sha256=IQfkdLVg0feAUd-xQzkS4-3Lp5_Ld9V9-SqiXeWLRtA,1335
 thds/adls/etag.py,sha256=ct7jpHhNFcKzbekn5rZ3m6DhjK48A7qOZGwDiHkc-pc,242
-thds/adls/file_properties.py,sha256=i3O2VNyAsySRcZwc9Si-V4r0TPdFAbVblZ6Nl_b3KfE,506
+thds/adls/file_properties.py,sha256=JFkobkxcAaRGAh3TJ9mfc3X872gqrGuw4Xk_HSTj1sY,1582
 thds/adls/fqn.py,sha256=pGPBAWKrXorXs1DGkH-XXeDgwwQk3S_heQE7VkPAbW4,5598
-thds/adls/global_client.py,sha256=kl1miL5W7r0k6P0_E3U__zuIEfBbMQlURCokjoli0Rs,3433
+thds/adls/global_client.py,sha256=f4VJw5y_Yh__8gQUcdSYTh1aU6iEPlauMchVirSAwDQ,3716
 thds/adls/impl.py,sha256=x1nSqc8W4NeuX8-JGOp2MRkK8ff6GnelTWedGxPs-qY,42494
 thds/adls/md5.py,sha256=qOX4_7WUj1QkbH_IwREcQNHvvZccOj-HpHZBfsKn1gY,1846
-thds/adls/meta.json,sha256=khEuo_U7-2O_vNDBQrDYJgsyGGw16NlYe4UgaHCzvuA,195
+thds/adls/meta.json,sha256=orDAfHUVcBKGEXl3Q-KNeW71DtfLtvKSp8mipWILANA,195
 thds/adls/named_roots.py,sha256=7SLbAoQQpV_mrFZaUPjYoS-F9dxQxN5Hg4M3YPirF_w,751
 thds/adls/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 thds/adls/ro_cache.py,sha256=F0uXol0t95mcRuBukNg3A7wt7XXQxpD5Sy09d9sl8f0,4825
+thds/adls/sas_tokens.py,sha256=tO7uUh5EVp2F_NITzz1Vks78KCZwVxT7C0faLj8UJqA,1889
 thds/adls/shared_credential.py,sha256=-x42aXoIM001KW59oS8PpuXQd4-F2vg-1gB6OMHlpk4,4602
 thds/adls/source.py,sha256=1JYliDqafqpLewOQ2XMKp4T2WD7b5Rl7XYViyLYtzX8,2437
 thds/adls/uri.py,sha256=pDH956p_VEHnjLLUnjWY6sGgRqksp9gdpc9KOW4gEP0,1205
@@ -29,8 +31,8 @@ thds/adls/resource/up_down.py,sha256=3uNlTvm2gVhSyYdQTBwsGecOgwtINQfINckR-awwV0Y
 thds/adls/tools/download.py,sha256=ZZ8t1g9MRRy-aENGjr10E4aAqyVkKq2OqpyH70pCD_8,965
 thds/adls/tools/ls.py,sha256=OgEaIfTK359twlZIj-A0AW_nv81Z6zi0b9Tw6OJJfWA,1083
 thds/adls/tools/upload.py,sha256=eMk4pdug1aCMPDDWpIE3Zoq77i5APp9Uuh-sVCCDNJE,493
-thds.adls-3.1.20250204215147.dist-info/METADATA,sha256=7wYzQn1dTM8JLN4lGYyBGr5jZZPNf3jwEl3Y89vW0q8,397
-thds.adls-3.1.20250204215147.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-thds.adls-3.1.20250204215147.dist-info/entry_points.txt,sha256=uTqreT1AIwqJboMfLv5w6sviM8mNbAkln765gIjzoA4,152
-thds.adls-3.1.20250204215147.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
-thds.adls-3.1.20250204215147.dist-info/RECORD,,
+thds.adls-3.1.20250205222211.dist-info/METADATA,sha256=pRGWlQvbXlQGBSdYpf_AZRl3KpRCXAY-Jifi5sbf9pw,397
+thds.adls-3.1.20250205222211.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+thds.adls-3.1.20250205222211.dist-info/entry_points.txt,sha256=uTqreT1AIwqJboMfLv5w6sviM8mNbAkln765gIjzoA4,152
+thds.adls-3.1.20250205222211.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
+thds.adls-3.1.20250205222211.dist-info/RECORD,,