27707168593

Committed 17 Jun 2026 05:23PM UTC coverage: 95.701% (-0.1%) from 95.806%

Build # 27707168593

Build Type

Pull #887

github

Committed by

web-flow

Commit Message

Merge a648fe80d into 24eb628b0

Pull Request Pull Request #887: [WDL 1.2] narrow FileCoercion lint criteria

Coverage Stats

51 of 61 new or added lines in 2 files covered. (83.61%)

2 existing lines in 2 files now uncovered.

8704 of 9095 relevant lines covered (95.7%)

0.96 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

88.8

/WDL/runtime/download.py

"""
Downloading input files from URIs, with plugin modules for different URI schemes

Download URI plugins are installed & registered using the setuptools entry point group
"miniwdl.plugin.file_download", with name equal to the URI scheme (e.g. "gs" or "s3").

The plugin entry point should be a context manager, which the runtime keeps open for the duration of
the download operation. Given the desired URI, it should quickly yield a tuple with:
    1. source code of a WDL 1.0 task to perform the download
    2. dict of Cromwell-style JSON inputs to give to the task
miniwdl then executes this specified operation, expecting it to produce an output "File file" with
the downloaded file. By doing the heavy lifting in a WDL task, the operation gets to inherit all
the functionality of miniwdl's task runtime, e.g. pulling docker image with binary dependencies,
resource scheduling & isolation, logging, error/signal handling, retry, etc.
The Python context manager itself might be used to obtain and manage the lifetime of any needed
security credentials.
"""

import os
import logging
import traceback
import tempfile
import hashlib
import shlex
from contextlib import ExitStack
from typing import Optional, Generator, Dict, Any, Tuple, Callable
from . import config
from .cache import CallCache
from .._util import compose_coroutines
from .._util import StructuredLogMessage as _


def _load(cfg: config.Loader):
    if getattr(cfg, "_downloaders", None):
        return

    # default public URI downloaders
    file_downloaders = {
        "https": aria2c_downloader,
        "http": aria2c_downloader,
        "ftp": aria2c_downloader,
    }
    directory_downloaders = {}

    # plugins
    for plugin_name, plugin_fn in config.load_plugins(cfg, "file_download"):
        file_downloaders[plugin_name] = plugin_fn
    for plugin_name, plugin_fn in config.load_plugins(cfg, "directory_download"):
        directory_downloaders[plugin_name] = plugin_fn

    setattr(cfg, "_downloaders", (file_downloaders, directory_downloaders))


def _downloader(
    cfg: config.Loader, uri: str, directory: bool = False
) -> Optional[Callable[..., Generator[Dict[str, Any], Dict[str, Any], None]]]:
    _load(cfg)
    colon = uri.find(":")
    if colon <= 0:
        return None
    scheme = uri[:colon]
    return getattr(cfg, "_downloaders")[1 if directory else 0].get(scheme, None)


def able(cfg: config.Loader, uri: Optional[str], directory: bool = False) -> bool:
    """
    Returns True if uri appears to be a URI we know how to download
    """
    return bool(uri and _downloader(cfg, uri, directory=directory) is not None)


def run(
    cfg: config.Loader, logger: logging.Logger, uri: str, directory: bool = False, **kwargs
) -> str:
    """
    Download the URI and return the local filename.

    kwargs are passed through to ``run_local_task``, so ``run_dir`` and ``logger_prefix`` may be
    useful in particular.
    """

    from .error import RunFailed, DownloadFailed, Terminated, error_json
    from .task import run_local_task
    from .. import parse_document, values_from_json, values_to_json, Walker

    gen = _downloader(cfg, uri, directory=directory)
    assert gen
    try:
        logger.info(_(f"start {'directory ' if directory else ''}download", uri=uri))
        with compose_coroutines([lambda kwargs: gen(cfg, logger, **kwargs)], {"uri": uri}) as cor:
            recv = next(cor)

            if "task_wdl" in recv:
                task_wdl, inputs = (recv[k] for k in ["task_wdl", "inputs"])

                doc = parse_document(task_wdl, version="development")
                assert len(doc.tasks) == 1 and not doc.workflow
                doc.typecheck()
                Walker.SetParents()(doc)
                task = doc.tasks[0]
                inputs = values_from_json(inputs, task.available_inputs)  # type: ignore[arg-type]
                subdir, outputs_env = run_local_task(
                    cfg, task, inputs, run_id=("download-" + task.name), **kwargs
                )

                recv = cor.send(
                    {"outputs": values_to_json(outputs_env), "dir": subdir}  # type: ignore[arg-type]
                )

            ans = recv["outputs"]["directory" if directory else "file"]
            assert isinstance(ans, str) and os.path.exists(ans)
            logger.info(_(f"downloaded{' directory' if directory else ' file'}", uri=uri, file=ans))
            return ans

    except RunFailed as exn:
        if isinstance(exn.__cause__, Terminated):
            raise exn.__cause__ from None
        raise DownloadFailed(uri) from exn.__cause__
    except Exception as exn:
        logger.debug(traceback.format_exc())
        logger.error(_("downloader error", uri=uri, **error_json(exn)))
        raise DownloadFailed(uri) from exn


def run_cached(
    cfg,
    logger: logging.Logger,
    cache: CallCache,
    uri: str,
    run_dir: str,
    directory: bool = False,
    **kwargs,
) -> Tuple[bool, str]:
    """
    Cached download logic: returns the file from the cache if available; otherwise, runs the
    download and puts it into the cache before returning
    """
    cached = cache.get_download(uri, directory=directory, logger=logger)
    if cached:
        return True, cached
    cache_path = cache.download_cacheable(uri, directory=directory)
    cache_path_preexists = cache_path and os.path.exists(cache_path)
    if cache_path and not cache_path_preexists:
        # run the download in a holding area under the cache directory, so that it can later be
        # moved atomically into its final cache location
        run_dir = os.path.join(
            cfg["file_io"]["root"], os.path.join(cfg["download_cache"]["dir"], "ops")
        )
    filename = run(cfg, logger, uri, directory=directory, run_dir=run_dir, **kwargs)
    if cache_path_preexists:
        # a cache entry had already existed, but we didn't use it (--no-cache).
        # FIXME: it'd be better to replace the old copy...but what if another workflow is using it?
        logger.warning(
            _(
                "ignored a previously-cached download, which remains in the cache",
                uri=uri,
                downloaded=filename,
                cache_path=cache_path,
            )
        )
        # use the newly downloaded copy in the current run directory
        cache.memo_download(uri, filename, directory=directory)
        return False, filename
    return False, cache.put_download(
        uri, os.path.realpath(filename), directory=directory, logger=logger
    )


# WDL tasks for downloading a file based on its URI scheme


def aria2c_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    wdl = r"""
    task aria2c {
        input {
            String uri
            String docker
            Int connections = 10
        }
        command <<<
            set -euxo pipefail
            mkdir __out
            cd __out
            aria2c -x ~{connections} -s ~{connections} \
                --file-allocation=none --retry-wait=2 --stderr=true --enable-color=false \
                "~{uri}"
        >>>
        output {
            File file = glob("__out/*")[0]
        }
        runtime {
            docker: docker
        }
    }
    """
    recv = yield {
        "task_wdl": wdl,
        "inputs": {"uri": uri, "docker": cfg["download_aria2c"]["docker"]},
    }
    yield recv


def awscli_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    inputs: Dict[str, Any] = {"uri": uri, "docker": cfg["download_awscli"]["docker"]}
    with ExitStack() as cleanup:
        inputs["aws_credentials"] = prepare_aws_credentials(cfg, logger, cleanup)

        wdl = r"""
        task aws_s3_cp {
            input {
                String uri
                String docker
                File? aws_credentials
            }

            command <<<
                set -euo pipefail
                if [ -n "~{aws_credentials}" ]; then
                    source "~{aws_credentials}"
                fi
                export AWS_RETRY_MODE=standard
                export AWS_MAX_ATTEMPTS=5
                set -x
                mkdir __out
                if ! aws s3 cp "~{uri}" __out/ ; then
                    # Retry with --no-sign-request in case the object is public. Without this flag,
                    # the previous invocation could have failed either because (i) no AWS
                    # credentials are available or (ii) the IAM policy restricts accessible S3
                    # buckets regardless of whether the desired object is public.
                    rm -f __out/*
                    >&2 echo 'Retrying with --no-sign-request in case the object is public.' \
                         ' If the overall operation fails, the real error may precede this message.'
                    aws s3 cp --no-sign-request "~{uri}" __out/
                fi
            >>>

            output {
                File file = glob("__out/*")[0]
            }

            runtime {
                docker: docker
            }
        }
        """
        recv = yield {"task_wdl": wdl, "inputs": inputs}
    yield recv


