toil 7.0.0__py3-none-any.whl → 8.1.0b1__py3-none-any.whl

toil/jobStores/googleJobStore.py

@@ -20,20 +20,24 @@ import uuid
 from contextlib import contextmanager
 from functools import wraps
 from io import BytesIO
-from typing import IO, List, Optional
+from typing import IO, Optional
 from urllib.parse import ParseResult
 
-from google.api_core.exceptions import (GoogleAPICallError,
-                                        InternalServerError,
-                                        ServiceUnavailable)
+from google.api_core.exceptions import (
+    GoogleAPICallError,
+    InternalServerError,
+    ServiceUnavailable,
+)
 from google.auth.exceptions import DefaultCredentialsError
 from google.cloud import exceptions, storage
 
-from toil.jobStores.abstractJobStore import (AbstractJobStore,
-                                             JobStoreExistsException,
-                                             NoSuchFileException,
-                                             NoSuchJobException,
-                                             NoSuchJobStoreException)
+from toil.jobStores.abstractJobStore import (
+    AbstractJobStore,
+    JobStoreExistsException,
+    NoSuchFileException,
+    NoSuchJobException,
+    NoSuchJobStoreException,
+)
 from toil.jobStores.utils import ReadablePipe, WritablePipe
 from toil.lib.compatibility import compat_bytes
 from toil.lib.io import AtomicFileCreate
@@ -42,7 +46,7 @@ from toil.lib.retry import old_retry
 
 log = logging.getLogger(__name__)
 
-GOOGLE_STORAGE = 'gs'
+GOOGLE_STORAGE = "gs"
 
 MAX_BATCH_SIZE = 1000
 
@@ -75,19 +79,22 @@ def google_retry(f):
 
     It should wrap any function that makes use of the Google Client API
     """
+
     @wraps(f)
     def wrapper(*args, **kwargs):
-        for attempt in old_retry(delays=truncExpBackoff(),
-                                 timeout=300,
-                                 predicate=google_retry_predicate):
+        for attempt in old_retry(
+            delays=truncExpBackoff(), timeout=300, predicate=google_retry_predicate
+        ):
             with attempt:
                 return f(*args, **kwargs)
+
     return wrapper
 
 
 class GoogleJobStore(AbstractJobStore):
 
-    nodeServiceAccountJson = '/root/service_account.json'
+    nodeServiceAccountJson = "/root/service_account.json"
+
     def __init__(self, locator: str) -> None:
         super().__init__(locator)
 
@@ -99,20 +106,19 @@ class GoogleJobStore(AbstractJobStore):
             projectID = None
 
         self.projectID = projectID
-        self.bucketName = namePrefix+"--toil"
+        self.bucketName = namePrefix + "--toil"
         log.debug("Instantiating google jobStore with name: %s", self.bucketName)
 
         # this is a :class:`~google.cloud.storage.bucket.Bucket`
         self.bucket = None
 
-        self.statsBaseID = 'f16eef0c-b597-4b8b-9b0c-4d605b4f506c'
-        self.statsReadPrefix = '_'
-        self.readStatsBaseID = self.statsReadPrefix+self.statsBaseID
+        self.statsBaseID = "f16eef0c-b597-4b8b-9b0c-4d605b4f506c"
+        self.statsReadPrefix = "_"
+        self.readStatsBaseID = self.statsReadPrefix + self.statsBaseID
 
         self.sseKey = None
         self.storageClient = self.create_client()
 
-
     @classmethod
     def create_client(cls) -> storage.Client:
         """
@@ -127,28 +133,36 @@ class GoogleJobStore(AbstractJobStore):
         # Determine if we have an override environment variable for our credentials.
         # We get the path to check existence, but Google Storage works out what
         # to use later by looking at the environment again.
-        credentials_path: Optional[str] = os.getenv('GOOGLE_APPLICATION_CREDENTIALS', None)
+        credentials_path: Optional[str] = os.getenv(
+            "GOOGLE_APPLICATION_CREDENTIALS", None
+        )
         if credentials_path is not None and not os.path.exists(credentials_path):
             # If the file is missing, complain.
             # This variable holds a file name and not any sensitive data itself.
