From 2c0fee2faec33185ca7fcd2276901977857e2c64 Mon Sep 17 00:00:00 2001
From: Cade Mirchandani <cmirchan@ucsc.edu>
Date: Fri, 4 Mar 2022 09:14:36 -0800
Subject: [PATCH] fix: work around segfault with >100 jobs in google life
 sciences backend (#1451)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* refactor _retry_request to avoid segfault

* fmt

* reworking _get_services to use auth_http

* fix mistake in retry_request

* remove imports outside of func

Co-authored-by: Johannes Köster <johannes.koester@tu-dortmund.de>
---
 snakemake/executors/google_lifesciences.py | 59 +++++++++++++++++-----
 1 file changed, 47 insertions(+), 12 deletions(-)

diff --git a/snakemake/executors/google_lifesciences.py b/snakemake/executors/google_lifesciences.py
index 7c2b21685..45c514cb8 100644
--- a/snakemake/executors/google_lifesciences.py
+++ b/snakemake/executors/google_lifesciences.py
@@ -23,6 +23,7 @@
 from snakemake.common import get_container_image, get_file_hash
 from snakemake.resources import DefaultResources
 
+
 # https://github.com/googleapis/google-api-python-client/issues/299#issuecomment-343255309
 logging.getLogger("googleapiclient.discovery_cache").setLevel(logging.ERROR)
 
@@ -139,28 +140,56 @@ def _get_services(self):
         for storage.
         """
         from googleapiclient.discovery import build as discovery_build
-        from oauth2client.client import (
-            GoogleCredentials,
-            ApplicationDefaultCredentialsError,
-        )
         from google.cloud import storage
+        import google.auth
+        import google_auth_httplib2
+        import httplib2
+        import googleapiclient
 
         # Credentials must be exported to environment
         try:
-            creds = GoogleCredentials.get_application_default()
-        except ApplicationDefaultCredentialsError as ex:
+            # oauth2client is deprecated, see: https://google-auth.readthedocs.io/en/master/oauth2client-deprecation.html
+            # google.auth is replacement
+            # not sure about scopes here. this cover all cloud services
+            creds, _ = google.auth.default(
+                scopes=["https://www.googleapis.com/auth/cloud-platform"]
+            )
+        except google.auth.DefaultCredentialsError as ex:
             log_verbose_traceback(ex)
             raise ex
 
+        def build_request(http, *args, **kwargs):
+            """
+            See https://googleapis.github.io/google-api-python-client/docs/thread_safety.html
+            """
+            new_http = google_auth_httplib2.AuthorizedHttp(creds, http=httplib2.Http())
+            return googleapiclient.http.HttpRequest(new_http, *args, **kwargs)
+
         # Discovery clients for Google Cloud Storage and Life Sciences API
+        # create authorized http for building services
+        authorized_http = google_auth_httplib2.AuthorizedHttp(
+            creds, http=httplib2.Http()
+        )
         self._storage_cli = discovery_build(
-            "storage", "v1", credentials=creds, cache_discovery=False
+            "storage",
+            "v1",
+            cache_discovery=False,
+            requestBuilder=build_request,
+            http=authorized_http,
         )
         self._compute_cli = discovery_build(
-            "compute", "v1", credentials=creds, cache_discovery=False
+            "compute",
+            "v1",
+            cache_discovery=False,
+            requestBuilder=build_request,
+            http=authorized_http,
         )
         self._api = discovery_build(
-            "lifesciences", "v2beta", credentials=creds, cache_discovery=False
+            "lifesciences",
+            "v2beta",
+            cache_discovery=False,
+            requestBuilder=build_request,
+            http=authorized_http,
         )
         self._bucket_service = storage.Client()
 
@@ -903,18 +932,24 @@ def _retry_request(self, request, timeout=2, attempts=3):
         except BrokenPipeError as ex:
             if attempts > 0:
                 time.sleep(timeout)
-                return self._retry_request(request, timeout * 2, attempts - 1)
+                return self._retry_request(
+                    request, timeout=timeout * 2, attempts=attempts - 1
+                )
             raise ex
         except googleapiclient.errors.HttpError as ex:
             if attempts > 0:
                 time.sleep(timeout)
-                return self._retry_request(request, timeout * 2, attempts - 1)
+                return self._retry_request(
+                    request, timeout=timeout * 2, attempts=attempts - 1
+                )
             log_verbose_traceback(ex)
             raise ex
         except Exception as ex:
             if attempts > 0:
                 time.sleep(timeout)
-                return self._retry_request(request, timeout * 2, attempts - 1)
+                return self._retry_request(
+                    request, timeout=timeout * 2, attempts=attempts - 1
+                )
             log_verbose_traceback(ex)
             raise ex