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 7afeb7a
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 14 deletions.
40 changes: 26 additions & 14 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 @@ -119,45 +120,50 @@ 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):
"""Perform local cleanup of any remainders of the storage object."""
...

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

Please sign in to comment.