Skip to content

Commit

Permalink
Add option to modify how a query is printed
Browse files Browse the repository at this point in the history
  • Loading branch information
MattMonk committed Sep 21, 2024
1 parent 19a5f80 commit 026cb90
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
19 changes: 12 additions & 7 deletions snakemake_interface_storage_plugins/storage_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__()

Expand Down Expand Up @@ -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:
Expand All @@ -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
)


Expand All @@ -189,14 +190,18 @@ 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):
Expand All @@ -219,4 +224,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)
8 changes: 8 additions & 0 deletions snakemake_interface_storage_plugins/storage_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down

0 comments on commit 026cb90

Please sign in to comment.