-            log.warning("File '%s' from GOOGLE_APPLICATION_CREDENTIALS is unavailable! "
-                        "We may not be able to authenticate!",
-                        credentials_path)
+            log.warning(
+                "File '%s' from GOOGLE_APPLICATION_CREDENTIALS is unavailable! "
+                "We may not be able to authenticate!",
+                credentials_path,
+            )
 
         if credentials_path is None and os.path.exists(cls.nodeServiceAccountJson):
             try:
                 # load credentials from a particular file on GCE nodes if an override path is not set
-                return storage.Client.from_service_account_json(cls.nodeServiceAccountJson)
+                return storage.Client.from_service_account_json(
+                    cls.nodeServiceAccountJson
+                )
             except OSError:
                 # Probably we don't have permission to use the file.
-                log.warning("File '%s' exists but didn't work to authenticate!",
-                            cls.nodeServiceAccountJson)
+                log.warning(
+                    "File '%s' exists but didn't work to authenticate!",
+                    cls.nodeServiceAccountJson,
+                )
 
         # Either a filename is specified, or our fallback file isn't there.
         try:
             # See if Google can work out how to authenticate.
             return storage.Client()
-        except (DefaultCredentialsError, EnvironmentError):
+        except (DefaultCredentialsError, OSError):
             # Depending on which Google codepath or module version (???)
             # realizes we have no credentials, we can get an EnvironemntError,
             # or the new DefaultCredentialsError we are supposedly specced to
@@ -158,7 +172,6 @@ class GoogleJobStore(AbstractJobStore):
             # This is likely to happen all the time so don't warn.
             return storage.Client.create_anonymous_client()
 
-
     @google_retry
     def initialize(self, config=None):
         try:
@@ -169,7 +182,7 @@ class GoogleJobStore(AbstractJobStore):
 
         # set up sever side encryption after we set up config in super
         if self.config.sseKey is not None:
-            with open(self.config.sseKey, 'rb') as f:
+            with open(self.config.sseKey, "rb") as f:
                 self.sseKey = compat_bytes(f.read())
                 assert len(self.sseKey) == 32
 
@@ -199,13 +212,13 @@ class GoogleJobStore(AbstractJobStore):
             count = 0
             while count < len(blobs_to_delete):
                 with self.storageClient.batch():
-                    for blob in blobs_to_delete[count:count + MAX_BATCH_SIZE]:
+                    for blob in blobs_to_delete[count : count + MAX_BATCH_SIZE]:
                         blob.delete()
                     count = count + MAX_BATCH_SIZE
             self.bucket.delete()
 
     def _new_job_id(self):
-        return f'job-{uuid.uuid4()}'
+        return f"job-{uuid.uuid4()}"
 
     def assign_job_id(self, job_description):
         jobStoreID = self._new_job_id()
@@ -219,12 +232,17 @@ class GoogleJobStore(AbstractJobStore):
 
     def create_job(self, job_description):
         job_description.pre_update_hook()
-        self._write_bytes(job_description.jobStoreID, pickle.dumps(job_description, protocol=pickle.HIGHEST_PROTOCOL))
+        self._write_bytes(
+            job_description.jobStoreID,
+            pickle.dumps(job_description, protocol=pickle.HIGHEST_PROTOCOL),
+        )
         return job_description
 
     @google_retry
     def job_exists(self, job_id):
-        return self.bucket.blob(compat_bytes(job_id), encryption_key=self.sseKey).exists()
+        return self.bucket.blob(
+            compat_bytes(job_id), encryption_key=self.sseKey
+        ).exists()
 
     @google_retry
     def get_public_url(self, fileName):
@@ -251,7 +269,11 @@ class GoogleJobStore(AbstractJobStore):
 
     def update_job(self, job):
         job.pre_update_hook()
-        self._write_bytes(job.jobStoreID, pickle.dumps(job, protocol=pickle.HIGHEST_PROTOCOL), update=True)
+        self._write_bytes(
+            job.jobStoreID,
+            pickle.dumps(job, protocol=pickle.HIGHEST_PROTOCOL),
+            update=True,
+        )
 
     @google_retry
     def delete_job(self, job_id):
