From ed2e41e2d351c49f174b0be06a4dc2e7143edd25 Mon Sep 17 00:00:00 2001
From: mmonk <matthew.david.monk@cern.ch>
Date: Wed, 18 Sep 2024 09:50:37 +0100
Subject: [PATCH] Add option to modify how a query is printed

---
 .../storage_object.py                             | 15 ++++++++-------
 .../storage_provider.py                           |  8 ++++++++
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/snakemake_interface_storage_plugins/storage_object.py b/snakemake_interface_storage_plugins/storage_object.py
index 66f3293..141295e 100644
--- a/snakemake_interface_storage_plugins/storage_object.py
+++ b/snakemake_interface_storage_plugins/storage_object.py
@@ -68,6 +68,7 @@ def __init__(
         self.keep_local = keep_local
         self.retrieve = retrieve
         self.provider = provider
+        self.print_query = self.provider.safe_print(self.query)
         self._overwrite_local_path = None
         self.__post_init__()
 
@@ -143,21 +144,21 @@ async def managed_size(self) -> int:
             async with self._rate_limiter(Operation.SIZE):
                 return self.size()
         except Exception as e:
-            raise WorkflowError(f"Failed to get size of {self.query}", e)
+            raise WorkflowError(f"Failed to get size of {self.print_query}", e)
 
     async def managed_mtime(self) -> float:
         try:
             async with self._rate_limiter(Operation.MTIME):
                 return self.mtime()
         except Exception as e:
-            raise WorkflowError(f"Failed to get mtime of {self.query}", e)
+            raise WorkflowError(f"Failed to get mtime of {self.print_query}", e)
 
     async def managed_exists(self) -> bool:
         try:
             async with self._rate_limiter(Operation.EXISTS):
                 return self.exists()
         except Exception as e:
-            raise WorkflowError(f"Failed to check existence of {self.query}", e)
+            raise WorkflowError(f"Failed to check existence of {self.print_query}", e)
 
     async def managed_retrieve(self):
         try:
@@ -173,7 +174,7 @@ async def managed_retrieve(self):
                 else:
                     os.remove(local_path)
             raise WorkflowError(
-                f"Failed to retrieve storage object from {self.query}", e
+                f"Failed to retrieve storage object from {self.print_query}", e
             )
 
 
@@ -189,14 +190,14 @@ async def managed_remove(self):
             async with self._rate_limiter(Operation.REMOVE):
                 self.remove()
         except Exception as e:
-            raise WorkflowError(f"Failed to remove storage object {self.query}", e)
+            raise WorkflowError(f"Failed to remove storage object {self.print_query}", e)
 
     async def managed_store(self):
         try:
             async with self._rate_limiter(Operation.STORE):
                 self.store_object()
         except Exception as e:
-            raise WorkflowError(f"Failed to store output in storage {self.query}", e)
+            raise WorkflowError(f"Failed to store output in storage {self.print_query}", e)
 
 
 class StorageObjectGlob(StorageObjectBase):
@@ -219,4 +220,4 @@ async def managed_touch(self):
             async with self._rate_limiter(Operation.TOUCH):
                 self.touch()
         except Exception as e:
-            raise WorkflowError(f"Failed to touch storage object {self.query}", e)
+            raise WorkflowError(f"Failed to touch storage object {self.print_query}", e)
diff --git a/snakemake_interface_storage_plugins/storage_provider.py b/snakemake_interface_storage_plugins/storage_provider.py
index 47d072f..3c0b4ba 100644
--- a/snakemake_interface_storage_plugins/storage_provider.py
+++ b/snakemake_interface_storage_plugins/storage_provider.py
@@ -141,6 +141,14 @@ def postprocess_query(self, query: str) -> str:
         """
         return query
 
+    def safe_print(self, query: str) -> str:
+        """Process the query to remove potentially sensitive information when printing.
+
+        Useful if the query is URL-like and can contain authentication tokens in the
+        parameters and/or usernames/passwords.
+        """
+        return query
+
     @property
     def is_read_write(self) -> bool:
         from snakemake_interface_storage_plugins.storage_object import (