def awscli_directory_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    inputs: Dict[str, Any] = {"uri": uri, "docker": cfg["download_awscli"]["docker"]}
    with ExitStack() as cleanup:
        inputs["aws_credentials"] = prepare_aws_credentials(cfg, logger, cleanup)

        wdl = r"""
        task aws_s3_cp_directory {
            input {
                String uri
                String docker
                File? aws_credentials
            }

            String dnm = basename(uri, "/")

            command <<<
                set -euo pipefail
                if [ -n "~{aws_credentials}" ]; then
                    source "~{aws_credentials}"
                fi
                export AWS_RETRY_MODE=standard
                export AWS_MAX_ATTEMPTS=5
                set -x
                mkdir -p "__out/~{dnm}/"
                if ! aws s3 cp --recursive "~{uri}" "__out/~{dnm}/" ; then
                    # Retry with --no-sign-request in case the object is public. Without this flag,
                    # the previous invocation could have failed either because (i) no AWS
                    # credentials are available or (ii) the IAM policy restricts accessible S3
                    # buckets regardless of whether the desired object is public.
                    rm -f "__out/~{dnm}/*"
                    >&2 echo 'Retrying with --no-sign-request in case the folder is public.' \
                        ' If the overall operation fails, the real error may precede this message.'
                    aws s3 cp --recursive --no-sign-request "~{uri}" "__out/~{dnm}/"
                fi
            >>>

            output {
                Directory directory = "__out/" + dnm
            }

            runtime {
                docker: docker
            }
        }
        """
        recv = yield {"task_wdl": wdl, "inputs": inputs}
    yield recv


def prepare_aws_credentials(
    cfg: config.Loader, logger: logging.Logger, cleanup: ExitStack
) -> Optional[str]:
    host_aws_credentials = {}
    if "AWS_EC2_METADATA_DISABLED" in os.environ:
        # https://github.com/aws/aws-cli/issues/5623
        host_aws_credentials["AWS_EC2_METADATA_DISABLED"] = os.environ["AWS_EC2_METADATA_DISABLED"]
    # get AWS credentials from boto3 (unless prevented by configuration)
    if cfg["download_awscli"].get_bool("host_credentials"):
        import boto3  # type: ignore

        try:
            b3creds = boto3.session.Session().get_credentials()
            host_aws_credentials["AWS_ACCESS_KEY_ID"] = b3creds.access_key
            host_aws_credentials["AWS_SECRET_ACCESS_KEY"] = b3creds.secret_key
            host_aws_credentials["AWS_SESSION_TOKEN"] = b3creds.token
        except Exception:
            pass

    if host_aws_credentials:
        # write credentials to temp file that'll self-destruct afterwards
        host_aws_credentials_str = (
            "\n".join(f"export {k}={shlex.quote(v)}" for (k, v) in host_aws_credentials.items())
            + "\n"
        )
        aws_credentials_file = cleanup.enter_context(
            tempfile.NamedTemporaryFile(
                prefix=hashlib.sha256(host_aws_credentials_str.encode()).hexdigest(),
                delete=True,
                mode="w",
            )
        )
        print(host_aws_credentials_str, file=aws_credentials_file, flush=True)
        # make file group-readable to ensure it'll be usable if the docker image runs as non-root
        os.chmod(aws_credentials_file.name, os.stat(aws_credentials_file.name).st_mode | 0o40)
        logger.getChild("awscli_downloader").info("loaded host AWS credentials")
        return aws_credentials_file.name
    else:
        logger.getChild("awscli_downloader").info(
            "no AWS credentials available via host awscli/boto3; if needed, "
            "configure them and set [download_awscli] host_credentials=true. "
            "(On EC2: awscli might still assume role from instance metadata "
            "service.)"
        )
        return None


def gsutil_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    """
    Built-in downloader plugin for public gs:// URIs; registered by setup.cfg entry_points section

    TODO: adopt security credentials from runtime environment
    """
    if uri == "gs://8675309":
        # hook for test coverage of exception handler
        raise RuntimeError("don't change your number")
    wdl = r"""
    task gsutil_cp {
        input {
            String uri
            String docker
        }
        command <<<
            set -euxo pipefail
            mkdir __out/
            gsutil -q cp "~{uri}" __out/
        >>>
        output {
            File file = glob("__out/*")[0]
        }
        runtime {
            docker: docker
        }
    }
    """
    yield (
        yield {"task_wdl": wdl, "inputs": {"uri": uri, "docker": cfg["download_gsutil"]["docker"]}}
    )