@@ -269,32 +291,40 @@ class GoogleJobStore(AbstractJobStore):
 
         env = {}
 
-        credentials_path: Optional[str] = os.getenv('GOOGLE_APPLICATION_CREDENTIALS', None)
+        credentials_path: Optional[str] = os.getenv(
+            "GOOGLE_APPLICATION_CREDENTIALS", None
+        )
         if credentials_path is not None:
             # Send along the environment variable that points to the credentials file.
             # It must be available in the same place on all nodes.
-            env['GOOGLE_APPLICATION_CREDENTIALS'] = credentials_path
+            env["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
 
         return env
 
     @google_retry
     def jobs(self):
-        for blob in self.bucket.list_blobs(prefix=b'job-'):
+        for blob in self.bucket.list_blobs(prefix=b"job-"):
             jobStoreID = blob.name
             # TODO: do this better
-            if len(jobStoreID) == 40 and jobStoreID.startswith('job-'):  # 'job-' + uuid length
+            if len(jobStoreID) == 40 and jobStoreID.startswith(
+                "job-"
+            ):  # 'job-' + uuid length
                 yield self.load_job(jobStoreID)
 
     def write_file(self, local_path, job_id=None, cleanup=False):
         fileID = self._new_id(isFile=True, jobStoreID=job_id if cleanup else None)
-        with open(local_path, 'rb') as f:
+        with open(local_path, "rb") as f:
             self._write_file(fileID, f)
         return fileID
 
     @contextmanager
-    def write_file_stream(self, job_id=None, cleanup=False, basename=None, encoding=None, errors=None):
+    def write_file_stream(
+        self, job_id=None, cleanup=False, basename=None, encoding=None, errors=None
+    ):
         fileID = self._new_id(isFile=True, jobStoreID=job_id if cleanup else None)
-        with self._upload_stream(fileID, update=False, encoding=encoding, errors=errors) as writable:
+        with self._upload_stream(
+            fileID, update=False, encoding=encoding, errors=errors
+        ) as writable:
             yield writable, fileID
 
     def get_empty_file_store_id(self, jobStoreID=None, cleanup=False, basename=None):
@@ -309,16 +339,19 @@ class GoogleJobStore(AbstractJobStore):
         if not self.file_exists(file_id):
             raise NoSuchFileException(file_id)
         with AtomicFileCreate(local_path) as tmpPath:
-            with open(tmpPath, 'wb') as writeable:
-                blob = self.bucket.get_blob(compat_bytes(file_id), encryption_key=self.sseKey)
+            with open(tmpPath, "wb") as writeable:
+                blob = self.bucket.get_blob(
+                    compat_bytes(file_id), encryption_key=self.sseKey
+                )
                 blob.download_to_file(writeable)
-        if getattr(file_id, 'executable', False):
+        if getattr(file_id, "executable", False):
             os.chmod(local_path, os.stat(local_path).st_mode | stat.S_IXUSR)
 
     @contextmanager
     def read_file_stream(self, file_id, encoding=None, errors=None):
-        with self.read_shared_file_stream(file_id, isProtected=True, encoding=encoding,
-                                          errors=errors) as readable:
+        with self.read_shared_file_stream(
+            file_id, isProtected=True, encoding=encoding, errors=errors
+        ) as readable:
             yield readable
 
     def delete_file(self, file_id):
@@ -326,32 +359,49 @@ class GoogleJobStore(AbstractJobStore):
 
     @google_retry
     def file_exists(self, file_id):
-        return self.bucket.blob(compat_bytes(file_id), encryption_key=self.sseKey).exists()
+        return self.bucket.blob(
+            compat_bytes(file_id), encryption_key=self.sseKey
+        ).exists()
 
     @google_retry
     def get_file_size(self, file_id):
         if not self.file_exists(file_id):
             return 0
-        return self.bucket.get_blob(compat_bytes(file_id), encryption_key=self.sseKey).size
+        return self.bucket.get_blob(
+            compat_bytes(file_id), encryption_key=self.sseKey
+        ).size
 
     def update_file(self, file_id, local_path):
-        with open(local_path, 'rb') as f:
+        with open(local_path, "rb") as f:
             self._write_file(file_id, f, update=True)
 
     @contextmanager
     def update_file_stream(self, file_id, encoding=None, errors=None):
-        with self._upload_stream(file_id, update=True, encoding=encoding, errors=errors) as writable:
+        with self._upload_stream(
+            file_id, update=True, encoding=encoding, errors=errors
+        ) as writable:
             yield writable
 
     @contextmanager
-    def write_shared_file_stream(self, shared_file_name, encrypted=True, encoding=None, errors=None):
-        with self._upload_stream(shared_file_name, encrypt=encrypted, update=True, encoding=encoding,
-                                 errors=errors) as writable:
+    def write_shared_file_stream(
+        self, shared_file_name, encrypted=True, encoding=None, errors=None
+    ):
+        with self._upload_stream(
+            shared_file_name,
+            encrypt=encrypted,
+            update=True,
+            encoding=encoding,
+            errors=errors,
+        ) as writable:
             yield writable
 
     @contextmanager
-    def read_shared_file_stream(self, shared_file_name, isProtected=True, encoding=None, errors=None):
-        with self._download_stream(shared_file_name, encrypt=isProtected, encoding=encoding, errors=errors) as readable:
+    def read_shared_file_stream(
+        self, shared_file_name, isProtected=True, encoding=None, errors=None
+    ):
+        with self._download_stream(
+            shared_file_name, encrypt=isProtected, encoding=encoding, errors=errors
+        ) as readable:
             yield readable
 
     @classmethod
@@ -374,7 +424,7 @@ class GoogleJobStore(AbstractJobStore):
         fileName = url.path
 
         # remove leading '/', which can cause problems if fileName is a path
-        if fileName.startswith('/'):
+        if fileName.startswith("/"):
             fileName = fileName[1:]
 
         storageClient = cls.create_client()
@@ -383,7 +433,7 @@ class GoogleJobStore(AbstractJobStore):
 
         if exists:
             if not blob.exists():
-                raise NoSuchFileException
+                raise NoSuchFileException(fileName)
             # sync with cloud so info like size is available
             blob.reload()
         return blob
@@ -413,7 +463,7 @@ class GoogleJobStore(AbstractJobStore):
 
     @classmethod
     def _supports_url(cls, url, export=False):
-        return url.scheme.lower() == 'gs'
+        return url.scheme.lower() == "gs"
 
     @classmethod
     def _write_to_url(cls, readable: bytes, url: str, executable: bool = False) -> None:
@@ -421,12 +471,16 @@ class GoogleJobStore(AbstractJobStore):
         blob.upload_from_file(readable)
 
     @classmethod
-    def _list_url(cls, url: ParseResult) -> List[str]:
-        raise NotImplementedError("Listing files in Google buckets is not yet implemented!")
+    def _list_url(cls, url: ParseResult) -> list[str]:
+        raise NotImplementedError(
+            "Listing files in Google buckets is not yet implemented!"
+        )
 
     @classmethod
     def _get_is_directory(cls, url: ParseResult) -> bool:
-        raise NotImplementedError("Checking directory status in Google buckets is not yet implemented!")
+        raise NotImplementedError(
+            "Checking directory status in Google buckets is not yet implemented!"
+        )
 
     @google_retry
     def write_logs(self, msg: bytes) -> None:
@@ -456,7 +510,9 @@ class GoogleJobStore(AbstractJobStore):
                     if not read_all:
                         # rename this file by copying it and deleting the old version to avoid
                         # rereading it
-                        newID = self.readStatsBaseID + blob.name[len(self.statsBaseID):]
+                        newID = (
+                            self.readStatsBaseID + blob.name[len(self.statsBaseID) :]
+                        )
                         # NOTE: just copies then deletes old.
                         self.bucket.rename_blob(blob, compat_bytes(newID))
                 except NoSuchFileException:
@@ -472,7 +528,7 @@ class GoogleJobStore(AbstractJobStore):
                 if lastTry:
                     # this was our second try, we are reasonably sure there aren't any stats
                     # left to gather
-                        break
+                    break
                 # Try one more time in a couple seconds
                 time.sleep(5)
                 lastTry = True
@@ -486,11 +542,11 @@ class GoogleJobStore(AbstractJobStore):
     @staticmethod
     def _new_id(isFile=False, jobStoreID=None):
         if isFile and jobStoreID:  # file associated with job
-            return jobStoreID+str(uuid.uuid4())
+            return jobStoreID + str(uuid.uuid4())
         elif isFile:  # nonassociated file
             return str(uuid.uuid4())
         else:  # job id
-            return f'job-{uuid.uuid4()}'
+            return f"job-{uuid.uuid4()}"
 
     @google_retry
     def _delete(self, jobStoreFileID):
@@ -514,8 +570,12 @@ class GoogleJobStore(AbstractJobStore):
         return job.download_as_string()
 
     @google_retry
-    def _write_file(self, jobStoreID: str, fileObj: bytes, update=False, encrypt=True) -> None:
-        blob = self.bucket.blob(compat_bytes(jobStoreID), encryption_key=self.sseKey if encrypt else None)
+    def _write_file(
+        self, jobStoreID: str, fileObj: bytes, update=False, encrypt=True
+    ) -> None:
+        blob = self.bucket.blob(
+            compat_bytes(jobStoreID), encryption_key=self.sseKey if encrypt else None
+        )
         if not update:
             # TODO: should probably raise a special exception and be added to all jobStores
             assert not blob.exists()
@@ -529,7 +589,9 @@ class GoogleJobStore(AbstractJobStore):
 
     @contextmanager
     @google_retry
-    def _upload_stream(self, fileName, update=False, encrypt=True, encoding=None, errors=None):
+    def _upload_stream(
+        self, fileName, update=False, encrypt=True, encoding=None, errors=None
+    ):
         """
         Yields a context manager that can be used to write to the bucket
         with a stream. See :class:`~toil.jobStores.utils.WritablePipe` for an example.
@@ -555,7 +617,10 @@ class GoogleJobStore(AbstractJobStore):
         :return: an instance of WritablePipe.
         :rtype: :class:`~toil.jobStores.utils.writablePipe`
         """
-        blob = self.bucket.blob(compat_bytes(fileName), encryption_key=self.sseKey if encrypt else None)
+        blob = self.bucket.blob(
+            compat_bytes(fileName), encryption_key=self.sseKey if encrypt else None
+        )
+
         class UploadPipe(WritablePipe):
             def readFrom(self, readable):
                 if not update:
@@ -591,7 +656,9 @@ class GoogleJobStore(AbstractJobStore):
         :rtype: :class:`~toil.jobStores.utils.ReadablePipe`
         """
 
-        blob = self.bucket.get_blob(compat_bytes(fileName), encryption_key=self.sseKey if encrypt else None)
+        blob = self.bucket.get_blob(
+            compat_bytes(fileName), encryption_key=self.sseKey if encrypt else None
+        )
         if blob is None:
             raise NoSuchFileException(fileName)
 

toil/jobStores/utils.py

@@ -10,6 +10,7 @@ from toil.lib.threading import ExceptionalThread
 
 log = logging.getLogger(__name__)
 
+
 class WritablePipe(ABC):
     """
     An object-oriented wrapper for os.pipe. Clients should subclass it, implement
@@ -84,7 +85,7 @@ class WritablePipe(ABC):
         raise NotImplementedError()
 
     def _reader(self):
-        with os.fdopen(self.readable_fh, 'rb') as readable:
+        with os.fdopen(self.readable_fh, "rb") as readable:
             # TODO: If the reader somehow crashes here, both threads might try
             # to close readable_fh.  Fortunately we don't do anything that
             # should be able to fail here.
@@ -112,7 +113,12 @@ class WritablePipe(ABC):
 
     def __enter__(self):
         self.readable_fh, writable_fh = os.pipe()
-        self.writable = os.fdopen(writable_fh, 'wb' if self.encoding == None else 'wt', encoding=self.encoding, errors=self.errors)
+        self.writable = os.fdopen(
+            writable_fh,
+            "wb" if self.encoding == None else "wt",
+            encoding=self.encoding,
+            errors=self.errors,
+        )
         self.thread = ExceptionalThread(target=self._reader)
         self.thread.start()
         return self.writable
@@ -132,7 +138,9 @@ class WritablePipe(ABC):
                 # already an exception in the main thread
                 raise
             else:
-                log.error('Swallowing additional exception in reader thread: %s', str(e))
+                log.error(
+                    "Swallowing additional exception in reader thread: %s", str(e)
+                )
         finally:
             # The responsibility for closing the readable end is generally that of the reader
             # thread. To cover the small window before the reader takes over we also close it here.
@@ -217,7 +225,7 @@ class ReadablePipe(ABC):
 
     def _writer(self):
         try:
-            with os.fdopen(self.writable_fh, 'wb') as writable:
+            with os.fdopen(self.writable_fh, "wb") as writable:
                 self.writeTo(writable)
         except OSError as e:
             # The other side of the pipe may have been closed by the
@@ -244,7 +252,12 @@ class ReadablePipe(ABC):
 
     def __enter__(self):
         readable_fh, self.writable_fh = os.pipe()
-        self.readable = os.fdopen(readable_fh, 'rb' if self.encoding == None else 'rt', encoding=self.encoding, errors=self.errors)
+        self.readable = os.fdopen(
+            readable_fh,
+            "rb" if self.encoding == None else "rt",
+            encoding=self.encoding,
+            errors=self.errors,
+        )
         self.thread = ExceptionalThread(target=self._writer)
         self.thread.start()
         return self.readable
@@ -264,6 +277,7 @@ class ReadablePipe(ABC):
                 # already an exception in the main thread
                 raise
 
+
 class ReadableTransformingPipe(ReadablePipe):
     """
     A pipe which is constructed around a readable stream, and which provides a
@@ -296,7 +310,6 @@ class ReadableTransformingPipe(ReadablePipe):
 
     """
 
-
     def __init__(self, source, encoding=None, errors=None):
         """
         :param str encoding: the name of the encoding used to encode the file. Encodings are the same
@@ -323,15 +336,17 @@ class ReadableTransformingPipe(ReadablePipe):
     def writeTo(self, writable):
         self.transform(self.source, writable)
 
+
 class JobStoreUnavailableException(RuntimeError):
     """
     Raised when a particular type of job store is requested but can't be used.
     """
 
+
 def generate_locator(
     job_store_type: str,
     local_suggestion: Optional[str] = None,
-    decoration: Optional[str] = None
+    decoration: Optional[str] = None,
 ) -> str:
     """
     Generate a random locator for a job store of the given type. Raises an
@@ -347,7 +362,7 @@ def generate_locator(
     """
 
     # Prepare decoration for splicing into strings
-    decoration = ('-' + decoration) if decoration else ''
+    decoration = ("-" + decoration) if decoration else ""
 
     try:
         if job_store_type == "google":
@@ -363,6 +378,7 @@ def generate_locator(
         elif job_store_type == "aws":
             # Make sure we have AWS
             from toil.jobStores.aws.jobStore import AWSJobStore  # noqa
+
             # Find a region
             from toil.lib.aws import get_current_aws_region
 
@@ -370,7 +386,9 @@ def generate_locator(
 
             if not region:
                 # We can't generate an AWS job store without a region
-                raise JobStoreUnavailableException(f"{job_store_type} job store can't be made without a region")
+                raise JobStoreUnavailableException(
+                    f"{job_store_type} job store can't be made without a region"
+                )
 
             # Roll a random name
             return f"aws:{region}:toil{decoration}-{str(uuid.uuid4())}"
@@ -380,11 +398,14 @@ def generate_locator(
                 return local_suggestion
             else:
                 # Pick a temp path
-                return os.path.join(tempfile.gettempdir(), 'toil-' + str(uuid.uuid4()) + decoration)
+                return os.path.join(
+                    tempfile.gettempdir(), "toil-" + str(uuid.uuid4()) + decoration
+                )
         else:
-            raise JobStoreUnavailableException(f"{job_store_type} job store isn't known")
+            raise JobStoreUnavailableException(
+                f"{job_store_type} job store isn't known"
+            )
     except ImportError:
-        raise JobStoreUnavailableException(f"libraries for {job_store_type} job store are not installed")
-
-
-
+        raise JobStoreUnavailableException(
+            f"libraries for {job_store_type} job store are not installed"
+        )