feat: add support for logging into private docker registries

latch/types/__init__.py

@@ -2,6 +2,7 @@
 from latch.types.file import LatchFile, LatchOutputFile
 from latch.types.glob import file_glob
 from latch.types.metadata import (
+    DockerMetadata,
     Fork,
     ForkBranch,
     LatchAppearanceType,

latch/types/metadata.py

@@ -521,10 +521,17 @@ def _parameter_str(t: Tuple[str, LatchParameter]):
         ).strip("\n ")
 
 
+@dataclass
+class DockerMetadata:
+    username: str
+    secret_name: str
+
+
 @dataclass
 class SnakemakeMetadata(LatchMetadata):
     output_dir: Optional[LatchDir] = None
     name: Optional[str] = None
+    docker_metadata: Optional[DockerMetadata] = None
     parameters: Dict[str, SnakemakeParameter] = field(default_factory=dict)
 
     def __post_init__(self):

latch_cli/services/register/register.py

@@ -358,17 +358,21 @@ def register(
             "N/A",
         )
         click.echo(
-            " ".join([
-                click.style("Target workspace:", fg="bright_blue"),
-                ws_name,
-                f"({current_workspace()})",
-            ])
+            " ".join(
+                [
+                    click.style("Target workspace:", fg="bright_blue"),
+                    ws_name,
+                    f"({current_workspace()})",
+                ]
+            )
         )
         click.echo(
-            " ".join([
-                click.style("Workflow root:", fg="bright_blue"),
-                str(ctx.default_container.pkg_dir),
-            ])
+            " ".join(
+                [
+                    click.style("Workflow root:", fg="bright_blue"),
+                    str(ctx.default_container.pkg_dir),
+                ]
+            )
         )
 
         if use_new_centromere:
@@ -386,10 +390,12 @@ def register(
         scp = None
 
         click.echo(
-            " ".join([
-                click.style("Docker Image:", fg="bright_blue"),
-                ctx.default_container.image_name,
-            ])
+            " ".join(
+                [
+                    click.style("Docker Image:", fg="bright_blue"),
+                    ctx.default_container.image_name,
+                ]
+            )
         )
         click.echo()
 

latch_cli/snakemake/workflow.py

@@ -526,7 +526,15 @@ def get_fn_code(
             print(f"JIT Workflow Version: {{jit_wf_version}}")
             print(f"JIT Execution Display Name: {{jit_exec_display_name}}")
 
-            wf = extract_snakemake_workflow(pkg_root, snakefile, jit_wf_version, jit_exec_display_name, local_to_remote_path_mapping, non_blob_parameters, {self.cache_tasks})
+            wf = extract_snakemake_workflow(
+                pkg_root,
+                snakefile,
+                jit_wf_version,
+                jit_exec_display_name,
+                local_to_remote_path_mapping,
+                non_blob_parameters,
+                {self.cache_tasks},
+            )
             wf_name = wf.name
             generate_snakemake_entrypoint(wf, pkg_root, snakefile, {repr(remote_output_url)}, non_blob_parameters)
 
@@ -659,6 +667,7 @@ def __init__(
         self._input_parameters = None
         self._dag = dag
         self._cache_tasks = cache_tasks
+        self._docker_metadata = metadata._snakemake_metadata.docker_metadata
         self.snakemake_tasks: List[SnakemakeJobTask] = []
 
         workflow_metadata = WorkflowMetadata(
@@ -1385,6 +1394,48 @@ def get_fn_code(
                 except Exception:
                     traceback.print_exc()
             lp.upload(compiled, "latch:///.snakemake_latch/workflows/{self.wf.name}/compiled_tasks/{self.name}.py")
+            """,
+            1,
+        )
+
+        if (
+            self.wf._docker_metadata is not None
+            and self.job.container_img_url is not None
+        ):
+            code_block += reindent(
+                rf"""
+
+                print("\n\n\nLogging into Docker\n")
+                from latch.functions.secrets import get_secret
+                docker_usr = "{self.wf._docker_metadata.username}"
+                try:
+                    docker_pwd = get_secret("{self.wf._docker_metadata.secret_name}")
+                except ValueError as e:
+                    print("Failed to get Docker credentials:", e)
+                    sys.exit(1)
+
+                try:
+                    subprocess.run(
+                        [
+                            "docker",
+                            "login",
+                            "--username",
+                            docker_usr,
+                            "--password",
+                            docker_pwd,
+                        ],
+                        check=True,
+                    )
+                except CalledProcessError as e:
+                    print("Failed to login to Docker")
+                except Exception:
+                    traceback.print_exc()
+                """,
+                1,
+            )
+
+        code_block += reindent(
+            rf"""
 
             print("\n\n\nRunning snakemake task\n")
             try: