From 7afeb7ac22a9b4f7e41d04ac1a271aa6e9785355 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                         | 40 ++++++++++++-------
 .../storage_provider.py                       |  8 ++++
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/snakemake_interface_storage_plugins/storage_object.py b/snakemake_interface_storage_plugins/storage_object.py
index 66f3293..8485a2f 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__()
 
@@ -119,7 +120,8 @@ async def inventory(self, cache: IOCacheStorageInterface):
         ...
 
     @abstractmethod
-    def get_inventory_parent(self) -> Optional[str]: ...
+    def get_inventory_parent(self) -> Optional[str]:
+        ...
 
     @abstractmethod
     def cleanup(self):
@@ -127,37 +129,41 @@ def cleanup(self):
         ...
 
     @abstractmethod
-    def exists(self) -> bool: ...
+    def exists(self) -> bool:
+        ...
 
     @abstractmethod
-    def mtime(self) -> float: ...
+    def mtime(self) -> float:
+        ...
 
     @abstractmethod
-    def size(self) -> int: ...
+    def size(self) -> int:
+        ...
 
     @abstractmethod
-    def retrieve_object(self): ...
+    def retrieve_object(self):
+        ...
 
     async def managed_size(self) -> int:
         try:
             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,30 +179,36 @@ 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
             )
 
 
 class StorageObjectWrite(StorageObjectBase):
     @abstractmethod
-    def store_object(self): ...
+    def store_object(self):
+        ...
 
     @abstractmethod
-    def remove(self): ...
+    def remove(self):
+        ...
 
     async def managed_remove(self):
         try:
             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 +231,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 (