rules/merge_inputs: Switch to shared remote file support

Removes the need to custom rules to download from AWS S3 and adds support for 
https, but requires Snakemake >=8.0.0.

Author: joverlee521

rules/main.smk

@@ -1,3 +1,8 @@
+from snakemake.utils import min_version
+
+# Minimum Snakemake version needed for the storage plugins used in remote_files.smk
+min_version("8.0.0")
+
 # constrain the wildcards to not include `_` which we use to separate "parts" of filenames (where a part may be a wildcard itself)
 wildcard_constraints:
     subtype = "[^_/]+",
@@ -14,7 +19,6 @@ SEGMENTS = ["pb2", "pb1", "pa", "ha","np", "na", "mp", "ns"]
 # (2) Filter the other segments by simply force-including the same strains as (1)
 SAME_STRAINS = bool(config.get('same_strains_per_segment', False))
 
-NEXTSTRAIN_PUBLIC_BUCKET = "s3://nextstrain-data/"
 
 rule all:
     input:
@@ -33,6 +37,7 @@ rule test_target:
 class InvalidConfigError(Exception):
     pass
 
+include: "../shared/vendored/snakemake/remote_files.smk"
 # This uses the `InvalidConfigError` defined above
 include: "merge_inputs.smk"
 

rules/merge_inputs.smk

@@ -13,60 +13,50 @@ def _parse_config_input(input):
       in which case it must not include the wildcard substring.
 
     Returns a dictionary with optional keys:
-    - metadata:string - the relative path to the metadata file. If the original data was remote then this represents
-      the output of a rule which downloads the file
-    - metadata_location:string - the URI for the remote file if applicable else `None`
-    - sequences:function. Takes in wildcards and returns the relative path to the sequences FASTA for the provided
+    - metadata:string - path or url to the metadata file.
+    - sequences:function. Takes in wildcards and returns path or url to the sequences FASTA for the provided
       segment wildcard, or returns `None` if this input doesn't define sequences for the provided segment.
-    - sequences_location:function. Takes in wildcards and returns the URI for the remote file, or `None`, where applicable.
 
     Raises InvalidConfigError
     """
-    name = input['name']
-    lambda_none = lambda w: None
-
-    info = {'metadata': None, 'metadata_location': None, 'sequences': lambda_none, 'sequences_location': lambda_none}
-
-    def _source(uri, *,  s3, local):
-        if uri.startswith('s3://'):
-            return s3
-        elif uri.lower().startswith('http://') or uri.lower().startswith('https://'):
-            raise InvalidConfigError("Workflow cannot yet handle HTTP[S] inputs")
-        return local
-
-    if location:=input.get('metadata', False):
-        info['metadata'] = _source(location,  s3=f"data/{name}/metadata.tsv", local=location)
-        info['metadata_location'] = _source(location,  s3=location, local=None)
+    info = {
+        "name": input["name"],
+        "metadata": path_or_url(input["metadata"]) if input.get("metadata") else None,
+        "sequences": None,
+    }
 
     if location:=input.get('sequences', False):
         if isinstance(location, dict):
-            info['sequences'] = lambda w: _source(location[w.segment],  s3=f"data/{name}/sequences_{w.segment}.fasta", local=location[w.segment]) \
-                if w.segment in location \
-                else None
-            info['sequences_location'] = lambda w: _source(location[w.segment], s3=location[w.segment], local=None) \
+            info['sequences'] = lambda w: path_or_url(location[w.segment]) \
                 if w.segment in location \
                 else None
         elif isinstance(location, str):
-            info['sequences'] = _source(location, s3=lambda w: f"data/{name}/sequences_{w.segment}.fasta", local=lambda w: location.format(segment=w.segment))
-            info['sequences_location'] = _source(location,  s3=lambda w: location.format(segment=w.segment), local=lambda_none)
+            info['sequences'] = lambda w: path_or_url(location)
         else:
-            raise InvalidConfigError(f"Config input for {name} specifies sequences in an unknown format; must be dict or string")
+            raise InvalidConfigError(f"Config input for {info['name']} specifies sequences in an unknown format; must be dict or string")
 
     return info
 
-
 def _gather_inputs():
     all_inputs = [*config['inputs'], *config.get('additional_inputs', [])]
 
     if len(all_inputs)==0:
         raise InvalidConfigError("Config must define at least one element in config.inputs or config.additional_inputs lists")
     if not all([isinstance(i, dict) for i in all_inputs]):
-        raise InvalidConfigError("All of the elements in config.inputs and config.additional_inputs lists must be dictionaries"
+        raise InvalidConfigError("All of the elements in config.inputs and config.additional_inputs lists must be dictionaries. "
             "If you've used a command line '--config' double check your quoting.")
     if len({i['name'] for i in all_inputs})!=len(all_inputs):
         raise InvalidConfigError("Names of inputs (config.inputs and config.additional_inputs) must be unique")
     if not all(['name' in i and ('sequences' in i or 'metadata' in i) for i in all_inputs]):
         raise InvalidConfigError("Each input (config.inputs and config.additional_inputs) must have a 'name' and 'metadata' and/or 'sequences'")
+    if not any(['metadata' in i for i in all_inputs]):
+        raise InvalidConfigError("At least one input must have 'metadata'")
+    if not any (['sequences' in i for i in all_inputs]):
+        raise InvalidConfigError("At least one input must have 'sequences'")
+
+    available_keys = set(['name', 'metadata', 'sequences'])
+    if any([len(set(el.keys())-available_keys)>0 for el in all_inputs]):
+        raise InvalidConfigError(f"Each input (config.inputs and config.additional_inputs) can only include keys of {', '.join(available_keys)}")
 
     return {i['name']: _parse_config_input(i) for i in all_inputs}
 
@@ -80,32 +70,6 @@ def input_sequences(wildcards):
     inputs = list(filter(None, [info['sequences'](wildcards) for info in input_sources.values() if info.get('sequences', None)]))
     return inputs[0] if len(inputs)==1 else "results/sequences_merged_{segment}.fasta"
 
-rule download_s3_sequences:
-    output:
-        sequences = "data/{input_name}/sequences_{segment}.fasta",
-    params:
-        address = lambda w: input_sources[w.input_name]['sequences_location'](w),
-        no_sign_request=lambda w: "--no-sign-request" \
-            if input_sources[w.input_name]['sequences_location'](w).startswith(NEXTSTRAIN_PUBLIC_BUCKET) \
-            else "",
-    shell:
-        """
-        aws s3 cp {params.no_sign_request:q} {params.address:q} - | zstd -d > {output.sequences}
-        """
-
-rule download_s3_metadata:
-    output:
-        metadata = "data/{input_name}/metadata.tsv",
-    params:
-        address = lambda w: input_sources[w.input_name]['metadata_location'],
-        no_sign_request=lambda w: "--no-sign-request" \
-            if input_sources[w.input_name]['metadata_location'].startswith(NEXTSTRAIN_PUBLIC_BUCKET) \
-            else "",
-    shell:
-        """
-        aws s3 cp {params.no_sign_request:q} {params.address:q} - | zstd -d > {output.metadata}
-        """
-
 rule merge_metadata:
     """
     This rule should only be invoked if there are multiple defined metadata inputs