def gsutil_directory_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    """
    Built-in downloader plugin for public gs:// URIs; registered by setup.cfg entry_points section

    TODO: adopt security credentials from runtime environment
    """
    wdl = r"""
    task gsutil_cp {
        input {
            String uri
            String docker
        }

        String dnm = basename(uri, "/")

        command <<<
            set -euxo pipefail
            mkdir __out/
            gsutil -q -m cp -r "~{uri}" __out/
        >>>
        output {
            Directory directory = "__out/" + dnm
        }
        runtime {
            docker: docker
        }
    }
    """
    yield (
        yield {"task_wdl": wdl, "inputs": {"uri": uri, "docker": cfg["download_gsutil"]["docker"]}}
    )

1	"""
2	Downloading input files from URIs, with plugin modules for different URI schemes
3
4	Download URI plugins are installed & registered using the setuptools entry point group
5	"miniwdl.plugin.file_download", with name equal to the URI scheme (e.g. "gs" or "s3").
6
7	The plugin entry point should be a context manager, which the runtime keeps open for the duration of
8	the download operation. Given the desired URI, it should quickly yield a tuple with:
9	1. source code of a WDL 1.0 task to perform the download
10	2. dict of Cromwell-style JSON inputs to give to the task
11	miniwdl then executes this specified operation, expecting it to produce an output "File file" with
12	the downloaded file. By doing the heavy lifting in a WDL task, the operation gets to inherit all
13	the functionality of miniwdl's task runtime, e.g. pulling docker image with binary dependencies,
14	resource scheduling & isolation, logging, error/signal handling, retry, etc.
15	The Python context manager itself might be used to obtain and manage the lifetime of any needed
16	security credentials.
17	"""
18
19	import os	1✔
20	import logging	1✔
21	import traceback	1✔
22	import tempfile	1✔
23	import hashlib	1✔
24	import shlex	1✔
25	from contextlib import ExitStack	1✔
26	from typing import Optional, Generator, Dict, Any, Tuple, Callable	1✔
27	from . import config	1✔
28	from .cache import CallCache	1✔
29	from .._util import compose_coroutines	1✔
30	from .._util import StructuredLogMessage as _	1✔
31
32
33	def _load(cfg: config.Loader):	1✔
34	if getattr(cfg, "_downloaders", None):	1✔
35	return	1✔
36
37	# default public URI downloaders
38	file_downloaders = {	1✔
39	"https": aria2c_downloader,
40	"http": aria2c_downloader,
41	"ftp": aria2c_downloader,
42	}
43	directory_downloaders = {}	1✔
44
45	# plugins
46	for plugin_name, plugin_fn in config.load_plugins(cfg, "file_download"):	1✔
47	file_downloaders[plugin_name] = plugin_fn	1✔
48	for plugin_name, plugin_fn in config.load_plugins(cfg, "directory_download"):	1✔
49	directory_downloaders[plugin_name] = plugin_fn	1✔
50
51	setattr(cfg, "_downloaders", (file_downloaders, directory_downloaders))	1✔
52
53
54	def _downloader(	1✔
55	cfg: config.Loader, uri: str, directory: bool = False
56	) -> Optional[Callable[..., Generator[Dict[str, Any], Dict[str, Any], None]]]:
57	_load(cfg)	1✔
58	colon = uri.find(":")	1✔
59	if colon <= 0:	1✔
60	return None	1✔
61	scheme = uri[:colon]	1✔
62	return getattr(cfg, "_downloaders")[1 if directory else 0].get(scheme, None)	1✔
63
64
65	def able(cfg: config.Loader, uri: Optional[str], directory: bool = False) -> bool:	1✔
66	"""
67	Returns True if uri appears to be a URI we know how to download
68	"""
69	return bool(uri and _downloader(cfg, uri, directory=directory) is not None)	1✔
70
71
72	def run(	1✔
73	cfg: config.Loader, logger: logging.Logger, uri: str, directory: bool = False, **kwargs
74	) -> str:
75	"""
76	Download the URI and return the local filename.
77
78	kwargs are passed through to ``run_local_task``, so ``run_dir`` and ``logger_prefix`` may be
79	useful in particular.
80	"""
81
82	from .error import RunFailed, DownloadFailed, Terminated, error_json	1✔
83	from .task import run_local_task	1✔
84	from .. import parse_document, values_from_json, values_to_json, Walker	1✔
85
86	gen = _downloader(cfg, uri, directory=directory)	1✔
87	assert gen	1✔
88	try:	1✔
89	logger.info(_(f"start {'directory ' if directory else ''}download", uri=uri))	1✔
90	with compose_coroutines([lambda kwargs: gen(cfg, logger, **kwargs)], {"uri": uri}) as cor:	1✔
91	recv = next(cor)	1✔
92
93	if "task_wdl" in recv:	1✔
94	task_wdl, inputs = (recv[k] for k in ["task_wdl", "inputs"])	1✔
95
96	doc = parse_document(task_wdl, version="development")	1✔
97	assert len(doc.tasks) == 1 and not doc.workflow	1✔
98	doc.typecheck()	1✔
99	Walker.SetParents()(doc)	1✔
100	task = doc.tasks[0]	1✔
101	inputs = values_from_json(inputs, task.available_inputs) # type: ignore[arg-type]	1✔
102	subdir, outputs_env = run_local_task(	1✔
103	cfg, task, inputs, run_id=("download-" + task.name), **kwargs
104	)
105
106	recv = cor.send(	1✔
107	{"outputs": values_to_json(outputs_env), "dir": subdir} # type: ignore[arg-type]
108	)
109
110	ans = recv["outputs"]["directory" if directory else "file"]	1✔
111	assert isinstance(ans, str) and os.path.exists(ans)	1✔
112	logger.info(_(f"downloaded{' directory' if directory else ' file'}", uri=uri, file=ans))	1✔
113	return ans	1✔
114
115	except RunFailed as exn:	1✔
116	if isinstance(exn.__cause__, Terminated):	1✔
UNCOV 117	raise exn.__cause__ from None	×
118	raise DownloadFailed(uri) from exn.__cause__	1✔
119	except Exception as exn:	1✔
120	logger.debug(traceback.format_exc())	1✔
121	logger.error(_("downloader error", uri=uri, **error_json(exn)))	1✔
122	raise DownloadFailed(uri) from exn	1✔
123
124
125	def run_cached(	1✔
126	cfg,
127	logger: logging.Logger,
128	cache: CallCache,
129	uri: str,
130	run_dir: str,
131	directory: bool = False,
132	**kwargs,
133	) -> Tuple[bool, str]:
134	"""
135	Cached download logic: returns the file from the cache if available; otherwise, runs the
136	download and puts it into the cache before returning
137	"""
138	cached = cache.get_download(uri, directory=directory, logger=logger)	1✔
139	if cached:	1✔
140	return True, cached	1✔
141	cache_path = cache.download_cacheable(uri, directory=directory)	1✔
142	cache_path_preexists = cache_path and os.path.exists(cache_path)	1✔
143	if cache_path and not cache_path_preexists:	1✔
144	# run the download in a holding area under the cache directory, so that it can later be
145	# moved atomically into its final cache location
146	run_dir = os.path.join(	1✔
147	cfg["file_io"]["root"], os.path.join(cfg["download_cache"]["dir"], "ops")
148	)
149	filename = run(cfg, logger, uri, directory=directory, run_dir=run_dir, **kwargs)	1✔
150	if cache_path_preexists:	1✔
151	# a cache entry had already existed, but we didn't use it (--no-cache).
152	# FIXME: it'd be better to replace the old copy...but what if another workflow is using it?
153	logger.warning(	×
154	_(
155	"ignored a previously-cached download, which remains in the cache",
156	uri=uri,
157	downloaded=filename,
158	cache_path=cache_path,
159	)
160	)
161	# use the newly downloaded copy in the current run directory
162	cache.memo_download(uri, filename, directory=directory)	×
163	return False, filename	×
164	return False, cache.put_download(	1✔
165	uri, os.path.realpath(filename), directory=directory, logger=logger
166	)
167
168
169	# WDL tasks for downloading a file based on its URI scheme
170
171
172	def aria2c_downloader(	1✔
173	cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
174	) -> Generator[Dict[str, Any], Dict[str, Any], None]:
175	wdl = r"""	1✔
176	task aria2c {
177	input {
178	String uri
179	String docker
180	Int connections = 10
181	}
182	command <<<
183	set -euxo pipefail
184	mkdir __out
185	cd __out
186	aria2c -x ~{connections} -s ~{connections} \
187	--file-allocation=none --retry-wait=2 --stderr=true --enable-color=false \
188	"~{uri}"
189	>>>
190	output {
191	File file = glob("__out/*")[0]
192	}
193	runtime {
194	docker: docker
195	}
196	}
197	"""
198	recv = yield {	1✔
199	"task_wdl": wdl,
200	"inputs": {"uri": uri, "docker": cfg["download_aria2c"]["docker"]},
201	}
202	yield recv	1✔
203
204
205	def awscli_downloader(	1✔
206	cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
207	) -> Generator[Dict[str, Any], Dict[str, Any], None]:
208	inputs: Dict[str, Any] = {"uri": uri, "docker": cfg["download_awscli"]["docker"]}	1✔
209	with ExitStack() as cleanup:	1✔
210	inputs["aws_credentials"] = prepare_aws_credentials(cfg, logger, cleanup)	1✔
211
212	wdl = r"""	1✔
213	task aws_s3_cp {
214	input {
215	String uri
216	String docker
217	File? aws_credentials
218	}
219
220	command <<<
221	set -euo pipefail
222	if [ -n "~{aws_credentials}" ]; then
223	source "~{aws_credentials}"
224	fi
225	export AWS_RETRY_MODE=standard
226	export AWS_MAX_ATTEMPTS=5
227	set -x
228	mkdir __out
229	if ! aws s3 cp "~{uri}" __out/ ; then
230	# Retry with --no-sign-request in case the object is public. Without this flag,
231	# the previous invocation could have failed either because (i) no AWS
232	# credentials are available or (ii) the IAM policy restricts accessible S3
233	# buckets regardless of whether the desired object is public.
234	rm -f __out/*
235	>&2 echo 'Retrying with --no-sign-request in case the object is public.' \
236	' If the overall operation fails, the real error may precede this message.'
237	aws s3 cp --no-sign-request "~{uri}" __out/
238	fi
239	>>>
240
241	output {
242	File file = glob("__out/*")[0]
243	}
244
245	runtime {
246	docker: docker
247	}
248	}
249	"""
250	recv = yield {"task_wdl": wdl, "inputs": inputs}	1✔
251	yield recv	1✔
252
253
254	def awscli_directory_downloader(	1✔
255	cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
256	) -> Generator[Dict[str, Any], Dict[str, Any], None]:
257	inputs: Dict[str, Any] = {"uri": uri, "docker": cfg["download_awscli"]["docker"]}	1✔
258	with ExitStack() as cleanup:	1✔
259	inputs["aws_credentials"] = prepare_aws_credentials(cfg, logger, cleanup)	1✔
260
261	wdl = r"""	1✔
262	task aws_s3_cp_directory {
263	input {
264	String uri
265	String docker
266	File? aws_credentials
267	}
268
269	String dnm = basename(uri, "/")
270
271	command <<<
272	set -euo pipefail
273	if [ -n "~{aws_credentials}" ]; then
274	source "~{aws_credentials}"
275	fi
276	export AWS_RETRY_MODE=standard
277	export AWS_MAX_ATTEMPTS=5
278	set -x
279	mkdir -p "__out/~{dnm}/"
280	if ! aws s3 cp --recursive "~{uri}" "__out/~{dnm}/" ; then
281	# Retry with --no-sign-request in case the object is public. Without this flag,
282	# the previous invocation could have failed either because (i) no AWS
283	# credentials are available or (ii) the IAM policy restricts accessible S3
284	# buckets regardless of whether the desired object is public.
285	rm -f "__out/~{dnm}/*"
286	>&2 echo 'Retrying with --no-sign-request in case the folder is public.' \
287	' If the overall operation fails, the real error may precede this message.'
288	aws s3 cp --recursive --no-sign-request "~{uri}" "__out/~{dnm}/"
289	fi
290	>>>
291
292	output {
293	Directory directory = "__out/" + dnm
294	}
295
296	runtime {
297	docker: docker
298	}
299	}
300	"""
301	recv = yield {"task_wdl": wdl, "inputs": inputs}	1✔
302	yield recv	1✔
303
304
305	def prepare_aws_credentials(	1✔
306	cfg: config.Loader, logger: logging.Logger, cleanup: ExitStack
307	) -> Optional[str]:
308	host_aws_credentials = {}	1✔
309	if "AWS_EC2_METADATA_DISABLED" in os.environ:	1✔
310	# https://github.com/aws/aws-cli/issues/5623
311	host_aws_credentials["AWS_EC2_METADATA_DISABLED"] = os.environ["AWS_EC2_METADATA_DISABLED"]	1✔
312	# get AWS credentials from boto3 (unless prevented by configuration)
313	if cfg["download_awscli"].get_bool("host_credentials"):	1✔
314	import boto3 # type: ignore	×
315
316	try:	×
317	b3creds = boto3.session.Session().get_credentials()	×
318	host_aws_credentials["AWS_ACCESS_KEY_ID"] = b3creds.access_key	×
319	host_aws_credentials["AWS_SECRET_ACCESS_KEY"] = b3creds.secret_key	×
320	host_aws_credentials["AWS_SESSION_TOKEN"] = b3creds.token	×
321	except Exception:	×
322	pass	×
323
324	if host_aws_credentials:	1✔
325	# write credentials to temp file that'll self-destruct afterwards
326	host_aws_credentials_str = (	1✔
327	"\n".join(f"export {k}={shlex.quote(v)}" for (k, v) in host_aws_credentials.items())
328	+ "\n"
329	)
330	aws_credentials_file = cleanup.enter_context(	1✔
331	tempfile.NamedTemporaryFile(
332	prefix=hashlib.sha256(host_aws_credentials_str.encode()).hexdigest(),
333	delete=True,
334	mode="w",
335	)
336	)
337	print(host_aws_credentials_str, file=aws_credentials_file, flush=True)	1✔
338	# make file group-readable to ensure it'll be usable if the docker image runs as non-root
339	os.chmod(aws_credentials_file.name, os.stat(aws_credentials_file.name).st_mode \| 0o40)	1✔
340	logger.getChild("awscli_downloader").info("loaded host AWS credentials")	1✔
341	return aws_credentials_file.name	1✔
342	else:
343	logger.getChild("awscli_downloader").info(	×
344	"no AWS credentials available via host awscli/boto3; if needed, "
345	"configure them and set [download_awscli] host_credentials=true. "
346	"(On EC2: awscli might still assume role from instance metadata "
347	"service.)"
348	)
349	return None	×
350
351
352	def gsutil_downloader(	1✔
353	cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
354	) -> Generator[Dict[str, Any], Dict[str, Any], None]:
355	"""
356	Built-in downloader plugin for public gs:// URIs; registered by setup.cfg entry_points section
357
358	TODO: adopt security credentials from runtime environment
359	"""
360	if uri == "gs://8675309":	1✔
361	# hook for test coverage of exception handler
362	raise RuntimeError("don't change your number")	1✔
363	wdl = r"""	1✔
364	task gsutil_cp {
365	input {
366	String uri
367	String docker
368	}
369	command <<<
370	set -euxo pipefail
371	mkdir __out/
372	gsutil -q cp "~{uri}" __out/
373	>>>
374	output {
375	File file = glob("__out/*")[0]
376	}
377	runtime {
378	docker: docker
379	}
380	}
381	"""
382	yield (	1✔
383	yield {"task_wdl": wdl, "inputs": {"uri": uri, "docker": cfg["download_gsutil"]["docker"]}}
384	)
385
386
387	def gsutil_directory_downloader(	1✔
388	cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
389	) -> Generator[Dict[str, Any], Dict[str, Any], None]:
390	"""
391	Built-in downloader plugin for public gs:// URIs; registered by setup.cfg entry_points section
392
393	TODO: adopt security credentials from runtime environment
394	"""
395	wdl = r"""	1✔
396	task gsutil_cp {
397	input {
398	String uri
399	String docker
400	}
401
402	String dnm = basename(uri, "/")
403
404	command <<<
405	set -euxo pipefail
406	mkdir __out/
407	gsutil -q -m cp -r "~{uri}" __out/
408	>>>
409	output {
410	Directory directory = "__out/" + dnm
411	}
412	runtime {
413	docker: docker
414	}
415	}
416	"""
417	yield (	1✔
418	yield {"task_wdl": wdl, "inputs": {"uri": uri, "docker": cfg["download_gsutil"]["docker"]}}
419	)

chanzuckerberg / miniwdl / 27707168593

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous