thds.adls 3.0.20250116223841__py3-none-any.whl

thds/adls/shared_credential.py

@@ -0,0 +1,107 @@
+import json
+import logging
+import os
+import platform
+import random
+import threading
+import time
+from pathlib import Path
+from typing import Dict
+
+from azure.core.credentials import AccessToken, TokenCredential
+from azure.identity import AzureCliCredential, DefaultAzureCredential
+from azure.identity._constants import EnvironmentVariables
+
+from thds.core import config, files, log
+from thds.core.lazy import Lazy
+
+# suppress noisy Azure identity logs
+logging.getLogger("azure.identity").setLevel(logging.WARNING)
+logger = log.getLogger(__name__)
+DISABLE_FAST_CACHED_CREDENTIAL = config.item("disable_fast", default=False, parse=config.tobool)
+
+
+class ThreadSafeAzureCredential(DefaultAzureCredential):
+    """This is a workaround for how terrible the Azure SDK is.
+
+    They do _not_ play nicely with many threads all trying to load the
+    credentials at once. Thankfully we can just shim in there and make
+    sure that we're not making thousands of parallel requests to the
+    refresh token API...
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.lock = threading.Lock()
+        super().__init__(*args, **kwargs)
+
+    def get_token(self, *args, **kwargs):
+        with self.lock:
+            return super().get_token(*args, **kwargs)
+
+    def close(self):
+        with self.lock:
+            return super().close()
+
+
+_CACHED_TOKEN_PATH = "~/.azure/thds-cache-azure-cli-tokens"
+
+
+class FastCachedAzureCliCredential(AzureCliCredential):
+    # AzureCliCredential only allows a single scope per token, so even though
+    # the abstract concept might name *args as *scopes, for our implementation
+    # it would be an error to try to fetch more than one scope at a time anyway.
+    def get_token(self, scope: str, *args, **kwargs):
+        cached_path = Path(os.path.expanduser(_CACHED_TOKEN_PATH)) / scope.replace("/", "|")
+        # the | (pipe) is just to make the local file paths not-weird.
+        # i chose it essentially at random, and while technically this is a form of compression,
+        # I very much doubt any collisions are possible in the real set of scopes Azure offers.
+        try:
+            token_dict = json.loads(open(cached_path).read())
+            expires_on = token_dict["expires_on"]
+            if expires_on - random.randint(50, 150) < time.time():
+                # we conservatively grab a new token even if we're within 50 to 150 seconds
+                # _before_ the current token expires.
+                # The randomness helps avoid a whole bunch of processes all grabbing it at once,
+                # assuming there's a _lot_ of Azure activity going on in parallel.
+                # There is no correctness concern with this; just optimizations.
+                raise ValueError("Expired Token")
+            return AccessToken(**token_dict)
+        except (FileNotFoundError, TypeError, KeyError, ValueError):
+            fresh_token = super().get_token(scope, *args, **kwargs)
+            # we write atomically b/c we know other processes may be reading or writing
+            # this file.  it's really just a best practice for almost all file writes
+            # where there's any chance of other readers or writers.
+            with files.atomic_write_path(cached_path) as wp, open(wp, "w") as wf:
+                wf.write(json.dumps(fresh_token._asdict()))
+            return fresh_token
+        except Exception as e:
+            logger.exception(f"failed to get fast credential: {e}")
+            raise
+
+
+def _has_workload_identity_creds() -> bool:
+    workload_identity_vars = [
+        EnvironmentVariables.AZURE_TENANT_ID,
+        EnvironmentVariables.AZURE_CLIENT_ID,
+        EnvironmentVariables.AZURE_FEDERATED_TOKEN_FILE,
+    ]
+    return all(var in os.environ for var in workload_identity_vars)
+
+
+def get_credential_kwargs() -> Dict[str, bool]:
+    if _has_workload_identity_creds():
+        # in K8s, we use various forms of credentials, but not the EnvironmentCredential,
+        # and that one gets tried early and then warns us about something we don't care about.
+        return dict(exclude_environment_credential=True)
+    # exclusion due to storage explorer credentials issue when trying to hit East SAs
+    return dict(exclude_shared_token_cache_credential=True)
+
+
+def _SharedCredential() -> TokenCredential:
+    if platform.system() == "Darwin" and not DISABLE_FAST_CACHED_CREDENTIAL():
+        # only try this crazy optimization on our local laptops
+        return FastCachedAzureCliCredential()  # type: ignore
+    return ThreadSafeAzureCredential(**get_credential_kwargs())
+
+
+SharedCredential = Lazy(_SharedCredential)

thds/adls/source.py

@@ -0,0 +1,66 @@
+import base64
+import typing as ty
+from functools import partial
+from pathlib import Path
+
+from thds.core import source
+from thds.core.hashing import Hash, b64
+
+from .cached_up_down import download_to_cache
+from .fqn import AdlsFqn
+from .resource import AdlsHashedResource
+from .uri import resolve_any, resolve_uri
+
+
+def _adls_uri_source_download_handler(uri: str) -> ty.Optional[source.Downloader]:
+    fqn = resolve_uri(uri)
+    if not fqn:
+        return None
+
+    def download(hash: ty.Optional[Hash]) -> Path:
+        assert fqn
+        if hash and hash.algo == "md5":
+            # this 'extra' check just allows us to short-circuit a download
+            # where the hash at this URI is known not to match what we expect.
+            # It's no safer than the non-md5 hash check that Source performs after download.
+            return download_to_cache(fqn, b64(hash.bytes))
+
+        # we don't validate this hash, because we already have md5 validation
+        # happening inside the download_to_cache function. the Source hash
+        # is actually mostly for use by systems that want to do content addressing,
+        # and not necessarily intended to be a runtime check in all scenarios.
+        return download_to_cache(fqn)
+
+    return download
+
+
+source.register_download_handler("thds.adls", _adls_uri_source_download_handler)
+
+
+def from_adls(
+    uri_or_fqn_or_ahr: ty.Union[str, AdlsFqn, AdlsHashedResource], hash: ty.Optional[Hash] = None
+) -> source.Source:
+    """Flexible, public interface to creating Sources from any ADLS-like reference.
+
+    Does NOT automatically fetch an MD5 hash from the ADLS URI if it's not provided. If
+    you know you want to include that, combine this with `resource.get`:
+    `source.from_adls(resource.get(uri))`
+    """
+    if isinstance(uri_or_fqn_or_ahr, AdlsHashedResource):
+        fqn = uri_or_fqn_or_ahr.fqn
+        res_hash = Hash("md5", base64.b64decode(uri_or_fqn_or_ahr.md5b64))
+        if hash and hash != res_hash:
+            raise ValueError(f"Resource Hash mismatch for {fqn}: {hash} != {res_hash}")
+        hash = res_hash
+    else:
+        r_fqn = resolve_any(uri_or_fqn_or_ahr)
+        if not r_fqn:
+            raise ValueError(f"Could not resolve {uri_or_fqn_or_ahr} to an ADLS FQN")
+        fqn = r_fqn
+
+    return source.Source(str(fqn), hash)
+
+
+source.register_from_uri_handler(
+    "thds.adls", lambda uri: partial(from_adls, uri) if resolve_uri(uri) else None
+)

thds/adls/tools/download.py

@@ -0,0 +1,35 @@
+import argparse
+from pathlib import Path
+
+from thds.adls.cached_up_down import download_directory, download_to_cache
+from thds.adls.file_properties import get_file_properties, is_directory
+from thds.adls.uri import resolve_uri
+from thds.core.link import link
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("uri", type=resolve_uri, help="A fully qualified path to an ADLS location")
+    parser.add_argument(
+        "--copy-to",
+        "-c",
+        type=Path,
+        help="This will create a link to the cached download at the specified location",
+    )
+
+    args = parser.parse_args()
+
+    if is_directory(get_file_properties(args.uri)):
+        cache_path = download_directory(args.uri)
+    else:
+        cache_path = download_to_cache(args.uri)
+
+    if args.copy_to:
+        link(cache_path, args.copy_to)
+        print(args.copy_to.resolve())
+    else:
+        print(cache_path.resolve())
+
+
+if __name__ == "__main__":
+    main()

thds/adls/tools/ls.py

@@ -0,0 +1,38 @@
+import argparse
+
+from thds.adls import ADLSFileSystem
+from thds.adls.file_properties import get_file_properties, is_directory
+from thds.adls.uri import resolve_uri
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("uri", type=resolve_uri, help="A fully qualified path to an ADLS location")
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        default=False,
+        action="store_true",
+        help="Verbose mode will display ADLS metadata associated with contents of the specified URI",
+    )
+
+    args = parser.parse_args()
+
+    fs = ADLSFileSystem(args.uri.sa, args.uri.container)
+
+    files_info = (
+        fs.get_directory_info(args.uri.path)
+        if is_directory(get_file_properties(args.uri))
+        else fs.get_files_info([args.uri.path])
+    )
+    for f in files_info:
+        if args.verbose:
+            print(f["name"])
+            [print(f"    {key}: {value}") for key, value in sorted(f.items()) if key != "name"]
+            print("\n")
+        else:
+            print(f["name"])
+
+
+if __name__ == "__main__":
+    main()

thds/adls/uri.py

@@ -0,0 +1,38 @@
+import typing as ty
+
+from . import abfss, dbfs, fqn
+
+HTTPS_SCHEME = "https"
+
+UriIsh = ty.Union[str, fqn.AdlsFqn]
+# Could be an `AdlsFqn`, a supported URI scheme, or potentially any string.
+
+
+def resolve_uri(uri: str) -> ty.Optional[fqn.AdlsFqn]:
+    if uri.startswith(fqn.ADLS_SCHEME):
+        return fqn.parse(uri)
+    elif uri.startswith(abfss.ABFSS_SCHEME):
+        return abfss.to_adls_fqn(uri)
+    elif uri.startswith(HTTPS_SCHEME):
+        uri = uri.replace("https://", "adls://")
+        uri = uri.replace(".blob.core.windows.net", ".dfs.core.windows.net")
+        uri = uri.replace(".dfs.core.windows.net", "")
+        return fqn.parse(uri)
+    elif uri.startswith(dbfs.DBFS_SCHEME):
+        return dbfs.to_adls_fqn(uri)
+    return None
+
+
+def parse_uri(uri: str) -> fqn.AdlsFqn:
+    """Strict/raises error if not a known ADLS URI"""
+    if name := resolve_uri(uri):
+        return name
+    raise fqn.NotAdlsUri(uri)
+
+
+def resolve_any(fqn_or_uri: UriIsh) -> ty.Optional[fqn.AdlsFqn]:
+    return resolve_uri(fqn_or_uri) if isinstance(fqn_or_uri, str) else fqn_or_uri
+
+
+def parse_any(fqn_or_uri: UriIsh) -> fqn.AdlsFqn:
+    return parse_uri(fqn_or_uri) if isinstance(fqn_or_uri, str) else fqn_or_uri

thds.adls-3.0.20250116223841.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.2
+Name: thds.adls
+Version: 3.0.20250116223841
+Summary: ADLS tools
+Author: Trilliant Health
+Description-Content-Type: text/markdown
+Requires-Dist: aiohttp>=3.8.1
+Requires-Dist: aiostream>=0.4.5
+Requires-Dist: azure-identity>=1.9
+Requires-Dist: azure-storage-file-datalake>=12.6
+Requires-Dist: filelock>=3.0
+Requires-Dist: thds.core>=1.31
+
+# adls Library
+
+A port of `core.adls`.

thds.adls-3.0.20250116223841.dist-info/RECORD

@@ -0,0 +1,35 @@
+thds/adls/__init__.py,sha256=20oVbjjKkXWMsuhZo9I5NxQDUl5fL4r10hD-jIKPjeM,715
+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/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/fqn.py,sha256=pGPBAWKrXorXs1DGkH-XXeDgwwQk3S_heQE7VkPAbW4,5598
+thds/adls/global_client.py,sha256=kl1miL5W7r0k6P0_E3U__zuIEfBbMQlURCokjoli0Rs,3433
+thds/adls/impl.py,sha256=x1nSqc8W4NeuX8-JGOp2MRkK8ff6GnelTWedGxPs-qY,42494
+thds/adls/md5.py,sha256=qOX4_7WUj1QkbH_IwREcQNHvvZccOj-HpHZBfsKn1gY,1846
+thds/adls/meta.json,sha256=d0EYRqOCSexOHBq31_qHQTXHzst_Bt6QEBg0GjWpbbU,218
+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/shared_credential.py,sha256=-x42aXoIM001KW59oS8PpuXQd4-F2vg-1gB6OMHlpk4,4602
+thds/adls/source.py,sha256=1JYliDqafqpLewOQ2XMKp4T2WD7b5Rl7XYViyLYtzX8,2437
+thds/adls/uri.py,sha256=pDH956p_VEHnjLLUnjWY6sGgRqksp9gdpc9KOW4gEP0,1205
+thds/adls/resource/__init__.py,sha256=IZ7_aRf1b3jEp7wXOxqHop0gV2gUcf9SOLeEEjIWlCU,1669
+thds/adls/resource/core.py,sha256=uU31tAVkwDcLC_TerZWq58xi4D791vNLDUfP0Y1BAoo,2716
+thds/adls/resource/file_pointers.py,sha256=PLru_3lwut_ZvrX5Keu-wJkPOt5o7UGf-OOT4ixaXME,2049
+thds/adls/resource/up_down.py,sha256=3uNlTvm2gVhSyYdQTBwsGecOgwtINQfINckR-awwV0Y,9907
+thds/adls/tools/download.py,sha256=ZZ8t1g9MRRy-aENGjr10E4aAqyVkKq2OqpyH70pCD_8,965
+thds/adls/tools/ls.py,sha256=OgEaIfTK359twlZIj-A0AW_nv81Z6zi0b9Tw6OJJfWA,1083
+thds.adls-3.0.20250116223841.dist-info/METADATA,sha256=kpsfheQUb9dA-EW5PlWdgInYStAnIv1VkCNrTN7NnJo,397
+thds.adls-3.0.20250116223841.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+thds.adls-3.0.20250116223841.dist-info/entry_points.txt,sha256=-dLxohfrxO0LTHZPGLF8vVXhb05qujJbF8c4fcKqpLQ,106
+thds.adls-3.0.20250116223841.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
+thds.adls-3.0.20250116223841.dist-info/RECORD,,

thds.adls-3.0.20250116223841.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

thds.adls-3.0.20250116223841.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+adls-download-uri = thds.adls.tools.download:main
+adls-ls-uri = thds.adls.tools.ls:main

thds.adls-3.0.20250116223841.dist-info/top_level.txt

@@ -0,0 +1 @@
+thds